174 lines
6.6 KiB
Python
174 lines
6.6 KiB
Python
|
#!/usr/bin/env python3
|
|||
|
|
"""光湖 代码模型(7B)->1.5B铸渊蒸馏 v1
|
|||
|
|
teacher: Qwen2.5-Coder-7B (COS下载)
|
|||
|
|
student: Qwen2.5-1.5B-Instruct
|
|||
|
|
data: sft.jsonl
|
|||
|
|
"""
|
|||
|
|
import os, json, torch, sys
|
|||
|
|
os.environ['CUDA_VISIBLE_DEVICES'] = '0'
|
|||
|
|
os.environ['TOKENIZERS_PARALLELISM'] = 'false'
|
|||
|
|
os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'expandable_segments:True'
|
|||
|
|
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路径（从COS下载）
|
|||
|
|
TCH = '/root/autodl-tmp/output/coder-7b-sft-cache/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-coder-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', 'AKIDkQuBQhoiS2OYXWebXLwMbdT7cvAScbbU')
|
|||
|
|
CS = os.environ.get('ZY_OSS_SECRET', 'nPoZKArgUJBA4nJenjSxJSQBj5FCj3A4')
|
|||
|
|
os.makedirs(OUT, exist_ok=True)
|
|||
|
|
os.makedirs(TCH, exist_ok=True)
|
|||
|
|
|
|||
|
|
# 1. Tokenize
|
|||
|
|
print('[1/5] Tokenize...')
|
|||
|
|
with open(DATA) as f:
|
|||
|
|
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
|
|||
|
|
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')
|
|||
|
|
|
|||
|
|
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]
|
|||
|
|
|
|||
|
|
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()}
|
|||
|
|
|
|||
|
|
loader = DataLoader(DS(data), B, shuffle=True, collate_fn=collate, num_workers=0)
|
|||
|
|
|
|||
|
|
# 2. Load teacher from COS
|
|||
|
|
print('[2/5] Load models...')
|
|||
|
|
torch.cuda.empty_cache()
|
|||
|
|
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 (code model) from COS...')
|
|||
|
|
cfg = CosConfig(Region=RG, SecretId=CK, SecretKey=CS)
|
|||
|
|
cl = CosS3Client(cfg)
|
|||
|
|
for obj in cl.list_objects(Bucket=BKT, Prefix='models/qwen25-coder-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' Downloading {fn}...', flush=True)
|
|||
|
|
cl.download_file(Bucket=BKT, Key=obj['Key'], DestFilePath=loc)
|
|||
|
|
|
|||
|
|
print(' Teacher (code model)...', 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)
|
|||
|
|
|
|||
|
|
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)
|
|||
|
|
|
|||
|
|
# 3. Train (knowledge distillation)
|
|||
|
|
print('[3/5] Train (KD)...', 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
|
|||
|
|
global_step = 0
|
|||
|
|
|
|||
|
|
for ep in range(E):
|
|||
|
|
print(f'\n===== Epoch {ep+1}/{E} =====', flush=True)
|
|||
|
|
loader.dataset.d = data # reshuffle won't hurt
|
|||
|
|
for batch in tqdm(loader, total=steps_per_epoch):
|
|||
|
|
with torch.no_grad():
|
|||
|
|
t_logits = teacher(**batch).logits
|
|||
|
|
# Fix vocab mismatch: teacher(152064) vs student(151936)
|
|||
|
|
t_logits = t_logits[:, :, :151936]
|
|||
|
|
s_logits = student(**batch).logits
|
|||
|
|
|
|||
|
|
# SFT loss
|
|||
|
|
sft_loss = F.cross_entropy(
|
|||
|
|
s_logits.view(-1, s_logits.size(-1)),
|
|||
|
|
batch['labels'].view(-1),
|
|||
|
|
ignore_index=-100, reduction='mean')
|
|||
|
|
|
|||
|
|
# KL divergence (distillation loss) — use cached t_logits
|
|||
|
|
T = TEMP
|
|||
|
|
t_prob = F.log_softmax(s_logits[:, :-1] / T, dim=-1)
|
|||
|
|
s_prob = F.softmax(t_logits[:, :-1] / T, dim=-1)
|
|||
|
|
# Only on non-pad tokens
|
|||
|
|
mask = (batch['input_ids'][:, 1:] != tok.pad_token_id).unsqueeze(-1)
|
|||
|
|
kl_loss = (F.kl_div(t_prob, s_prob, reduction='none') * mask).sum() / mask.sum() * (T ** 2)
|
|||
|
|
|
|||
|
|
loss = sft_loss + ALPHA * kl_loss
|
|||
|
|
loss.backward()
|
|||
|
|
|
|||
|
|
if (global_step + 1) % GA == 0:
|
|||
|
|
torch.nn.utils.clip_grad_norm_(student.parameters(), 1.0)
|
|||
|
|
opt.step()
|
|||
|
|
opt.zero_grad()
|
|||
|
|
|
|||
|
|
if global_step % 50 == 0:
|
|||
|
|
print(f' step={global_step} loss={loss.item():.4f} sft={sft_loss.item():.4f} kl={kl_loss.item():.4f}', flush=True)
|
|||
|
|
|
|||
|
|
# Periodic cache clear
|
|||
|
|
if global_step % 500 == 0 and global_step > 0:
|
|||
|
|
torch.cuda.empty_cache()
|
|||
|
|
|
|||
|
|
global_step += 1
|
|||
|
|
|
|||
|
|
# Save epoch checkpoint
|
|||
|
|
ckpt = os.path.join(OUT, f'ep{ep+1}')
|
|||
|
|
student.save_pretrained(ckpt)
|
|||
|
|
tok.save_pretrained(ckpt)
|
|||
|
|
torch.cuda.empty_cache()
|
|||
|
|
print(f' Checkpoint: {ckpt}', flush=True)
|
|||
|
|
|
|||
|
|
# 4. Save final
|
|||
|
|
print('[4/5] Save final...', flush=True)
|
|||
|
|
fnl = os.path.join(OUT, 'final')
|
|||
|
|
student.save_pretrained(fnl)
|
|||
|
|
tok.save_pretrained(fnl)
|
|||
|
|
print(f' Final: {fnl}', flush=True)
|
|||
|
|
|
|||
|
|
# 5. Upload to COS
|
|||
|
|
print('[5/5] Upload to COS...', flush=True)
|
|||
|
|
cfg = CosConfig(Region=RG, SecretId=CK, SecretKey=CS)
|
|||
|
|
cl = CosS3Client(cfg)
|
|||
|
|
for f in sorted(os.listdir(fnl)):
|
|||
|
|
fp = os.path.join(fnl, f)
|
|||
|
|
if os.path.isfile(fp):
|
|||
|
|
mb = os.path.getsize(fp) / 1024 / 1024
|
|||
|
|
print(f' {f} ({mb:.0f}MB)', flush=True)
|
|||
|
|
cl.upload_file(Bucket=BKT, Key=f'models/qwen25-15b-coder-distill/{f}', LocalFilePath=fp)
|
|||
|
|
|
|||
|
|
print('ALL DONE', flush=True)
|