distill_mother.py v6: fix vocab_size mismatch (7B=152064, 1.5B=151936)
This commit is contained in:
bingshuo
parent
ec8939f5c0
commit
7d53003fb1
@ -1,229 +1,167 @@
|
|||||||
#!/usr/bin/env python3
|
#!/usr/bin/env python3
|
||||||
"""光湖母模型→1.5B霜砚模板 蒸馏脚本
|
"""光湖 母模型(7B)->1.5B霜砚 蒸馏 v6
|
||||||
Teacher: Qwen2.5-7B (SFT后的母模型)
|
修复：vocab_size不匹配 (teacher=152064, student=151936)
|
||||||
Student: Qwen2.5-1.5B (将学会霜砚的思维方式)
|
|
||||||
|
|
||||||
蒸馏方法：软蒸馏 (KL散度) + 混合SFT
|
|
||||||
|
|
||||||
使用方法：
|
|
||||||
nohup python3 -u distill_mother.py > distill_mother.log 2>&1 &
|
|
||||||
|
|
||||||
配置：
|
|
||||||
- 模型路径需根据实际存储位置修改
|
|
||||||
- Teacher路径：本地或COS上的SFT输出
|
|
||||||
- Student路径：ModelScope/HuggingFace原始模型
|
|
||||||
"""
|
"""
|
||||||
|
|
||||||
import os, json, torch, sys
|
import os, json, torch, sys
|
||||||
os.environ["CUDA_VISIBLE_DEVICES"] = "0"
|
os.environ['CUDA_VISIBLE_DEVICES'] = '0'
|
||||||
os.environ["TOKENIZERS_PARALLELISM"] = "false"
|
os.environ['TOKENIZERS_PARALLELISM'] = 'false'
|
||||||
|
from transformers import AutoModelForCausalLM, AutoTokenizer
|
||||||
from transformers import AutoModelForCausalLM, AutoTokenizer, TrainingArguments, Trainer
|
from torch.utils.data import Dataset, DataLoader
|
||||||
from datasets import Dataset
|
|
||||||
from tqdm import tqdm
|
from tqdm import tqdm
|
||||||
import torch.nn.functional as F
|
import torch.nn.functional as F
|
||||||
|
from qcloud_cos import CosConfig, CosS3Client
|
||||||
|
|
||||||
# ========== 配置 ==========
|
TCH = '/root/autodl-tmp/output/qwen25-7b-sft/final'
|
||||||
TEACHER_PATH = "/root/autodl-tmp/output/qwen25-7b-sft/final" # 母模型SFT输出
|
STU = '/root/autodl-tmp/models/Qwen/Qwen2___5-1___5B-Instruct'
|
||||||
STUDENT_PATH = "/root/autodl-tmp/cache/Qwen/Qwen2___5-1___5B" # 1.5B学生
|
DATA = '/root/autodl-tmp/data/sft.jsonl'
|
||||||
DATA = "/root/autodl-tmp/data/sft.jsonl" # 主语料（也可用shuangyan专属语料）
|
OUT = '/root/autodl-tmp/output/qwen25-15b-shuangyan-distill'
|
||||||
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
|
||||||
EPOCHS = 3
|
BKT, RG = 'sy-finetune-corpus-1317346199', 'ap-guangzhou'
|
||||||
BS = 4 # 1.5B可以更大batch
|
CK = os.environ.get('ZY_OSS_KEY')
|
||||||
GA = 8
|
CS = os.environ.get('ZY_OSS_SECRET')
|
||||||
LR = 1e-5
|
|
||||||
MAX_LEN = 2048
|
|
||||||
TEMP = 2.0 # 蒸馏温度（越高分布越平滑）
|
|
||||||
ALPHA = 0.7 # 蒸馏loss权重 (0.7蒸馏 + 0.3SFT)
|
|
||||||
|
|
||||||
os.makedirs(OUT, exist_ok=True)
|
os.makedirs(OUT, exist_ok=True)
|
||||||
|
|
||||||
# ========== 1. 加载数据 ==========
|
# 1. Tokenize
|
||||||
print("[1/6] Loading data...")
|
print('[1/5] Tokenize...')
|
||||||
with open(DATA) as f:
|
with open(DATA) as f:
|
||||||
raw = [json.loads(line) for line in f]
|
raw = [json.loads(l) for l in f]
|
||||||
raw = [{"messages": [m for m in obj["messages"] if m["role"] != "system"]} for obj in raw]
|
tok = AutoTokenizer.from_pretrained(STU, trust_remote_code=True)
|
||||||
print(f" {len(raw)} examples")
|
tok.pad_token = tok.eos_token
|
||||||
|
data = []
|
||||||
# ========== 2. 加载Teacher + Student ==========
|
for d in tqdm(raw):
|
||||||
print("[2/6] Loading teacher (7B) and student (1.5B)...")
|
msgs = [m for m in d['messages'] if m['role'] != 'system']
|
||||||
|
ii, ll = [], []
|
||||||
tokenizer = AutoTokenizer.from_pretrained(STUDENT_PATH, trust_remote_code=True)
|
for m in msgs:
|
||||||
tokenizer.pad_token = tokenizer.eos_token
|
c = m['content']
|
||||||
|
|
||||||
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"]
|
|
||||||
if not c.strip(): continue
|
if not c.strip(): continue
|
||||||
t = f"<|im_start|>{msg['role']}\n{c}<|im_end|>\n"
|
txt = '<|im_start|>' + m['role'] + '\n' + c + '<|im_end|>\n'
|
||||||
tok = tokenizer.encode(t, add_special_tokens=False)
|
tk = tok.encode(txt, add_special_tokens=False)
|
||||||
ids.extend(tok)
|
ii.extend(tk)
|
||||||
labs.extend(tok if msg["role"] == "assistant" else [-100] * len(tok))
|
ll.extend(tk if m['role'] == 'assistant' else [-100]*len(tk))
|
||||||
if len(ids) > MAX_LEN:
|
if len(ii) > ML:
|
||||||
ids, labs = ids[:MAX_LEN], labs[:MAX_LEN]
|
ii, ll = ii[:ML], ll[:ML]
|
||||||
|
data.append({'input_ids': ii, 'labels': ll})
|
||||||
|
print(f' {len(data)} examples')
|
||||||
|
|
||||||
# Teacher生成logits（蒸馏目标）
|
class DS(Dataset):
|
||||||
with torch.no_grad():
|
def __init__(self, d): self.d = d
|
||||||
inp = torch.tensor([ids]).cuda()
|
def __len__(self): return len(self.d)
|
||||||
t_out = teacher(input_ids=inp)
|
def __getitem__(self, i): return self.d[i]
|
||||||
t_logits = t_out.logits[0].float().cpu() # [seq_len, vocab_size]
|
|
||||||
|
|
||||||
processed.append({
|
def collate(batch):
|
||||||
"input_ids": ids, "labels": labs, "attention_mask": [1]*len(ids),
|
ml = max(len(x['input_ids']) for x in batch)
|
||||||
"teacher_logits": t_logits # 保存teacher的logits
|
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()}
|
||||||
|
|
||||||
ds = Dataset.from_list(processed)
|
loader = DataLoader(DS(data), B, shuffle=True, collate_fn=collate, num_workers=0)
|
||||||
total_tok = sum(len(d["input_ids"]) for d in processed)
|
|
||||||
print(f" Dataset: {len(ds)} ex, {total_tok:,} tokens")
|
|
||||||
|
|
||||||
# ========== 4. 配置训练 ==========
|
# 2. Load
|
||||||
print("[4/6] Training config...")
|
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)
|
||||||
|
|
||||||
def distill_collate(features):
|
print(' Teacher...', flush=True)
|
||||||
"""collate函数：处理padding + 蒸馏loss计算"""
|
teacher = AutoModelForCausalLM.from_pretrained(
|
||||||
max_len = max(len(f["input_ids"]) for f in features)
|
TCH, trust_remote_code=True, torch_dtype=torch.bfloat16, attn_implementation='sdpa').cuda()
|
||||||
batch = {}
|
teacher.eval()
|
||||||
for k in ["input_ids", "labels", "attention_mask"]:
|
for p in teacher.parameters(): p.requires_grad_(False)
|
||||||
pad = tokenizer.pad_token_id if k != "labels" else -100
|
print(f' T: {sum(p.numel() for p in teacher.parameters())/1e9:.2f}B', flush=True)
|
||||||
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
|
|
||||||
|
|
||||||
class DistillTrainer(Trainer):
|
print(' Student...', flush=True)
|
||||||
"""自定义Trainer：蒸馏loss + SFT loss混合"""
|
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)
|
||||||
|
|
||||||
def compute_loss(self, model, inputs, return_outputs=False, **kwargs):
|
# 3. Train
|
||||||
# 前向传播
|
print('[3/5] Train...', flush=True)
|
||||||
outputs = model(
|
opt = torch.optim.AdamW(student.parameters(), LR, weight_decay=0.01)
|
||||||
input_ids=inputs["input_ids"],
|
steps_per_epoch = (len(data) // B)
|
||||||
attention_mask=inputs["attention_mask"],
|
total_steps = steps_per_epoch * E
|
||||||
use_cache=False,
|
sch = torch.optim.lr_scheduler.CosineAnnealingLR(opt, T_max=total_steps // GA)
|
||||||
)
|
global_step = 0
|
||||||
student_logits = outputs.logits # [batch, seq_len, vocab_size]
|
|
||||||
|
|
||||||
# SFT loss (交叉熵，只计算assistant部分)
|
for ep in range(E):
|
||||||
shift_logits = student_logits[..., :-1, :].contiguous()
|
print(f' Epoch {ep+1}/{E}', flush=True)
|
||||||
shift_labels = inputs["labels"][..., 1:].contiguous()
|
for step, batch in enumerate(tqdm(loader, desc=f'Ep{ep+1}')):
|
||||||
sft_loss = F.cross_entropy(
|
global_step += 1
|
||||||
shift_logits.view(-1, shift_logits.size(-1)),
|
# Teacher forward (on-the-fly)
|
||||||
shift_labels.view(-1),
|
with torch.no_grad():
|
||||||
ignore_index=-100,
|
t_logits = teacher(input_ids=batch['input_ids'], attention_mask=batch['attention_mask']).logits
|
||||||
reduction="mean",
|
# Student forward
|
||||||
)
|
s_logits = student(input_ids=batch['input_ids'], attention_mask=batch['attention_mask']).logits
|
||||||
|
|
||||||
# KL蒸馏loss（teacher vs student）
|
# ⚠️ FIX: Truncate teacher logits to student vocab size
|
||||||
teacher_logits = inputs["teacher_logits"] # [batch, seq_len, vocab_size]
|
# Teacher(7B) vocab=152064, Student(1.5B) vocab=151936
|
||||||
# 只对assistant部分计算KL
|
t_logits = t_logits[..., :student.config.vocab_size]
|
||||||
mask = (inputs["labels"] != -100).unsqueeze(-1).float() # [batch, seq_len, 1]
|
|
||||||
|
|
||||||
# 软化分布
|
shift_s = s_logits[..., :-1, :].contiguous()
|
||||||
s_logits_soft = student_logits / TEMP
|
shift_l = batch['labels'][..., 1:].contiguous()
|
||||||
t_logits_soft = teacher_logits / TEMP
|
sft = F.cross_entropy(shift_s.view(-1, shift_s.size(-1)),
|
||||||
|
shift_l.view(-1), ignore_index=-100, reduction='mean')
|
||||||
|
|
||||||
kl_loss = F.kl_div(
|
mask = (batch['labels'] != -100).unsqueeze(-1).float()
|
||||||
F.log_softmax(s_logits_soft, dim=-1),
|
kls = F.kl_div(F.log_softmax(s_logits.float()/TEMP, -1),
|
||||||
F.softmax(t_logits_soft, dim=-1),
|
F.softmax(t_logits.float()/TEMP, -1), reduction='none')
|
||||||
reduction="none",
|
kl = (kls * mask).sum() / mask.sum() * TEMP**2
|
||||||
)
|
loss = (ALPHA*kl + (1-ALPHA)*sft) / GA
|
||||||
kl_loss = (kl_loss * mask).sum() / mask.sum()
|
loss.backward()
|
||||||
kl_loss = kl_loss * (TEMP ** 2) # 温度缩放
|
|
||||||
|
|
||||||
# 混合loss
|
if global_step % GA == 0:
|
||||||
total_loss = ALPHA * kl_loss + (1 - ALPHA) * sft_loss
|
torch.nn.utils.clip_grad_norm_(student.parameters(), 1.0)
|
||||||
|
opt.step(); opt.zero_grad(); sch.step()
|
||||||
|
|
||||||
return total_loss
|
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)
|
||||||
|
|
||||||
args = TrainingArguments(
|
ckpt = os.path.join(OUT, f'ep{ep+1}')
|
||||||
output_dir=OUT, num_train_epochs=EPOCHS,
|
os.makedirs(ckpt, exist_ok=True)
|
||||||
per_device_train_batch_size=BS, gradient_accumulation_steps=GA,
|
student.save_pretrained(ckpt); tok.save_pretrained(ckpt)
|
||||||
learning_rate=LR, warmup_ratio=0.05, lr_scheduler_type="cosine",
|
print(f' Checkpoint: {ckpt}', flush=True)
|
||||||
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(
|
# 4. Save
|
||||||
model=student, args=args,
|
print('[4/5] Save final...', flush=True)
|
||||||
train_dataset=ds, data_collator=distill_collate,
|
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. 启动训练 ==========
|
# 5. Upload
|
||||||
print("[5/6] Starting distillation!")
|
print('[5/5] Upload to COS...', flush=True)
|
||||||
gpu = torch.cuda.get_device_name(0)
|
cfg = CosConfig(Region=RG, SecretId=CK, SecretKey=CS)
|
||||||
mem = torch.cuda.get_device_properties(0).total_memory / 1e9
|
cl = CosS3Client(cfg)
|
||||||
t_params = sum(p.numel() for p in teacher.parameters())
|
import glob
|
||||||
s_params = sum(p.numel() for p in student.parameters())
|
for fp in sorted(glob.glob(fnl + '/**/*', recursive=True)):
|
||||||
print(f" GPU: {gpu} ({mem:.1f}GB)")
|
if os.path.isfile(fp):
|
||||||
print(f" Teacher: {t_params/1e9:.2f}B | Student: {s_params/1e9:.2f}B")
|
name = os.path.relpath(fp, fnl)
|
||||||
print(f" Temp={TEMP}, Alpha={ALPHA}, Eff batch={BS*GA}, LR={LR}")
|
sz = os.path.getsize(fp) / 1024 / 1024
|
||||||
sys.stdout.flush()
|
print(f' {name} ({sz:.0f}MB)', flush=True)
|
||||||
|
cl.upload_file(Bucket=BKT, Key='models/shuangyan-15b-distill/' + name, LocalFilePath=fp)
|
||||||
trainer.train()
|
print('ALL DONE')
|
||||||
|
|
||||||
# ========== 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}")
|
|
||||||
print(f" Peak VRAM: {peak:.2f}GB / {mem:.1f}GB")
|
|
||||||
print("DONE!")
|
|
||||||
|
|||||||
Loading…
x
Reference in New Issue
Block a user