diff --git a/scripts/distill_coder.py b/scripts/distill_coder.py index 8e696d6..f97eb5d 100644 --- a/scripts/distill_coder.py +++ b/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") - kl_loss = (kl_loss * mask).sum() / mask.sum() * (TEMP ** 2) + 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) + 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!")