Upload folder using huggingface_hub

src/mathcompose/train/train.py

@@ -30,16 +30,19 @@ def _presets_for(vram_gb: float) -> dict:
     """Per-GPU QLoRA defaults for a ~1.5B model. Explicit CLI flags override these;
     these override the yaml config."""
     big = vram_gb >= 70  # 80GB A100/H100 vs 40GB
+    # Gradient checkpointing stays ON everywhere: with Qwen's 152k vocab the logits
+    # + activations at seq 4096 dwarf the (tiny, 4-bit) weights, so checkpointing is
+    # what keeps even an 80GB card in budget. Effective batch is ~16 via accumulation.
     return {
         # T4 has no bf16 and its fp16 GradScaler has been unreliable -> amp off (fp32).
-        "t4":   dict(batch_size=1,  grad_accum=16, max_length=2048, grad_ckpt="on",  amp="off"),
-        "l4":   dict(batch_size=4,  grad_accum=4,  max_length=4096, grad_ckpt="on",  amp="auto"),
-        "a100": dict(batch_size=16 if big else 8, grad_accum=1 if big else 2,
-                     max_length=4096, grad_ckpt="off", amp="auto"),
-        "h100": dict(batch_size=16, grad_accum=1,  max_length=4096, grad_ckpt="off", amp="auto"),
-        "generic_big":   dict(batch_size=8, grad_accum=2,  max_length=4096, grad_ckpt="off", amp="auto"),
-        "generic_mid":   dict(batch_size=4, grad_accum=4,  max_length=3072, grad_ckpt="on",  amp="auto"),
-        "generic_small": dict(batch_size=1, grad_accum=16, max_length=2048, grad_ckpt="on",  amp="auto"),
+        "t4":   dict(batch_size=1,  grad_accum=16, max_length=2048, grad_ckpt="on", amp="off"),
+        "l4":   dict(batch_size=2,  grad_accum=8,  max_length=3072, grad_ckpt="on", amp="auto"),
+        "a100": dict(batch_size=8 if big else 4, grad_accum=2 if big else 4,
+                     max_length=4096, grad_ckpt="on", amp="auto"),
+        "h100": dict(batch_size=8,  grad_accum=2,  max_length=4096, grad_ckpt="on", amp="auto"),
+        "generic_big":   dict(batch_size=4, grad_accum=4,  max_length=4096, grad_ckpt="on", amp="auto"),
+        "generic_mid":   dict(batch_size=2, grad_accum=8,  max_length=3072, grad_ckpt="on", amp="auto"),
+        "generic_small": dict(batch_size=1, grad_accum=16, max_length=2048, grad_ckpt="on", amp="auto"),
     }
 
 
@@ -71,6 +74,9 @@ def detect_gpu_preset(force: str = "auto") -> tuple[str, dict]:
 
 
 def build_and_train(args) -> str:
+    import os
+    # reduce fragmentation OOMs (set before CUDA initializes)
+    os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")
     import torch
     from datasets import load_dataset
     from transformers import AutoModelForCausalLM