diff --git a/scripts/distill/finetune_shuangyan_v2.py b/scripts/distill/finetune_shuangyan_v2.py new file mode 100644 index 0000000..3d58095 --- /dev/null +++ b/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)