diff --git a/scripts/distill_mother.py b/scripts/distill_mother.py new file mode 100644 index 0000000..e772363 --- /dev/null +++ b/scripts/distill_mother.py @@ -0,0 +1,125 @@ +#!/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 & +""" + +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 +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" +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 + +os.makedirs(OUT, exist_ok=True) + +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") + +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 +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") + +print("[3/6] Tokenizing + generating teacher logits...") +processed = [] +for d in tqdm(raw, desc="Tokenize"): + ids, labs = [], [] + for msg in d["messages"]: + c = msg["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] + 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}) + +print("[4/6] Training config...") +def distill_collate(features): + 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]) + 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]) + 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) + student_logits = outputs.logits + 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 + +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, +) + +trainer = DistillTrainer(model=student, args=args, train_dataset=Dataset.from_list(processed), data_collator=distill_collate) + +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}") +sys.stdout.flush() +trainer.train() + +print("[6/6] Saving...") +final = os.path.join(OUT, "final") +trainer.save_model(final) +tokenizer.save_pretrained(final) +peak = torch.cuda.max_memory_allocated() / 1e9 +print(f" Model: {final} | Peak VRAM: {peak:.2f}GB / {mem:.1f}GB") +print("DONE!")