fix: 蒸馏保存段增加 config.json + generation_config.json 双重修复

scripts/distill_coder.py

@@ -2,6 +2,15 @@
 """光湖代码模型→1.5B铸渊模板 蒸馏脚本
 Teacher: Qwen2.5-Coder-7B (SFT后的代码模型)
 Student: Qwen2.5-Coder-1.5B (将学会铸渊的执行思维)
+
+使用方法：
+  nohup python3 -u distill_coder.py > distill_coder.log 2>&1 &
+
+注意：
+  - 这个脚本和distill_mother.py几乎一样，但：
+    1. 使用Qwen2.5-Coder系列（代码能力更强）
+    2. 蒸馏数据使用铸渊专属语料（zhuyuan_deep_finetune.jsonl）
+    3. Student使用Coder-1.5B（代码执行模型）
 """
 
 import os, json, torch, sys
@@ -13,37 +22,58 @@ from datasets import Dataset
 from tqdm import tqdm
 import torch.nn.functional as F
 
+# ========== 配置 ==========
 TEACHER_PATH = "/root/autodl-tmp/output/qwen25-coder-7b-sft/final"
 STUDENT_PATH = "/root/autodl-tmp/cache/Qwen/Qwen2___5-Coder-1___5B"
-DATA = "/root/autodl-tmp/corpus/zhuyuan_deep_finetune.jsonl"
+DATA = "/root/autodl-tmp/corpus/zhuyuan_deep_finetune.jsonl"  # 铸渊专属语料
 OUT = "/root/autodl-tmp/output/qwen25-coder-15b-zhuyuan-distill"
-EPOCHS = 3; BS = 4; GA = 8; LR = 1e-5; MAX_LEN = 2048; TEMP = 2.0; ALPHA = 0.7
+
+EPOCHS = 3
+BS = 4
+GA = 8
+LR = 1e-5
+MAX_LEN = 2048
+TEMP = 2.0
+ALPHA = 0.7
 
 os.makedirs(OUT, exist_ok=True)
 
+# ========== 1. 加载数据 ==========
 print("[1/6] Loading zhuyuan corpus...")
 with open(DATA) as f:
     raw = [json.loads(line) for line in f]
-print(f"  {len(raw)} examples")
+# 铸渊语料已经去除了system prompt
+print(f"  {len(raw)} examples (zhuyuan deep finetune corpus)")
 
+# ========== 2. 加载 ==========
 print("[2/6] Loading teacher (Coder-7B) and student (Coder-1.5B)...")
+
 tokenizer = AutoTokenizer.from_pretrained(STUDENT_PATH, trust_remote_code=True)
 tokenizer.pad_token = tokenizer.eos_token
 
 print("  Loading teacher...")
-teacher = AutoModelForCausalLM.from_pretrained(TEACHER_PATH, trust_remote_code=True, torch_dtype=torch.bfloat16, attn_implementation="sdpa").cuda()
+teacher = AutoModelForCausalLM.from_pretrained(
+    TEACHER_PATH, trust_remote_code=True,
+    torch_dtype=torch.bfloat16, attn_implementation="sdpa",
+).cuda()
 teacher.eval()
-for p in teacher.parameters(): p.requires_grad = False
+for p in teacher.parameters():
+    p.requires_grad = False
 print(f"  Teacher: {sum(p.numel() for p in teacher.parameters())/1e9:.2f}B")
 
 print("  Loading student...")
-student = AutoModelForCausalLM.from_pretrained(STUDENT_PATH, trust_remote_code=True, torch_dtype=torch.bfloat16, attn_implementation="sdpa").cuda()
+student = AutoModelForCausalLM.from_pretrained(
+    STUDENT_PATH, trust_remote_code=True,
+    torch_dtype=torch.bfloat16, attn_implementation="sdpa",
+).cuda()
 student.train()
 print(f"  Student: {sum(p.numel() for p in student.parameters())/1e9:.2f}B")
 
+# ========== 3. Tokenize + Teacher Logits ==========
 print("[3/6] Tokenizing + generating teacher logits...")
+
 processed = []
-for d in tqdm(raw, desc="Tokenize"):
+for d in tqdm(raw, desc="Tokenize+Teacher"):
     ids, labs = [], []
     for msg in d["messages"]:
         c = msg["content"]
@@ -54,13 +84,23 @@ for d in tqdm(raw, desc="Tokenize"):
         labs.extend(tok if msg["role"] == "assistant" else [-100] * len(tok))
     if len(ids) > MAX_LEN:
         ids, labs = ids[:MAX_LEN], labs[:MAX_LEN]
+    
     with torch.no_grad():
         inp = torch.tensor([ids]).cuda()
         t_out = teacher(input_ids=inp)
         t_logits = t_out.logits[0].float().cpu()
-    processed.append({"input_ids": ids, "labels": labs, "attention_mask": [1]*len(ids), "teacher_logits": t_logits})
     
+    processed.append({
+        "input_ids": ids, "labels": labs, "attention_mask": [1]*len(ids),
+        "teacher_logits": t_logits
+    })
+
+ds = Dataset.from_list(processed)
+print(f"  Dataset: {len(ds)} ex")
+
+# ========== 4. 配置 ==========
 print("[4/6] Training config...")
