90 lines
3.1 KiB
Python
90 lines
3.1 KiB
Python
#!/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)
|