| """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() |
