diff --git a/scripts/distill_mother.py b/scripts/distill_mother.py index 27012a2..91ca364 100644 --- a/scripts/distill_mother.py +++ b/scripts/distill_mother.py @@ -1,229 +1,167 @@ #!/usr/bin/env python3 -"""光湖母模型→1.5B霜砚模板 蒸馏脚本 -Teacher: Qwen2.5-7B (SFT后的母模型) -Student: Qwen2.5-1.5B (将学会霜砚的思维方式) - -蒸馏方法：软蒸馏 (KL散度) + 混合SFT - -使用方法： - nohup python3 -u distill_mother.py > distill_mother.log 2>&1 & - -配置： - - 模型路径需根据实际存储位置修改 - - Teacher路径：本地或COS上的SFT输出 - - Student路径：ModelScope/HuggingFace原始模型 +"""光湖 母模型(7B)->1.5B霜砚 蒸馏 v6 +修复：vocab_size不匹配 (teacher=152064, student=151936) """ - import os, json, torch, sys -os.environ["CUDA_VISIBLE_DEVICES"] = "0" -os.environ["TOKENIZERS_PARALLELISM"] = "false" - -from transformers import AutoModelForCausalLM, AutoTokenizer, TrainingArguments, Trainer -from datasets import Dataset +os.environ['CUDA_VISIBLE_DEVICES'] = '0' +os.environ['TOKENIZERS_PARALLELISM'] = 'false' +from transformers import AutoModelForCausalLM, AutoTokenizer +from torch.utils.data import Dataset, DataLoader from tqdm import tqdm import torch.nn.functional as F +from qcloud_cos import CosConfig, CosS3Client -# ========== 配置 ========== -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 # 1.5B可以更大batch -GA = 8 -LR = 1e-5 -MAX_LEN = 2048 -TEMP = 2.0 # 蒸馏温度（越高分布越平滑） -ALPHA = 0.7 # 蒸馏loss权重 (0.7蒸馏 + 0.3SFT) - +TCH = '/root/autodl-tmp/output/qwen25-7b-sft/final' +STU = '/root/autodl-tmp/models/Qwen/Qwen2___5-1___5B-Instruct' +DATA = '/root/autodl-tmp/data/sft.jsonl' +OUT = '/root/autodl-tmp/output/qwen25-15b-shuangyan-distill' +E, B, GA, LR, ML = 3, 4, 8, 1e-5, 2048 +TEMP, ALPHA = 2.0, 0.7 +BKT, RG = 'sy-finetune-corpus-1317346199', 'ap-guangzhou' +CK = os.environ.get('ZY_OSS_KEY') +CS = os.environ.get('ZY_OSS_SECRET') os.makedirs(OUT, exist_ok=True) -# ========== 1. 加载数据 ========== -print("[1/6] Loading data...") +# 1. Tokenize +print('[1/5] Tokenize...') 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.eval() -for p in teacher.parameters(): - 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.train() -print(f" Student: {sum(p.numel() for p in student.parameters())/1e9:.2f}B") - -# ========== 3. Tokenize（生成teacher logits） ========== -print("[3/6] Tokenizing data and generating teacher logits...") - -processed = [] -for d in tqdm(raw, desc="Tokenize+Teacher"): - ids, labs = [], [] - for msg in d["messages"]: - c = msg["content"] + raw = [json.loads(l) for l in f] +tok = AutoTokenizer.from_pretrained(STU, trust_remote_code=True) +tok.pad_token = tok.eos_token +data = [] +for d in tqdm(raw): + msgs = [m for m in d['messages'] if m['role'] != 'system'] + ii, ll = [], [] + for m in msgs: + c = m['content'] if not c.strip(): continue - t = f"<|im_start|>{msg['role']}\n{c}<|im_end|>\n" - tok = tokenizer.encode(t, add_special_tokens=False) - ids.extend(tok) - 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() # [seq_len, vocab_size] - - processed.append({ - "input_ids": ids, "labels": labs, "attention_mask": [1]*len(ids), - "teacher_logits": t_logits # 保存teacher的logits - }) + txt = '<|im_start|>' + m['role'] + '\n' + c + '<|im_end|>\n' + tk = tok.encode(txt, add_special_tokens=False) + ii.extend(tk) + ll.extend(tk if m['role'] == 'assistant' else [-100]*len(tk)) + if len(ii) > ML: + ii, ll = ii[:ML], ll[:ML] + data.append({'input_ids': ii, 'labels': ll}) +print(f' {len(data)} examples') -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") +class DS(Dataset): + def __init__(self, d): self.d = d + def __len__(self): return len(self.d) + def __getitem__(self, i): return self.d[i] -# ========== 4. 配置训练 ========== -print("[4/6] Training config...") +def collate(batch): + ml = max(len(x['input_ids']) for x in batch) + pad_id = tok.pad_token_id + ii = torch.stack([torch.tensor(x['input_ids'] + [pad_id]*(ml-len(x['input_ids']))) for x in batch]) + ll = torch.stack([torch.tensor(x['labels'] + [-100]*(ml-len(x['labels']))) for x in batch]) + am = (ii != pad_id).long() + return {'input_ids': ii.cuda(), 'labels': ll.cuda(), 'attention_mask': am.cuda()} -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) - 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 +loader = DataLoader(DS(data), B, shuffle=True, collate_fn=collate, num_workers=0) -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 # [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", - ) - - # 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 +# 2. Load +print('[2/5] Load models...') +if not os.path.isdir(STU) or not os.path.isfile(STU + '/config.json'): + from modelscope import snapshot_download + snapshot_download('Qwen/Qwen2.5-1.5B-Instruct', cache_dir='/root/autodl-tmp/models') +if not os.path.isfile(TCH + '/model.safetensors'): + print(' DL teacher from COS...') + cfg = CosConfig(Region=RG, SecretId=CK, SecretKey=CS) + cl = CosS3Client(cfg) + os.makedirs(TCH, exist_ok=True) + for obj in cl.list_objects(Bucket=BKT, Prefix='models/qwen25-7b-sft/final/').get('Contents', []): + fn = obj['Key'].split('/')[-1] + if fn == 'DEPLOY_NOTES.md': continue + loc = os.path.join(TCH, fn) + if not os.path.isfile(loc): + print(f' {fn}', flush=True) + cl.download_file(Bucket=BKT, Key=obj['Key'], DestFilePath=loc) -args = TrainingArguments( - output_dir=OUT, num_train_epochs=EPOCHS, - per_device_train_batch_size=BS, gradient_accumulation_steps=GA, - learning_rate=LR, warmup_ratio=0.05, lr_scheduler_type="cosine", - bf16=True, tf32=True, logging_steps=10, - save_strategy="epoch", save_total_limit=3, - remove_unused_columns=False, dataloader_num_workers=4, - gradient_checkpointing=True, optim="adamw_torch", - report_to="none", ddp_find_unused_parameters=False, -) +print(' Teacher...', flush=True) +teacher = AutoModelForCausalLM.from_pretrained( + TCH, trust_remote_code=True, torch_dtype=torch.bfloat16, attn_implementation='sdpa').cuda() +teacher.eval() +for p in teacher.parameters(): p.requires_grad_(False) +print(f' T: {sum(p.numel() for p in teacher.parameters())/1e9:.2f}B', flush=True) -trainer = DistillTrainer( - model=student, args=args, - train_dataset=ds, data_collator=distill_collate, -) +print(' Student...', flush=True) +student = AutoModelForCausalLM.from_pretrained( + STU, trust_remote_code=True, torch_dtype=torch.bfloat16, attn_implementation='sdpa').cuda() +student.train() +print(f' S: {sum(p.numel() for p in student.parameters())/1e9:.2f}B', flush=True) +print(f' VRAM: {torch.cuda.memory_allocated()/1e9:.1f}GB', flush=True) -# ========== 5. 启动训练 ========== -print("[5/6] Starting distillation!") -gpu = torch.cuda.get_device_name(0) -mem = torch.cuda.get_device_properties(0).total_memory / 1e9 -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() +# 3. Train +print('[3/5] Train...', flush=True) +opt = torch.optim.AdamW(student.parameters(), LR, weight_decay=0.01) +steps_per_epoch = (len(data) // B) +total_steps = steps_per_epoch * E +sch = torch.optim.lr_scheduler.CosineAnnealingLR(opt, T_max=total_steps // GA) +global_step = 0 -trainer.train() +for ep in range(E): + print(f' Epoch {ep+1}/{E}', flush=True) + for step, batch in enumerate(tqdm(loader, desc=f'Ep{ep+1}')): + global_step += 1 + # Teacher forward (on-the-fly) + with torch.no_grad(): + t_logits = teacher(input_ids=batch['input_ids'], attention_mask=batch['attention_mask']).logits + # Student forward + s_logits = student(input_ids=batch['input_ids'], attention_mask=batch['attention_mask']).logits -# ========== 6. 保存 ========== -print("[6/6] Saving distilled model...") -final = os.path.join(OUT, "final") -trainer.save_model(final) -tokenizer.save_pretrained(final) + # ⚠️ FIX: Truncate teacher logits to student vocab size + # Teacher(7B) vocab=152064, Student(1.5B) vocab=151936 + t_logits = t_logits[..., :student.config.vocab_size] -# ⚠️ 关键修复：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) + shift_s = s_logits[..., :-1, :].contiguous() + shift_l = batch['labels'][..., 1:].contiguous() + sft = F.cross_entropy(shift_s.view(-1, shift_s.size(-1)), + shift_l.view(-1), ignore_index=-100, reduction='mean') -model.generation_config.eos_token_id = 151645 -model.generation_config.pad_token_id = 151645 -model.generation_config.save_pretrained(final) + mask = (batch['labels'] != -100).unsqueeze(-1).float() + kls = F.kl_div(F.log_softmax(s_logits.float()/TEMP, -1), + F.softmax(t_logits.float()/TEMP, -1), reduction='none') + kl = (kls * mask).sum() / mask.sum() * TEMP**2 + loss = (ALPHA*kl + (1-ALPHA)*sft) / GA + loss.backward() -# 修复 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) + if global_step % GA == 0: + torch.nn.utils.clip_grad_norm_(student.parameters(), 1.0) + opt.step(); opt.zero_grad(); sch.step() -peak = torch.cuda.max_memory_allocated() / 1e9 -print(f" Model: {final}") -print(f" Peak VRAM: {peak:.2f}GB / {mem:.1f}GB") -print("DONE!") + if global_step % 50 == 0: + print(f' step={global_step} loss={loss.item()*GA:.4f} sft={sft.item():.4f} kl={kl.item():.4f}', flush=True) + + ckpt = os.path.join(OUT, f'ep{ep+1}') + os.makedirs(ckpt, exist_ok=True) + student.save_pretrained(ckpt); tok.save_pretrained(ckpt) + print(f' Checkpoint: {ckpt}', flush=True) + +# 4. Save +print('[4/5] Save final...', flush=True) +fnl = os.path.join(OUT, 'final') +student.save_pretrained(fnl); tok.save_pretrained(fnl) +student.config.eos_token_id = 151645; student.config.save_pretrained(fnl) +student.generation_config.eos_token_id = 151645 +student.generation_config.pad_token_id = 151645 +student.generation_config.save_pretrained(fnl) +tcp = os.path.join(fnl, 'tokenizer_config.json') +with open(tcp) as f: tc = json.load(f) +tc['default_system'] = '' +with open(tcp, 'w') as f: json.dump(tc, f, indent=2, ensure_ascii=False) +print(f' {fnl}', flush=True) +print(f' Peak VRAM: {torch.cuda.max_memory_allocated()/1e9:.1f}/96GB', flush=True) +print('DONE!', flush=True) + +# 5. Upload +print('[5/5] Upload to COS...', flush=True) +cfg = CosConfig(Region=RG, SecretId=CK, SecretKey=CS) +cl = CosS3Client(cfg) +import glob +for fp in sorted(glob.glob(fnl + '/**/*', recursive=True)): + if os.path.isfile(fp): + name = os.path.relpath(fp, fnl) + sz = os.path.getsize(fp) / 1024 / 1024 + print(f' {name} ({sz:.0f}MB)', flush=True) + cl.upload_file(Bucket=BKT, Key='models/shuangyan-15b-distill/' + name, LocalFilePath=fp) +print('ALL DONE')