File size: 3,675 Bytes
4d8daa4 6380390 4d8daa4 6380390 4d8daa4 | 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 | """CLERK LoRA SFT training.
Trains the consolidation policy (and the ledger-reading answer head) via
LoRA on a small instruct model. Written against the verified API surface of
``huggingface/trl`` examples/sft_qlora/sft_qlora.ipynb (SFTTrainer +
SFTConfig + LoraConfig, dataset with a single ``messages`` column,
``max_length`` for sequence length, ``push_to_hub`` for Hub upload).
Usage:
python -m clerk.train_sft --train data/sft_train.jsonl \
--val data/sft_val.jsonl --hub-model-id Falln87/clerk-qwen2.5-0.5b-lora
"""
from __future__ import annotations
import argparse
from datasets import load_dataset
from peft import LoraConfig
from trl import SFTConfig, SFTTrainer
BASE_MODEL = "Qwen/Qwen2.5-0.5B-Instruct"
def main() -> None:
ap = argparse.ArgumentParser(description="CLERK LoRA SFT")
ap.add_argument("--train", type=str, required=True)
ap.add_argument("--val", type=str, default=None)
ap.add_argument("--base-model", type=str, default=BASE_MODEL)
ap.add_argument("--hub-model-id", type=str,
default="Falln87/clerk-qwen2.5-0.5b-lora")
ap.add_argument("--output-dir", type=str, default="clerk-ckpt")
ap.add_argument("--lr", type=float, default=2e-4)
ap.add_argument("--epochs", type=float, default=2.0)
ap.add_argument("--max-length", type=int, default=2048)
ap.add_argument("--batch-size", type=int, default=2)
ap.add_argument("--grad-accum", type=int, default=8)
ap.add_argument("--lora-r", type=int, default=32)
ap.add_argument("--trackio-space-id", type=str, default=None)
ap.add_argument("--max-steps", type=int, default=-1)
args = ap.parse_args()
train = load_dataset("json", data_files=args.train, split="train")
train = train.remove_columns(
[c for c in train.column_names if c != "messages"])
assert train.column_names == ["messages"], train.column_names
assert len(train) > 0
if args.val:
val = load_dataset("json", data_files=args.val, split="train")
val = val.remove_columns(
[c for c in val.column_names if c != "messages"])
else:
val = None
peft_config = LoraConfig(
r=args.lora_r,
lora_alpha=args.lora_r,
target_modules=["q_proj", "k_proj", "v_proj", "o_proj",
"gate_proj", "up_proj", "down_proj"],
task_type="CAUSAL_LM",
)
training_args = SFTConfig(
model_init_kwargs={
"attn_implementation": "sdpa",
"dtype": "bfloat16",
"use_cache": True,
},
per_device_train_batch_size=args.batch_size,
gradient_accumulation_steps=args.grad_accum,
warmup_ratio=0.03,
num_train_epochs=args.epochs,
max_steps=args.max_steps,
learning_rate=args.lr,
optim="paged_adamw_8bit",
lr_scheduler_type="cosine",
logging_steps=20,
eval_strategy="steps" if val else "no",
eval_steps=200 if val else None,
per_device_eval_batch_size=args.batch_size,
report_to="trackio",
trackio_space_id=args.trackio_space_id,
output_dir=args.output_dir,
max_length=args.max_length,
packing=False,
push_to_hub=True,
hub_model_id=args.hub_model_id,
seed=137,
data_seed=137,
)
trainer = SFTTrainer(
model=args.base_model,
args=training_args,
train_dataset=train,
eval_dataset=val,
peft_config=peft_config,
)
trainer.train()
trainer.save_model(args.output_dir)
trainer.push_to_hub()
print("done: adapter pushed to", args.hub_model_id)
if __name__ == "__main__":
main() |