diff --git a/scripts/distill/distill_coder.py b/scripts/distill/distill_coder.py new file mode 100644 index 0000000..8d580b4 --- /dev/null +++ b/scripts/distill/distill_coder.py @@ -0,0 +1,173 @@ +#!/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)