fix: 蒸馏保存段增加 config.json + generation_config.json + tokenizer 三重修复
This commit is contained in:
bingshuo
parent
9f95ee891d
commit
9e7bc2000a
@ -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!")
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user