File size: 3,956 Bytes
796e424 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 | from __future__ import annotations
import argparse
from pathlib import Path
import torch
from datasets import load_dataset
from peft import LoraConfig
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig, TrainingArguments
from trl import SFTTrainer
def format_messages(example: dict, tokenizer) -> str:
messages = example["messages"]
if hasattr(tokenizer, "apply_chat_template") and tokenizer.chat_template:
return tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=False)
chunks = []
for msg in messages:
role = msg["role"]
content = msg["content"]
chunks.append(f"<|im_start|>{role}\n{content}<|im_end|>")
return "\n".join(chunks)
def main() -> int:
parser = argparse.ArgumentParser()
parser.add_argument("--model_id", default="openbmb/MiniCPM5-1B")
parser.add_argument("--train_file", type=Path, required=True)
parser.add_argument("--val_file", type=Path, required=True)
parser.add_argument("--output_dir", type=Path, required=True)
parser.add_argument("--max_seq_length", type=int, default=2048)
parser.add_argument("--epochs", type=float, default=2.0)
parser.add_argument("--learning_rate", type=float, default=2e-4)
parser.add_argument("--batch_size", type=int, default=2)
parser.add_argument("--grad_accum", type=int, default=8)
parser.add_argument("--lora_r", type=int, default=16)
parser.add_argument("--lora_alpha", type=int, default=32)
args = parser.parse_args()
dataset = load_dataset(
"json",
data_files={
"train": str(args.train_file),
"validation": str(args.val_file),
},
)
tokenizer = AutoTokenizer.from_pretrained(args.model_id, trust_remote_code=True)
if tokenizer.pad_token is None:
tokenizer.pad_token = tokenizer.eos_token
quant_config = BitsAndBytesConfig(
load_in_4bit=True,
bnb_4bit_compute_dtype=torch.bfloat16,
bnb_4bit_quant_type="nf4",
bnb_4bit_use_double_quant=True,
)
model = AutoModelForCausalLM.from_pretrained(
args.model_id,
quantization_config=quant_config,
device_map="auto",
trust_remote_code=True,
)
model.config.use_cache = False
peft_config = LoraConfig(
r=args.lora_r,
lora_alpha=args.lora_alpha,
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",
],
)
training_args = TrainingArguments(
output_dir=str(args.output_dir),
num_train_epochs=args.epochs,
per_device_train_batch_size=args.batch_size,
per_device_eval_batch_size=args.batch_size,
gradient_accumulation_steps=args.grad_accum,
learning_rate=args.learning_rate,
lr_scheduler_type="cosine",
warmup_ratio=0.05,
logging_steps=10,
eval_strategy="steps",
eval_steps=100,
save_steps=100,
save_total_limit=3,
bf16=True,
optim="paged_adamw_8bit",
report_to="none",
gradient_checkpointing=True,
)
def formatting_func(example):
return format_messages(example, tokenizer)
trainer = SFTTrainer(
model=model,
processing_class=tokenizer,
train_dataset=dataset["train"],
eval_dataset=dataset["validation"],
peft_config=peft_config,
formatting_func=formatting_func,
max_seq_length=args.max_seq_length,
args=training_args,
)
trainer.train()
trainer.save_model(str(args.output_dir))
tokenizer.save_pretrained(str(args.output_dir))
print(f"Saved LoRA adapter to {args.output_dir}")
return 0
if __name__ == "__main__":
raise SystemExit(main())
|