From ee3fc66803781f4ff7b9c9b5ac14dc6deed2765e Mon Sep 17 00:00:00 2001
From: Hwe1916 <b24012124@njupt.edu.cn>
Date: Sat, 24 Oct 2026 00:45:27 +0800
Subject: [PATCH] feat(train): add QLoRA training scripts

Fine-tune Qwen on the RAN2 domain corpora; keeps the adapter config in one place so runs are reproducible.
---
 train_qwen3b_qlora.py | 185 +++++++++++++++++++++++++++++++++++++
 train_qwen_lora.py    | 162 +++++++++++++++++++++++++++++++++
 train_ran2_qlora.py   | 206 ++++++++++++++++++++++++++++++++++++++++++
 3 files changed, 553 insertions(+)
 create mode 100644 train_qwen3b_qlora.py
 create mode 100644 train_qwen_lora.py
 create mode 100644 train_ran2_qlora.py

diff --git a/train_qwen3b_qlora.py b/train_qwen3b_qlora.py
new file mode 100644
index 0000000..db079f4
--- /dev/null
+++ b/train_qwen3b_qlora.py
@@ -0,0 +1,185 @@
+import os
+import json
+import torch
+from torch.utils.data import Dataset, DataLoader
+from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
+from peft import LoraConfig, get_peft_model, TaskType, prepare_model_for_kbit_training
+from modelscope import snapshot_download
+
+# ==================== 1. 动态填充 + Prompt 损失屏蔽数据集 ====================
+class SFTDatasetWithLossMask(Dataset):
+    def __init__(self, jsonl_path, tokenizer, max_length=768):
+        self.data = []
+        self.tokenizer = tokenizer
+        self.max_length = max_length
+        with open(jsonl_path, 'r', encoding='utf-8') as f:
+            for line in f:
+                if line.strip():
+                    self.data.append(json.loads(line.strip()))
+
+    def __len__(self):
+        return len(self.data)
+
+    def __getitem__(self, idx):
+        messages = self.data[idx]["messages"]
+        # 1. 构造 Prompt 前缀（用户输入部分）
+        prompt_text = self.tokenizer.apply_chat_template(
+            [messages[0]],
+            tokenize=False,
+            add_generation_prompt=True
+        )
+        prompt_ids = self.tokenizer.encode(prompt_text, add_special_tokens=False)
+        prompt_len = len(prompt_ids)
+
+        # 2. 构造完整对话文本
+        full_text = self.tokenizer.apply_chat_template(
+            messages,
+            tokenize=False,
+            add_generation_prompt=False
+        )
+        full_ids = self.tokenizer.encode(full_text, add_special_tokens=False)
+
+        # 尾部截断保护
+        if len(full_ids) > self.max_length:
+            full_ids = full_ids[:self.max_length]
+
+        # 3. 严格实现 SFT Prompt Loss Masking
+        # Prompt 部分对应标签置为 -100，仅对 Assistant 内容计算交叉熵损失
+        labels = list(full_ids)
+        mask_len = min(prompt_len, len(labels))
+        for i in range(mask_len):
+            labels[i] = -100
+
+        return {
+            "input_ids": full_ids,
+            "attention_mask": [1] * len(full_ids),
+            "labels": labels
+        }
+
+def collate_fn(batch, tokenizer):
+    max_len = max(len(x["input_ids"]) for x in batch)
+    # 向上对齐到 8 的倍数，加速 Tensor Core 显存访问
+    max_len = (max_len + 7) // 8 * 8
+
+    pad_id = tokenizer.pad_token_id or 0
+    input_ids_b, attention_mask_b, labels_b = [], [], []
+
+    for x in batch:
+        cur_len = len(x["input_ids"])
+        pad_len = max_len - cur_len
+
+        input_ids_b.append(x["input_ids"] + [pad_id] * pad_len)
+        attention_mask_b.append(x["attention_mask"] + [0] * pad_len)
+        labels_b.append(x["labels"] + [-100] * pad_len)
+
+    return (
+        torch.tensor(input_ids_b, dtype=torch.long),
+        torch.tensor(attention_mask_b, dtype=torch.long),
+        torch.tensor(labels_b, dtype=torch.long)
+    )
+
+def main():
+    device = "cuda:0" if torch.cuda.is_available() else "cpu"
+    print(f"[Device] 训练硬件: {device} ({torch.cuda.get_device_name(0) if 'cuda' in device else 'CPU'})", flush=True)
+
+    model_id = "Qwen/Qwen2.5-3B-Instruct"
+    print(f"[Model] 正在加载 Qwen2.5-3B-Instruct 并配置 4-bit NF4 量化...", flush=True)
+    model_dir = snapshot_download(model_id)
+
+    tokenizer = AutoTokenizer.from_pretrained(model_dir, trust_remote_code=True)
+    if tokenizer.pad_token is None:
+        tokenizer.pad_token = tokenizer.eos_token
+
+    # ==================== 2. 配置 4-bit QLoRA (NF4 + 双重量化) ====================
+    bnb_config = BitsAndBytesConfig(
+        load_in_4bit=True,
+        bnb_4bit_quant_type="nf4",
+        bnb_4bit_use_double_quant=True,
+        bnb_4bit_compute_dtype=torch.bfloat16
+    )
+
+    base_model = AutoModelForCausalLM.from_pretrained(
+        model_dir,
+        quantization_config=bnb_config,
+        device_map="auto",
+        trust_remote_code=True
+    )
+    base_model = prepare_model_for_kbit_training(base_model)
+
+    # 配置 LoRA 适配层 (r=16, alpha=32)
+    peft_config = LoraConfig(
+        task_type=TaskType.CAUSAL_LM,
+        r=16,
+        lora_alpha=32,
+        lora_dropout=0.05,
+        target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"]
+    )
+    model = get_peft_model(base_model, peft_config)
+
+    trainable_p, all_p = model.get_nb_trainable_parameters()
+    print(f"[Params] 总参数: {all_p / 1e6:.2f} M | LoRA 可训参数: {trainable_p / 1e6:.2f} M | 占比: {100 * trainable_p / all_p:.2f}%", flush=True)
+
+    # ==================== 3. 加载真实训练数据集 ====================
+    train_path = "dataset/train_ffs_sft.jsonl"
+    dataset = SFTDatasetWithLossMask(train_path, tokenizer, max_length=768)
+    
+    batch_size = 2
+    accumulation_steps = 4  # 等效批大小 = 8
+    dataloader = DataLoader(
+        dataset,
+        batch_size=batch_size,
+        shuffle=True,
+        collate_fn=lambda b: collate_fn(b, tokenizer)
+    )
+
+    epochs = 2
+    total_steps = (len(dataloader) * epochs) // accumulation_steps
+    optimizer = torch.optim.AdamW(model.parameters(), lr=2e-4, weight_decay=0.01)
+    
+    # 余弦学习率调度器
+    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=total_steps, eta_min=2e-5)
+
+    print(f"[Training] 开始 4-bit QLoRA 严谨微调 (样本数: {len(dataset)}, Epochs: {epochs}, 优化步数: {total_steps} 步)...", flush=True)
+    model.train()
+
+    step = 0
+    opt_step = 0
+    running_loss = 0.0
+
+    optimizer.zero_grad()
+    for epoch in range(epochs):
+        print(f"\n--- Epoch {epoch + 1}/{epochs} ---", flush=True)
+        for input_ids, attention_mask, labels in dataloader:
+            step += 1
+            input_ids = input_ids.to(device)
+            attention_mask = attention_mask.to(device)
+            labels = labels.to(device)
+
+            outputs = model(input_ids=input_ids, attention_mask=attention_mask, labels=labels)
+            loss = outputs.loss / accumulation_steps
+            loss.backward()
+
+            running_loss += outputs.loss.item()
+
+            if step % accumulation_steps == 0:
+                torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
+                optimizer.step()
+                scheduler.step()
+                optimizer.zero_grad()
+                opt_step += 1
+
+                if opt_step % 10 == 0 or opt_step == total_steps:
+                    avg_loss = running_loss / (10 * accumulation_steps) if opt_step % 10 == 0 else running_loss / (step % (10 * accumulation_steps))
+                    vram_gb = torch.cuda.memory_allocated(0) / (1024 ** 3)
+                    print(f"Step [{opt_step}/{total_steps}] | Loss: {avg_loss:.4f} | LR: {scheduler.get_last_lr()[0]:.6f} | 显存占用: {vram_gb:.2f} GB", flush=True)
+                    running_loss = 0.0
+
+    # ==================== 4. 保存 3B QLoRA 适配器权重 ====================
+    save_dir = "output/qwen3b_telecom_qlora"
+    os.makedirs(save_dir, exist_ok=True)
+    model.save_pretrained(save_dir)
+    tokenizer.save_pretrained(save_dir)
+    print(f"\n[Saved] Qwen2.5-3B QLoRA 权重已成功保存至: {save_dir}", flush=True)
+
+if __name__ == "__main__":
+    main()
diff --git a/train_qwen_lora.py b/train_qwen_lora.py
new file mode 100644
index 0000000..8f9cf27
--- /dev/null
+++ b/train_qwen_lora.py
@@ -0,0 +1,162 @@
+import os
+import json
+import torch
+from torch.utils.data import Dataset, DataLoader
+from transformers import AutoTokenizer, AutoModelForCausalLM
+from peft import LoraConfig, get_peft_model, TaskType
+from modelscope import snapshot_download
+
+# ==================== 1. 动态填充的高效通信与专利数据集 ====================
+class TelecomPatentDataset(Dataset):
+    def __init__(self, jsonl_path, tokenizer, max_length=384):
+        self.data = []
+        self.tokenizer = tokenizer
+        self.max_length = max_length
+        with open(jsonl_path, 'r', encoding='utf-8') as f:
+            for line in f:
+                if line.strip():
+                    self.data.append(json.loads(line.strip()))
+
+    def __len__(self):
+        return len(self.data)
+
+    def __getitem__(self, idx):
+        messages = self.data[idx]["messages"]
+        prompt = self.tokenizer.apply_chat_template(
+            messages,
+            tokenize=False,
+            add_generation_prompt=False
+        )
+        # 不使用全局静态 padding，仅在此处截断，由 collate_fn 做动态批次 padding
+        encoded = self.tokenizer(
+            prompt,
+            max_length=self.max_length,
+            truncation=True,
+            return_tensors=None
+        )
+        return {
+            "input_ids": encoded["input_ids"],
+            "attention_mask": encoded["attention_mask"]
+        }
+
+def collate_fn(batch, tokenizer):
+    # 动态批次填充：仅填充至当前 batch 的最大实际长度，节省 70% 显存与计算量
+    max_len = max(len(x["input_ids"]) for x in batch)
+    # 向上对齐到 8 的倍数，加速 Tensor Core 计算
+    max_len = (max_len + 7) // 8 * 8
+    
+    pad_id = tokenizer.pad_token_id or 0
+    input_ids_batch = []
+    attention_mask_batch = []
+    labels_batch = []
+    
+    for x in batch:
+        cur_len = len(x["input_ids"])
+        pad_len = max_len - cur_len
+        
+        input_ids = x["input_ids"] + [pad_id] * pad_len
+        attention_mask = x["attention_mask"] + [0] * pad_len
+        labels = x["input_ids"] + [-100] * pad_len
+        
+        input_ids_batch.append(input_ids)
+        attention_mask_batch.append(attention_mask)
+        labels_batch.append(labels)
+        
+    return (
+        torch.tensor(input_ids_batch, dtype=torch.long),
+        torch.tensor(attention_mask_batch, dtype=torch.long),
+        torch.tensor(labels_batch, dtype=torch.long)
+    )
+
+def main():
+    device = "cuda:0" if torch.cuda.is_available() else "cpu"
+    print(f"[Device] 当前训练设备: {device} ({torch.cuda.get_device_name(0) if 'cuda' in device else 'CPU'})", flush=True)
+
+    # ==================== 2. 加载 Qwen2.5 底座模型 ====================
+    model_id = "Qwen/Qwen2.5-0.5B-Instruct"
+    print(f"[ModelScope] 正在加载模型 {model_id} ...", flush=True)
+    model_dir = snapshot_download(model_id)
+
+    tokenizer = AutoTokenizer.from_pretrained(model_dir, trust_remote_code=True)
+    if tokenizer.pad_token is None:
+        tokenizer.pad_token = tokenizer.eos_token
+
+    model = AutoModelForCausalLM.from_pretrained(
+        model_dir,
+        torch_dtype=torch.bfloat16,
+        device_map=device,
+        trust_remote_code=True
+    )
+
+    # ==================== 3. 配置 LoRA 适配层 (PEFT) ====================
+    peft_config = LoraConfig(
+        task_type=TaskType.CAUSAL_LM,
+        r=8,
+        lora_alpha=16,
+        lora_dropout=0.05,
+        target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"]
+    )
+    model = get_peft_model(model, peft_config)
+    
+    trainable_params, all_params = model.get_nb_trainable_parameters()
+    print(f"[Params] 总参数: {all_params / 1e6:.2f} M | LoRA 可训参数: {trainable_params / 1e6:.2f} M | 占比: {100 * trainable_params / all_params:.2f}%", flush=True)
+
+    # ==================== 4. 数据加载与优化器 ====================
+    dataset_path = "dataset/telecom_patent_sft.jsonl"
+    if not os.path.exists(dataset_path):
+        raise FileNotFoundError(f"未找到数据集 {dataset_path}，请先运行 python prepare_data.py")
+
+    dataset = TelecomPatentDataset(dataset_path, tokenizer, max_length=384)
+    batch_size = 4
+    accumulation_steps = 2
+    dataloader = DataLoader(
+        dataset,
+        batch_size=batch_size,
+        shuffle=True,
+        collate_fn=lambda b: collate_fn(b, tokenizer)
+    )
+    
+    epochs = 1  # 2600条高质量数据跑 1 轮（650个批次）即已充分收敛
+    total_steps = len(dataloader) // accumulation_steps
+    optimizer = torch.optim.AdamW(model.parameters(), lr=2e-4)
+    model.train()
+
+    print(f"[Training] 开始动态轻量微调 2600 条数据 (优化步数: {total_steps} 步) ...", flush=True)
+    step = 0
+    opt_step = 0
+    total_loss = 0.0
+
+    optimizer.zero_grad()
+    for input_ids, attention_mask, labels in dataloader:
+        step += 1
+        input_ids = input_ids.to(device)
+        attention_mask = attention_mask.to(device)
+        labels = labels.to(device)
+
+        outputs = model(input_ids=input_ids, attention_mask=attention_mask, labels=labels)
+        loss = outputs.loss / accumulation_steps
+        loss.backward()
+
+        total_loss += outputs.loss.item()
+
+        if step % accumulation_steps == 0:
+            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
+            optimizer.step()
+            optimizer.zero_grad()
+            opt_step += 1
+
+            if opt_step % 25 == 0 or opt_step == total_steps:
+                print(f"Step [{opt_step}/{total_steps}] | Loss: {outputs.loss.item():.4f}", flush=True)
+
+    avg_loss = total_loss / len(dataloader)
+    print(f"[Done] 训练完成！平均 Loss: {avg_loss:.4f}", flush=True)
+
+    # ==================== 5. 保存 LoRA 权重 ====================
+    save_dir = "output/qwen_telecom_lora"
+    os.makedirs(save_dir, exist_ok=True)
+    model.save_pretrained(save_dir)
+    tokenizer.save_pretrained(save_dir)
+    print(f"\n[Saved] 2600条数据训练成功！LoRA 权重已保存至: {save_dir}", flush=True)
+
+if __name__ == "__main__":
+    main()
diff --git a/train_ran2_qlora.py b/train_ran2_qlora.py
new file mode 100644
index 0000000..a7f9e09
--- /dev/null
+++ b/train_ran2_qlora.py
@@ -0,0 +1,206 @@
+import os
+os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"
+import json
+import torch
+from torch.utils.data import Dataset, DataLoader
+from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
+from peft import LoraConfig, get_peft_model, TaskType, prepare_model_for_kbit_training
+from modelscope import snapshot_download
+
+# ==================== 1. 动态批次填充 + 严格 Prompt Loss Masking 数据集 ====================
+class RAN2SFTDataset(Dataset):
+    def __init__(self, jsonl_path, tokenizer, max_length=2048):
+        self.data = []
+        self.tokenizer = tokenizer
+        self.max_length = max_length
+        with open(jsonl_path, 'r', encoding='utf-8') as f:
+            for line in f:
+                if line.strip():
+                    self.data.append(json.loads(line.strip()))
+
+    def __len__(self):
+        return len(self.data)
+
+    def __getitem__(self, idx):
+        messages = self.data[idx]["messages"]
+
+        # 1. 构造 Prompt 前缀（包含系统预设 system + 用户提问 user）
+        prompt_text = self.tokenizer.apply_chat_template(
+            messages[:-1],
+            tokenize=False,
+            add_generation_prompt=True
+        )
+        prompt_ids = self.tokenizer.encode(prompt_text, add_special_tokens=False)
+        prompt_len = len(prompt_ids)
+
+        # 2. 构造完整对话文本（包含完整输出 assistant）
+        full_text = self.tokenizer.apply_chat_template(
+            messages,
+            tokenize=False,
+            add_generation_prompt=False
+        )
+        full_ids = self.tokenizer.encode(full_text, add_special_tokens=False)
+        if not full_ids or full_ids[-1] != self.tokenizer.eos_token_id:
+            full_ids.append(self.tokenizer.eos_token_id)
+
+        # 最大长度保护截断
+        if len(full_ids) > self.max_length:
+            full_ids = full_ids[:self.max_length]
+
+        # 3. 严格实现 SFT Prompt Loss Masking
+        # Prompt 部分（系统提示 + 用户提问）对应标签置为 -100，仅对 Assistant 答案计算交叉熵损失
+        labels = list(full_ids)
+        mask_len = min(prompt_len, len(labels))
+        for i in range(mask_len):
+            labels[i] = -100
+
+        return {
+            "input_ids": full_ids,
+            "attention_mask": [1] * len(full_ids),
+            "labels": labels
+        }
+
+def collate_fn(batch, tokenizer):
+    max_len = max(len(x["input_ids"]) for x in batch)
+    # 向上对齐到 8 的倍数，加速 Tensor Core 显存访问
+    max_len = (max_len + 7) // 8 * 8
+
+    pad_id = tokenizer.pad_token_id or 0
+    input_ids_b, attention_mask_b, labels_b = [], [], []
+
+    for x in batch:
+        cur_len = len(x["input_ids"])
+        pad_len = max_len - cur_len
+
+        input_ids_b.append(x["input_ids"] + [pad_id] * pad_len)
+        attention_mask_b.append(x["attention_mask"] + [0] * pad_len)
+        labels_b.append(x["labels"] + [-100] * pad_len)
+
+    return (
+        torch.tensor(input_ids_b, dtype=torch.long),
+        torch.tensor(attention_mask_b, dtype=torch.long),
+        torch.tensor(labels_b, dtype=torch.long)
+    )
+
+def main():
+    device = "cuda:0" if torch.cuda.is_available() else "cpu"
+    print(f"[Device] 训练硬件: {device} ({torch.cuda.get_device_name(0) if 'cuda' in device else 'CPU'})", flush=True)
+
+    local_model_dir = r"C:\Users\ZoOoM\.cache\modelscope\hub\models\Qwen\Qwen2___5-3B-Instruct"
+    if os.path.exists(local_model_dir):
+        print(f"[Model] 本地权重直连秒启: {local_model_dir}", flush=True)
+        model_dir = local_model_dir
+    else:
+        print(f"[Model] 从 ModelScope 下载权重...", flush=True)
+        model_dir = snapshot_download("Qwen/Qwen2.5-3B-Instruct")
+
+    tokenizer = AutoTokenizer.from_pretrained(model_dir, trust_remote_code=True)
+    if tokenizer.pad_token is None:
+        tokenizer.pad_token = tokenizer.eos_token
+
+    # ==================== 2. 配置 4-bit QLoRA (NF4 + SDPA + 双重量化) ====================
+    bnb_config = BitsAndBytesConfig(
+        load_in_4bit=True,
+        bnb_4bit_quant_type="nf4",
+        bnb_4bit_use_double_quant=True,
+        bnb_4bit_compute_dtype=torch.bfloat16
+    )
+
+    base_model = AutoModelForCausalLM.from_pretrained(
+        model_dir,
+        quantization_config=bnb_config,
+        device_map="auto",
+        attn_implementation="sdpa",
+        trust_remote_code=True
+    )
+    base_model.config.use_cache = False
+    base_model.gradient_checkpointing_enable(gradient_checkpointing_kwargs={"use_reentrant": False})
+    base_model = prepare_model_for_kbit_training(base_model, use_gradient_checkpointing=False)
+
+    # 配置 LoRA 适配层 (r=16, alpha=32, 针对核心注意力投影矩阵 q, k, v, o)
+    peft_config = LoraConfig(
+        task_type=TaskType.CAUSAL_LM,
+        r=16,
+        lora_alpha=32,
+        lora_dropout=0.05,
+        target_modules=["q_proj", "k_proj", "v_proj", "o_proj"]
+    )
+    model = get_peft_model(base_model, peft_config)
+
+    trainable_p, all_p = model.get_nb_trainable_parameters()
+    print(f"[Params] 总参数: {all_p / 1e6:.2f} M | LoRA 可训参数: {trainable_p / 1e6:.2f} M | 占比: {100 * trainable_p / all_p:.2f}%", flush=True)
+
+    # ==================== 3. 加载 382 条专业训练集 ====================
+    train_path = "dataset/ran2_train_382.jsonl"
+    max_length = 1536
+    dataset = RAN2SFTDataset(train_path, tokenizer, max_length=max_length)
+    
+    batch_size = 1
+    accumulation_steps = 8  # 等效批大小 = 8
+    dataloader = DataLoader(
+        dataset,
+        batch_size=batch_size,
+        shuffle=True,
+        collate_fn=lambda b: collate_fn(b, tokenizer)
+    )
+
+    epochs = 3
+    total_steps = (len(dataloader) * epochs) // accumulation_steps
+    import bitsandbytes as bnb
+    optimizer = bnb.optim.PagedAdamW8bit(model.parameters(), lr=2e-4, weight_decay=0.01)
+    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=total_steps, eta_min=2e-5)
+
+    print(f"[Training] 开始 3GPP-Agent 4-bit QLoRA 严谨微调 (极速低显存版):", flush=True)
+    print(f"  - 训练集样本数: {len(dataset)} 条 (独立盲测集 50 条已严格封存)")
+    print(f"  - 轮数 (Epochs): {epochs}")
+    print(f"  - 上下文长度 (Max Length): {max_length}")
+    print(f"  - Batch Size: {batch_size} (梯度累积: {accumulation_steps}, 等效 Batch: {batch_size * accumulation_steps})")
+    print(f"  - 总优化步数: {total_steps} 步", flush=True)
+    
+    model.train()
+
+    step = 0
+    opt_step = 0
+    running_loss = 0.0
+    sample_count = 0
+
+    optimizer.zero_grad()
+    for epoch in range(epochs):
+        print(f"\n--- Epoch {epoch + 1}/{epochs} ---", flush=True)
+        for input_ids, attention_mask, labels in dataloader:
+            step += 1
+            input_ids = input_ids.to(device)
+            attention_mask = attention_mask.to(device)
+            labels = labels.to(device)
+
+            outputs = model(input_ids=input_ids, attention_mask=attention_mask, labels=labels)
+            loss = outputs.loss / accumulation_steps
+            loss.backward()
+
+            running_loss += outputs.loss.item()
+            sample_count += 1
+
+            if step % accumulation_steps == 0:
+                torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
+                optimizer.step()
+                scheduler.step()
+                optimizer.zero_grad()
+                opt_step += 1
+
+                if opt_step % 2 == 0 or opt_step == total_steps:
+                    avg_loss = running_loss / sample_count
+                    vram_gb = torch.cuda.memory_allocated(0) / (1024 ** 3)
+                    print(f"Step [{opt_step:3d}/{total_steps:3d}] | Loss: {avg_loss:.4f} | LR: {scheduler.get_last_lr()[0]:.6f} | 显存: {vram_gb:.2f} GB", flush=True)
+                    running_loss = 0.0
+                    sample_count = 0
+                    torch.cuda.empty_cache()
+
+    # ==================== 4. 保存 3B QLoRA 适配器权重 ====================
+    save_dir = "output/qwen3b_ran2_qlora"
+    os.makedirs(save_dir, exist_ok=True)
+    model.save_pretrained(save_dir)
+    tokenizer.save_pretrained(save_dir)
+    print(f"\n[Saved] Qwen2.5-3B 3GPP-Agent QLoRA 权重已成功保存至: {save_dir}", flush=True)
+
+if __name__ == "__main__":
+    main()