Upload folder using huggingface_hub

src/mathcompose/train/train.py

@@ -70,7 +70,23 @@ def build_and_train(args) -> str:
     eval_ds = ds.get("validation")
 
     t = cfg["train"]
-    grad_ckpt = bool(t.get("gradient_checkpointing", True)) and use_4bit
+    # optional CLI overrides (getattr keeps programmatic callers/tests working)
+    if getattr(args, "epochs", None):
+        t["epochs"] = args.epochs
+    if getattr(args, "batch_size", None):
+        t["per_device_train_batch_size"] = args.batch_size
+    if getattr(args, "grad_accum", None):
+        t["gradient_accumulation_steps"] = args.grad_accum
+    if getattr(args, "max_length", None):
+        t["max_length"] = args.max_length
+
+    gc_mode = getattr(args, "grad_ckpt", "auto")
+    if gc_mode == "on":
+        grad_ckpt = True
+    elif gc_mode == "off":
+        grad_ckpt = False
+    else:
+        grad_ckpt = bool(t.get("gradient_checkpointing", True)) and use_4bit
     optim = t.get("optim", "paged_adamw_8bit") if use_4bit else "adamw_torch"
 
     sft_args = SFTConfig(
@@ -127,7 +143,13 @@ def main() -> None:
     ap.add_argument("--val-file", default=None)
     ap.add_argument("--output-dir", default=None)
     ap.add_argument("--no-4bit", action="store_true", help="disable 4-bit (CPU smoke / debugging)")
-    ap.add_argument("--max-steps", type=int, default=0, help="cap steps (smoke test)")
+    ap.add_argument("--max-steps", type=int, default=0, help="cap steps (GPU/CPU smoke)")
+    # per-run overrides for adapting to the GPU you got (T4 vs L4/A100/H100)
+    ap.add_argument("--epochs", type=int, default=None)
+    ap.add_argument("--batch-size", type=int, default=None, help="per-device train batch")
+    ap.add_argument("--grad-accum", type=int, default=None)
+    ap.add_argument("--max-length", type=int, default=None, help="sequence length (<=4096)")
+    ap.add_argument("--grad-ckpt", choices=["auto", "on", "off"], default="auto")
     ap.add_argument("--wandb", action="store_true")
     ap.add_argument("--push", action="store_true", help="push adapter to HF Hub")
     ap.add_argument("--hub-model-id", default=None)