import json, os, torch
from datasets import Dataset
from peft import LoraConfig, get_peft_model
from transformers import AutoModelForCausalLM, AutoTokenizer, set_seed
from trl import SFTConfig, SFTTrainer
BASE = os.environ["BASE_MODEL"]
DATA = os.environ["DATA"]
OUT = os.environ["OUTPUT_DIR"]
TRAIN_CHAT_TEMPLATE = (
"{{- bos_token }}"
"{%- for message in messages %}"
"{%- if message['role'] == 'system' %}"
"{{- '<|im_start|>system\\n' + message['content'] + '<|im_end|>\\n' }}"
"{%- elif message['role'] == 'user' %}"
"{{- '<|im_start|>user\\n' + message['content'] + '<|im_end|>\\n' }}"
"{%- elif message['role'] == 'assistant' %}"
"{{- '<|im_start|>assistant\\n' }}"
"{%- generation %}"
"{{- message['content'] + '<|im_end|>' }}"
"{%- endgeneration %}"
"{{- '\\n' }}"
"{%- endif %}"
"{%- endfor %}"
"{%- if add_generation_prompt %}"
"{{- '<|im_start|>assistant\\n' }}"
"{%- endif %}"
)
set_seed(42)
tok = AutoTokenizer.from_pretrained(BASE, use_fast=True)
if tok.pad_token is None:
tok.pad_token = tok.eos_token
tok.chat_template = TRAIN_CHAT_TEMPLATE
rows = [json.loads(l) for l in open(DATA, encoding="utf-8") if l.strip()]
ds = Dataset.from_list([{"messages": r["messages"]} for r in rows])
model = AutoModelForCausalLM.from_pretrained(BASE, torch_dtype=torch.bfloat16, attn_implementation="sdpa")
model.config.use_cache = False
model.gradient_checkpointing_enable(gradient_checkpointing_kwargs={"use_reentrant": False})
lora = LoraConfig(
r=16, lora_alpha=32, lora_dropout=0.05, bias="none", task_type="CAUSAL_LM",
target_modules=["q_proj","k_proj","v_proj","o_proj","gate_proj","up_proj","down_proj"],
)
model = get_peft_model(model, lora)
model.print_trainable_parameters()
trainer = SFTTrainer(
model=model,
args=SFTConfig(
output_dir=OUT,
num_train_epochs=2,
per_device_train_batch_size=4,
gradient_accumulation_steps=4,
learning_rate=2e-4,
warmup_ratio=0.03,
lr_scheduler_type="cosine",
bf16=True,
max_length=2048,
packing=False,
assistant_only_loss=True,
logging_steps=10,
save_steps=200,
save_total_limit=2,
report_to=["tensorboard"],
dataloader_num_workers=2,
remove_unused_columns=False,
seed=42,
),
train_dataset=ds,
processing_class=tok,
)
trainer.train()
trainer.model.save_pretrained(f"{OUT}/adapter_final")
trainable params: 11,206,656 || all params: 1,091,839,488 || trainable%: 1.0264
{'loss': 4.07, 'mean_token_accuracy': 0.29, 'epoch': 0.2}
{'loss': 3.52, 'mean_token_accuracy': 0.36, 'epoch': 1.0}
from peft import PeftModel
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch
base = AutoModelForCausalLM.from_pretrained("openbmb/MiniCPM5-1B", torch_dtype=torch.bfloat16, device_map="auto")
model = PeftModel.from_pretrained(base, "./runs/minicpm5_trl/adapter_final").eval()
tok = AutoTokenizer.from_pretrained("openbmb/MiniCPM5-1B")
inputs = tok.apply_chat_template([{"role":"user","content":"用一句话解释 GQA。"}],
add_generation_prompt=True, enable_thinking=True, return_tensors="pt").to(model.device)
out = model.generate(inputs, max_new_tokens=512, do_sample=True, temperature=0.9, top_p=0.95)
print(tok.decode(out[0][inputs.shape[-1]:], skip_special_tokens=True))
