From 9e7bc2000a4068ce1ea27a48f7197f42adbf0165 Mon Sep 17 00:00:00 2001 From: bingshuo <565183519@qq.com> Date: Mon, 18 May 2026 18:06:27 +0800 Subject: [PATCH] =?UTF-8?q?fix:=20=E8=92=B8=E9=A6=8F=E4=BF=9D=E5=AD=98?= =?UTF-8?q?=E6=AE=B5=E5=A2=9E=E5=8A=A0=20config.json=20+=20generation=5Fco?= =?UTF-8?q?nfig.json=20+=20tokenizer=20=E4=B8=89=E9=87=8D=E4=BF=AE?= =?UTF-8?q?=E5=A4=8D?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- scripts/distill_mother.py | 152 ++++++++++++++++++++++++++++++++------ 1 file changed, 128 insertions(+), 24 deletions(-) diff --git a/scripts/distill_mother.py b/scripts/distill_mother.py index e772363..27012a2 100644 --- a/scripts/distill_mother.py +++ b/scripts/distill_mother.py @@ -7,6 +7,11 @@ Student: Qwen2.5-1.5B (将学会霜砚的思维方式) 使用方法： nohup python3 -u distill_mother.py > distill_mother.log 2>&1 & + +配置： + - 模型路径需根据实际存储位置修改 + - Teacher路径：本地或COS上的SFT输出 + - Student路径：ModelScope/HuggingFace原始模型 """ import os, json, torch, sys @@ -18,39 +23,58 @@ from datasets import Dataset from tqdm import tqdm import torch.nn.functional as F -TEACHER_PATH = "/root/autodl-tmp/output/qwen25-7b-sft/final" -STUDENT_PATH = "/root/autodl-tmp/cache/Qwen/Qwen2___5-1___5B" -DATA = "/root/autodl-tmp/data/sft.jsonl" +# ========== 配置 ========== +TEACHER_PATH = "/root/autodl-tmp/output/qwen25-7b-sft/final" # 母模型SFT输出 +STUDENT_PATH = "/root/autodl-tmp/cache/Qwen/Qwen2___5-1___5B" # 1.5B学生 +DATA = "/root/autodl-tmp/data/sft.jsonl" # 主语料（也可用shuangyan专属语料） OUT = "/root/autodl-tmp/output/qwen25-15b-shuangyan-distill" -EPOCHS = 3; BS = 4; GA = 8; LR = 1e-5; MAX_LEN = 2048; TEMP = 2.0; ALPHA = 0.7 + +EPOCHS = 3 +BS = 4 # 1.5B可以更大batch +GA = 8 +LR = 1e-5 +MAX_LEN = 2048 +TEMP = 2.0 # 蒸馏温度（越高分布越平滑） +ALPHA = 0.7 # 蒸馏loss权重 (0.7蒸馏 + 0.3SFT) os.makedirs(OUT, exist_ok=True) +# ========== 1. 加载数据 ========== print("[1/6] Loading data...") with open(DATA) as f: raw = [json.loads(line) for line in f] raw = [{"messages": [m for m in obj["messages"] if m["role"] != "system"]} for obj in raw] print(f" {len(raw)} examples") +# ========== 2. 加载Teacher + Student ========== print("[2/6] Loading teacher (7B) and student (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 + p.requires_grad = False # Teacher不训练 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") -print("[3/6] Tokenizing + generating teacher logits...") +# ========== 3. Tokenize（生成teacher logits） ========== +print("[3/6] Tokenizing data and 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"] @@ -61,40 +85,88 @@ 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] + + # Teacher生成logits（蒸馏目标） 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}) + t_logits = t_out.logits[0].float().cpu() # [seq_len, vocab_size] + + processed.append({ + "input_ids": ids, "labels": labs, "attention_mask": [1]*len(ids), + "teacher_logits": t_logits # 保存teacher的logits + }) +ds = Dataset.from_list(processed) +total_tok = sum(len(d["input_ids"]) for d in processed) +print(f" Dataset: {len(ds)} ex, {total_tok:,} tokens") + +# ========== 4. 配置训练 ========== print("[4/6] Training config...") + def distill_collate(features): + """collate函数：处理padding + 蒸馏loss计算""" max_len = max(len(f["input_ids"]) for f in features) batch = {} for k in ["input_ids", "labels", "attention_mask"]: pad = tokenizer.pad_token_id if k != "labels" else -100 batch[k] = torch.tensor([f[k] + [pad]*(max_len-len(f[k])) for f in features]) + # teacher_logits需要特殊padding（用0填充） vocab_size = features[0]["teacher_logits"].size(-1) tl = [] 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): + """自定义Trainer：蒸馏loss + SFT loss混合""" + 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) - student_logits = outputs.logits + # 前向传播 + outputs = model( + input_ids=inputs["input_ids"], + attention_mask=inputs["attention_mask"], + use_cache=False, + ) + student_logits = outputs.logits # [batch, seq_len, vocab_size] + + # SFT loss (交叉熵，只计算assistant部分) 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") - 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") - kl_loss = (kl_loss * mask).sum() / mask.sum() * (TEMP ** 2) - return ALPHA * kl_loss + (1 - ALPHA) * sft_loss + 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 vs student） + teacher_logits = inputs["teacher_logits"] # [batch, seq_len, vocab_size] + # 只对assistant部分计算KL + mask = (inputs["labels"] != -100).unsqueeze(-1).float() # [batch, seq_len, 1] + + # 软化分布 + s_logits_soft = student_logits / TEMP + 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) # 温度缩放 + + # 混合loss + total_loss = ALPHA * kl_loss + (1 - ALPHA) * sft_loss + + return total_loss args = TrainingArguments( output_dir=OUT, num_train_epochs=EPOCHS, @@ -107,19 +179,51 @@ 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. 启动训练 ========== 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") +print(f" Temp={TEMP}, Alpha={ALPHA}, Eff batch={BS*GA}, LR={LR}") sys.stdout.flush() + trainer.train() -print("[6/6] Saving...") +# ========== 6. 保存 ========== +print("[6/6] Saving distilled model...") final = os.path.join(OUT, "final") trainer.save_model(final) tokenizer.save_pretrained(final) + +# ⚠️ 关键修复：Qwen chat template 使用 <|im_end|> (151645) 作为对话EOS +# 但默认 eos_token_id=151643 (<|endoftext|>) +# 不修复会导致部署时模型无限生成 → 死循环乱码 +# 注意：必须同时修复 config.json 和 generation_config.json！ +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（避免 "You are Qwen..."） +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!")