mirror of
https://github.com/allaunthefox/Research-Stack.git
synced 2026-07-31 03:05:21 +00:00
101 lines
3.6 KiB
Python
101 lines
3.6 KiB
Python
#!/usr/bin/env python3
|
|
"""Unsloth SFT scaffold for a physics/math routing LLM.
|
|
|
|
Use a trainable HF/safetensors Gemma-family checkpoint as --model-name. A GGUF
|
|
is a good deployment/teacher artifact, but this script expects a trainable
|
|
Transformers checkpoint because LoRA/QLoRA training needs model weights in that
|
|
ecosystem.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
from pathlib import Path
|
|
|
|
|
|
def main() -> int:
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--model-name", required=True, help="HF/Unsloth trainable model id or local safetensors dir")
|
|
parser.add_argument("--dataset", type=Path, default=Path("4-Infrastructure/shim/physics_math_llm_sft.jsonl"))
|
|
parser.add_argument("--out", type=Path, default=Path("4-Infrastructure/shim/physics_math_lora"))
|
|
parser.add_argument("--max-seq-length", type=int, default=4096)
|
|
parser.add_argument("--load-in-4bit", action="store_true")
|
|
parser.add_argument("--max-steps", type=int, default=120)
|
|
parser.add_argument("--learning-rate", type=float, default=2e-4)
|
|
parser.add_argument(
|
|
"--packing",
|
|
action=argparse.BooleanOptionalAction,
|
|
default=True,
|
|
help="Pack short SFT records to reduce padding waste; latest Unsloth/NVIDIA paths cache packed metadata.",
|
|
)
|
|
parser.add_argument("--dataset-num-proc", type=int, default=2)
|
|
args = parser.parse_args()
|
|
|
|
try:
|
|
from datasets import load_dataset
|
|
from trl import SFTTrainer, SFTConfig
|
|
from unsloth import FastLanguageModel
|
|
except ImportError as exc:
|
|
raise SystemExit(
|
|
"Missing training dependencies. Install Unsloth stack first, then rerun. "
|
|
"Expected: unsloth, trl, datasets."
|
|
) from exc
|
|
|
|
model, tokenizer = FastLanguageModel.from_pretrained(
|
|
model_name=str(args.model_name),
|
|
max_seq_length=args.max_seq_length,
|
|
dtype=None,
|
|
load_in_4bit=args.load_in_4bit,
|
|
)
|
|
model = FastLanguageModel.get_peft_model(
|
|
model,
|
|
r=16,
|
|
target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"],
|
|
lora_alpha=16,
|
|
lora_dropout=0,
|
|
bias="none",
|
|
use_gradient_checkpointing="unsloth",
|
|
random_state=3407,
|
|
)
|
|
|
|
dataset = load_dataset("json", data_files=str(args.dataset), split="train")
|
|
|
|
def formatting_prompts_func(examples):
|
|
texts = []
|
|
for messages in examples["messages"]:
|
|
texts.append(tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=False))
|
|
return {"text": texts}
|
|
|
|
dataset = dataset.map(formatting_prompts_func, batched=True)
|
|
|
|
trainer = SFTTrainer(
|
|
model=model,
|
|
tokenizer=tokenizer,
|
|
train_dataset=dataset,
|
|
args=SFTConfig(
|
|
output_dir=str(args.out),
|
|
dataset_text_field="text",
|
|
max_seq_length=args.max_seq_length,
|
|
packing=args.packing,
|
|
dataset_num_proc=args.dataset_num_proc,
|
|
per_device_train_batch_size=1,
|
|
gradient_accumulation_steps=8,
|
|
warmup_steps=5,
|
|
max_steps=args.max_steps,
|
|
learning_rate=args.learning_rate,
|
|
logging_steps=5,
|
|
save_steps=max(20, args.max_steps // 2),
|
|
optim="adamw_8bit",
|
|
seed=3407,
|
|
),
|
|
)
|
|
trainer.train()
|
|
args.out.mkdir(parents=True, exist_ok=True)
|
|
model.save_pretrained(str(args.out))
|
|
tokenizer.save_pretrained(str(args.out))
|
|
print(f"saved LoRA adapter to {args.out}")
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|