210 lines
7.0 KiB
Python
210 lines
7.0 KiB
Python
#!/usr/bin/env python3
|
||
"""光湖代码模型→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
|
||
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-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" # 铸渊专属语料
|
||
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
|
||
|
||
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]
|
||
# 铸渊语料已经去除了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.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")
|
||
|
||
# ========== 3. Tokenize + Teacher Logits ==========
|
||
print("[3/6] Tokenizing + generating teacher logits...")
|
||
|
||
processed = []
|
||
for d in tqdm(raw, desc="Tokenize+Teacher"):
|
||
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
|
||
})
|
||
|
||
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 = {}
|
||
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)
|
||
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,
|
||
)
|
||
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",
|
||
)
|
||
|
||
# KL loss
|
||
teacher_logits = inputs["teacher_logits"]
|
||
mask = (inputs["labels"] != -100).unsqueeze(-1).float()
|
||
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(
|
||
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=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
|
||
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}")
|
||
print(f" Peak VRAM: {peak:.2f}GB / {mem:.1f}GB")
|
||
print("DONE!")
|