CLERK LoRA SFT trainer (TRL SFTTrainer), verified against trl examples/sft_qlora; pushes adapter to the Hub

clerk/train_sft.py

@@ -0,0 +1,109 @@
+"""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:
+    import sys
+
+    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=2048,
+        packing=False,
+        push_to_hub=True,
+        hub_model_id=args.hub_model_id,
+        seed=137,
+        data_seed=137,
+    )
+
+    trainer = SFTTrainer(
+        model=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()