93 lines
3.2 KiB
Python
93 lines
3.2 KiB
Python
|
#!/usr/bin/env python3
|
|||
|
|
"""
|
|||
|
|
小铸渊LoRA微调 v1.0（线B）
|
|||
|
|
基座：1.5B代码蒸馏模板 (Track2)
|
|||
|
|
语料：zhuyuan_full_corpus.jsonl (铸渊对话+系统知识, 31条, 223K)
|
|||
|
|
"""
|
|||
|
|
import os, json, torch
|
|||
|
|
os.environ['CUDA_VISIBLE_DEVICES'] = '0'
|
|||
|
|
os.environ['TOKENIZERS_PARALLELISM'] = 'false'
|
|||
|
|
|
|||
|
|
# 基座模型用Track2的代码蒸馏模板
|
|||
|
|
BASE = '/root/autodl-tmp/output/qwen25-15b-coder-distill/final'
|
|||
|
|
OUT = '/root/autodl-tmp/output/zhuyuan-lora'
|
|||
|
|
DATA = '/root/autodl-tmp/corpus_work/zhuyuan_full_corpus.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 if l.strip()]
|
|||
|
|
# 去掉system message
|
|||
|
|
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
|
|||
|
|
|
|||
|
|
# Qwen2.5-Coder-1.5B蒸馏模板
|
|||
|
|
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=8,
|
|||
|
|
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)
|