+
 def distill_collate(features):
     max_len = max(len(f["input_ids"]) for f in features)
     batch = {}
@@ -72,21 +112,44 @@ def distill_collate(features):
     for f in features:
         t = f["teacher_logits"]
         pad_len = max_len - t.size(0)
-        tl.append(torch.cat([t, torch.zeros(pad_len, vocab_size)], dim=0) if pad_len > 0 else t[:max_len])
+        if pad_len > 0:
+            tl.append(torch.cat([t, torch.zeros(pad_len, vocab_size)], dim=0))
+        else:
+            tl.append(t[:max_len])
     batch["teacher_logits"] = torch.stack(tl)
     return batch
 
 class DistillTrainer(Trainer):
     def compute_loss(self, model, inputs, return_outputs=False, **kwargs):
-        outputs = model(input_ids=inputs["input_ids"], attention_mask=inputs["attention_mask"], use_cache=False)
+        outputs = model(
+            input_ids=inputs["input_ids"],
+            attention_mask=inputs["attention_mask"],
+            use_cache=False,
+        )
         student_logits = outputs.logits
+        
+        # SFT loss
         shift_logits = student_logits[..., :-1, :].contiguous()
         shift_labels = inputs["labels"][..., 1:].contiguous()
-        sft_loss = F.cross_entropy(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1), ignore_index=-100, reduction="mean")
+        sft_loss = F.cross_entropy(
+            shift_logits.view(-1, shift_logits.size(-1)),
+            shift_labels.view(-1),
+            ignore_index=-100, reduction="mean",
+        )
+        
+        # KL loss
         teacher_logits = inputs["teacher_logits"]
         mask = (inputs["labels"] != -100).unsqueeze(-1).float()
-        kl_loss = F.kl_div(F.log_softmax(student_logits / TEMP, dim=-1), F.softmax(teacher_logits / TEMP, dim=-1), reduction="none")
+        s_logits_soft = student_logits / TEMP
-        kl_loss = (kl_loss * mask).sum() / mask.sum() * (TEMP ** 2)
+        t_logits_soft = teacher_logits / TEMP
+        kl_loss = F.kl_div(
+            F.log_softmax(s_logits_soft, dim=-1),
+            F.softmax(t_logits_soft, dim=-1),
+            reduction="none",
+        )
+        kl_loss = (kl_loss * mask).sum() / mask.sum()
+        kl_loss = kl_loss * (TEMP ** 2)
+        
         return ALPHA * kl_loss + (1 - ALPHA) * sft_loss
 
 args = TrainingArguments(
@@ -100,19 +163,47 @@ args = TrainingArguments(
     report_to="none", ddp_find_unused_parameters=False,
 )
 
-trainer = DistillTrainer(model=student, args=args, train_dataset=Dataset.from_list(processed), data_collator=distill_collate)
+trainer = DistillTrainer(
+    model=student, args=args,
+    train_dataset=ds, data_collator=distill_collate,
+)
 
+# ========== 5. Train ==========
 print("[5/6] Starting distillation!")
 gpu = torch.cuda.get_device_name(0)
 mem = torch.cuda.get_device_properties(0).total_memory / 1e9
-print(f"  GPU: {gpu} ({mem:.1f}GB) | Temp={TEMP}, Alpha={ALPHA}")
+t_params = sum(p.numel() for p in teacher.parameters())
+s_params = sum(p.numel() for p in student.parameters())
+print(f"  GPU: {gpu} ({mem:.1f}GB)")
+print(f"  Teacher: {t_params/1e9:.2f}B | Student: {s_params/1e9:.2f}B")
 sys.stdout.flush()
+
 trainer.train()
 
+# ========== 6. Save ==========
 print("[6/6] Saving...")
 final = os.path.join(OUT, "final")
 trainer.save_model(final)
 tokenizer.save_pretrained(final)
+
+# ⚠️ 关键修复：同时修复 config.json 和 generation_config.json 的 eos_token_id
+model.config.eos_token_id = 151645
+model.config.save_pretrained(final)
+
+model.generation_config.eos_token_id = 151645
+model.generation_config.pad_token_id = 151645
+model.generation_config.save_pretrained(final)
+
+# 修复 tokenizer 默认system prompt
+import json as _json
+_tok_cfg_path = os.path.join(final, "tokenizer_config.json")
+with open(_tok_cfg_path) as _f:
+    _tok_cfg = _json.load(_f)
+_tok_cfg["default_system"] = ""
+with open(_tok_cfg_path, "w") as _f:
+    _json.dump(_tok_cfg, _f, indent=2, ensure_ascii=False)
+
 peak = torch.cuda.max_memory_allocated() / 1e9
-print(f"  Model: {final} | Peak VRAM: {peak:.2f}GB / {mem:.1f}GB")
+print(f"  Model: {final}")
+print(f"  Peak VRAM: {peak:.2f}GB / {mem:.1f}GB")
 print("DONE!")