madcat
98 lines
2.2 KiB
Python
98 lines
2.2 KiB
Python
"""BT-7274 LoRA v3 — balanced tool-call + direct-response.
|
|
|
|
Usage:
|
|
source ~/lora-train/bin/activate
|
|
cd ~/Projects/lora
|
|
python train_v3.py
|
|
"""
|
|
|
|
from unsloth import FastLanguageModel
|
|
from trl import SFTTrainer
|
|
from transformers import TrainingArguments
|
|
from datasets import load_dataset
|
|
|
|
# --- Config ---
|
|
MODEL = "unsloth/Qwen2.5-7B-Instruct-bnb-4bit"
|
|
MAX_SEQ = 4096
|
|
RANK = 16
|
|
ALPHA = 16
|
|
DATA = "./bt7274_v3.jsonl"
|
|
OUT = "./bt7274-lora-v3"
|
|
EPOCHS = 3
|
|
BATCH = 1
|
|
GRAD_ACCUM = 8
|
|
LR = 2e-4
|
|
|
|
# --- Load model (4-bit quantized) ---
|
|
model, tokenizer = FastLanguageModel.from_pretrained(
|
|
model_name=MODEL,
|
|
max_seq_length=MAX_SEQ,
|
|
load_in_4bit=True,
|
|
dtype=None,
|
|
)
|
|
|
|
# --- LoRA adapter ---
|
|
model = FastLanguageModel.get_peft_model(
|
|
model,
|
|
r=RANK,
|
|
lora_alpha=ALPHA,
|
|
lora_dropout=0,
|
|
target_modules=[
|
|
"q_proj", "k_proj", "v_proj", "o_proj",
|
|
"gate_proj", "up_proj", "down_proj",
|
|
],
|
|
bias="none",
|
|
use_gradient_checkpointing="unsloth",
|
|
random_state=42,
|
|
)
|
|
|
|
# --- Dataset ---
|
|
ds = load_dataset("json", data_files=DATA, split="train")
|
|
|
|
def to_chatml(ex):
|
|
text = tokenizer.apply_chat_template(
|
|
ex["messages"], tokenize=False, add_generation_prompt=False
|
|
)
|
|
return {"text": text}
|
|
|
|
ds = ds.map(to_chatml)
|
|
|
|
steps = (len(ds) * EPOCHS) // (BATCH * GRAD_ACCUM)
|
|
print(f"Dataset: {len(ds)} examples")
|
|
print(f"Epochs: {EPOCHS}, effective batch: {BATCH * GRAD_ACCUM}")
|
|
print(f"Estimated steps: {steps}")
|
|
print(f"LoRA: r={RANK}, alpha={ALPHA}")
|
|
print(f"Max seq: {MAX_SEQ}")
|
|
|
|
# --- Train ---
|
|
trainer = SFTTrainer(
|
|
model=model,
|
|
tokenizer=tokenizer,
|
|
train_dataset=ds,
|
|
args=TrainingArguments(
|
|
output_dir=OUT,
|
|
per_device_train_batch_size=BATCH,
|
|
gradient_accumulation_steps=GRAD_ACCUM,
|
|
num_train_epochs=EPOCHS,
|
|
learning_rate=LR,
|
|
bf16=True,
|
|
logging_steps=5,
|
|
save_steps=50,
|
|
save_total_limit=2,
|
|
warmup_steps=10,
|
|
optim="adamw_8bit",
|
|
seed=42,
|
|
report_to="none",
|
|
),
|
|
dataset_text_field="text",
|
|
max_seq_length=MAX_SEQ,
|
|
packing=False,
|
|
)
|
|
|
|
trainer.train()
|
|
|
|
# --- Save LoRA adapter ---
|
|
model.save_pretrained(OUT)
|
|
tokenizer.save_pretrained(OUT)
|
|
print(f"\nSaved LoRA adapter to {OUT}/")
|