#!/usr/bin/env python3
"""光湖 母模型(7B)->1.5B霜砚 蒸馏 v6
修复：vocab_size不匹配 (teacher=152064, student=151936)
"""
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

TCH = '/root/autodl-tmp/output/qwen25-7b-sft/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-shuangyan-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')
CS = os.environ.get('ZY_OSS_SECRET')
os.makedirs(OUT, 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
print('[2/5] Load models...')
torch.cuda.empty_cache()  # Clear any stale memory before loading
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)

print('  Teacher...', 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
print('[3/5] Train...', 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
sch = torch.optim.lr_scheduler.CosineAnnealingLR(opt, T_max=total_steps // GA)
global_step = 0

for ep in range(E):
    print(f'  Epoch {ep+1}/{E}', flush=True)
    for step, batch in enumerate(tqdm(loader, desc=f'Ep{ep+1}')):
        global_step += 1
        # Teacher forward (on-the-fly)
        with torch.no_grad():
            t_logits = teacher(input_ids=batch['input_ids'], attention_mask=batch['attention_mask']).logits
        # Student forward
        s_logits = student(input_ids=batch['input_ids'], attention_mask=batch['attention_mask']).logits

        # ⚠️ FIX: Truncate teacher logits to student vocab size
        # Teacher(7B) vocab=152064, Student(1.5B) vocab=151936
        t_logits = t_logits[..., :student.config.vocab_size]

        shift_s = s_logits[..., :-1, :].contiguous()
        shift_l = batch['labels'][..., 1:].contiguous()
        sft = F.cross_entropy(shift_s.view(-1, shift_s.size(-1)),
            shift_l.view(-1), ignore_index=-100, reduction='mean')

        mask = (batch['labels'] != -100).unsqueeze(-1).float()
        kls = F.kl_div(F.log_softmax(s_logits.float()/TEMP, -1),
            F.softmax(t_logits.float()/TEMP, -1), reduction='none')
        kl = (kls * mask).sum() / mask.sum() * TEMP**2
        loss = (ALPHA*kl + (1-ALPHA)*sft) / GA
        loss.backward()

        if global_step % GA == 0:
            torch.nn.utils.clip_grad_norm_(student.parameters(), 1.0)
            opt.step(); opt.zero_grad(); sch.step()

        if global_step % 500 == 0:
            # Prevent memory fragmentation
            torch.cuda.empty_cache()

        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)

    ckpt = os.path.join(OUT, f'ep{ep+1}')
    os.makedirs(ckpt, exist_ok=True)
    student.save_pretrained(ckpt); tok.save_pretrained(ckpt)
    print(f'  Checkpoint: {ckpt}', flush=True)
    # FIX: Clear GPU cache between epochs to prevent fragmentation OOM
    torch.cuda.empty_cache()

# 4. Save
print('[4/5] Save final...', flush=True)
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. Upload
print('[5/5] Upload to COS...', flush=True)
cfg = CosConfig(Region=RG, SecretId=CK, SecretKey=CS)
cl = CosS3Client(cfg)
import glob
for fp in sorted(glob.glob(fnl + '/**/*', recursive=True)):
    if os.path.isfile(fp):
        name = os.path.relpath(fp, fnl)
        sz = os.path.getsize(fp) / 1024 / 1024
        print(f'  {name} ({sz:.0f}MB)', flush=True)
        cl.upload_file(Bucket=BKT, Key='models/shuangyan-15b-distill/' + name, LocalFilePath=fp)
print('ALL DONE')