📜 蒸馏脚本：finetune_shuangyan_v2.py

scripts/distill/finetune_shuangyan_v2.py

@@ -0,0 +1,89 @@
+#!/usr/bin/env python3
+"""
+小霜砚LoRA微调 v2 — 用sft_v2.jsonl (1868条霜砚对话)
+基座：1.5B蒸馏模板 (Track1)
+"""
+import os, json, torch
+os.environ['CUDA_VISIBLE_DEVICES'] = '0'
+os.environ['TOKENIZERS_PARALLELISM'] = 'false'
+
+BASE = '/root/autodl-tmp/output/qwen25-15b-shuangyan-distill/final'
+OUT = '/root/autodl-tmp/output/shuangyan-v2'
+DATA = '/root/autodl-tmp/data/sft_v2.jsonl'
+os.makedirs(OUT, exist_ok=True)
+
+print('[1/4] 加载语料...', flush=True)
+from datasets import Dataset
+from transformers import AutoTokenizer, AutoModelForCausalLM, TrainingArguments, Trainer
+from peft import LoraConfig, get_peft_model
+
+with open(DATA) as f:
+    raw = [json.loads(l) for l in f]
+# 去掉system message (cc-002)
+data = []
+for d in raw:
+    msgs = [m for m in d['messages'] if m['role'] != 'system']
+    if len(msgs) >= 2:
+        data.append({'messages': msgs})
+print(f'  加载 {len(data)} 条对话', flush=True)
+
+print('[2/4] 加载模型...', flush=True)
+tokenizer = AutoTokenizer.from_pretrained(BASE, trust_remote_code=True)
+tokenizer.pad_token = tokenizer.eos_token
+
+model = AutoModelForCausalLM.from_pretrained(
+    BASE, trust_remote_code=True, torch_dtype=torch.bfloat16,
+    attn_implementation='sdpa').cuda()
+model.config.use_cache = False
+
+lora = LoraConfig(r=16, lora_alpha=32,
+    target_modules=['q_proj','k_proj','v_proj','o_proj'],
+    lora_dropout=0.05, bias='none', task_type='CAUSAL_LM')
+model = get_peft_model(model, lora)
+model.print_trainable_parameters()
+
+print('[3/4] Tokenize...', flush=True)
+MAX_LEN = 2048
+processed = []
+for d in data:
+    ii, ll = [], []
+    for m in d['messages']:
+        c = m['content']
+        if not c.strip(): continue
+        txt = '<|im_start|>' + m['role'] + '\n' + c + '<|im_end|>\n'
+        tk = tokenizer.encode(txt, add_special_tokens=False)
+        ii.extend(tk)
+        ll.extend(tk if m['role'] == 'assistant' else [-100]*len(tk))
+    if len(ii) > MAX_LEN:
+        ii, ll = ii[:MAX_LEN], ll[:MAX_LEN]
+    if len(ii) > 10:
+        processed.append({'input_ids': ii, 'labels': ll, 'attention_mask': [1]*len(ii)})
+
+print(f'  {len(processed)} 条训练数据', flush=True)
+ds = Dataset.from_list(processed)
+
+def collate(features):
+    ml = 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]*(ml-len(f[k])) for f in features])
+    return batch
+
+args = TrainingArguments(
+    output_dir=OUT, num_train_epochs=5,
+    per_device_train_batch_size=2, gradient_accumulation_steps=4,
+    learning_rate=2e-4, warmup_ratio=0.05, lr_scheduler_type='cosine',
+    bf16=True, logging_steps=10, save_strategy='epoch', save_total_limit=2,
+    remove_unused_columns=False, report_to='none',
+    gradient_checkpointing=True, optim='adamw_torch',
+)
+
+print('[4/4] 开始微调！', flush=True)
+trainer = Trainer(model=model, args=args, train_dataset=ds, data_collator=collate)
+trainer.train()
+
+fnl = os.path.join(OUT, 'final')
+model.save_pretrained(fnl)
+tokenizer.save_pretrained(fnl)
+print(f'✅ 小霜砚微调完成！{fnl}', flush=True)