Upload folder using huggingface_hub

src/mathcompose/train/train.py

@@ -26,6 +26,50 @@ from pathlib import Path
 DEFAULT_CONFIGS = {"v": "configs/verifier_v.yaml", "g": "configs/generator_g.yaml"}
 
 
+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
+    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"),
+    }
+
+
+def detect_gpu_preset(force: str = "auto") -> tuple[str, dict]:
+    """Return (label, preset dict). ``force`` in {auto,none,t4,l4,a100,h100}."""
+    import torch
+
+    if force == "none":
+        return "none", {}
+    if not torch.cuda.is_available():
+        return "cpu", {}
+    vram = torch.cuda.get_device_properties(0).total_memory / 1e9
+    presets = _presets_for(vram)
+    if force in ("t4", "l4", "a100", "h100"):
+        return force, presets[force]
+    name = torch.cuda.get_device_name(0).lower()
+    cap = torch.cuda.get_device_capability()
+    if "a100" in name:
+        key = "a100"
+    elif "h100" in name or "h200" in name:
+        key = "h100"
+    elif "l4" in name:
+        key = "l4"
+    elif "t4" in name or cap[0] < 8:
+        key = "t4"
+    else:
+        key = "generic_big" if vram >= 70 else ("generic_mid" if vram >= 20 else "generic_small")
+    return key, presets[key]
+
+
 def build_and_train(args) -> str:
     import torch
     from datasets import load_dataset
@@ -43,13 +87,28 @@ def build_and_train(args) -> str:
     val_file = args.val_file or cfg["data"].get("val_path")
 
     has_cuda = torch.cuda.is_available()
-    # bf16 needs Ampere+ (sm_80): A100/L4/H100 yes, T4 (sm_75) no -> use fp16 there.
     bf16_ok = has_cuda and torch.cuda.get_device_capability()[0] >= 8
-    compute_dtype = torch.bfloat16 if bf16_ok else (torch.float16 if has_cuda else torch.float32)
+
+    label, preset = detect_gpu_preset(getattr(args, "gpu_preset", "auto"))
+    amp_off = getattr(args, "no_amp", False) or preset.get("amp") == "off"
+
+    # Precision: bf16 on Ampere+ (A100/L4/H100); fp16 AMP on T4; fp32 when AMP is
+    # disabled on a non-bf16 GPU (avoids fp16 gradient underflow without a scaler).
+    if bf16_ok:
+        compute_dtype = torch.bfloat16
+    elif has_cuda and not amp_off:
+        compute_dtype = torch.float16
+    else:
+        compute_dtype = torch.float32
+
     if has_cuda:
         cap = torch.cuda.get_device_capability()
-        print(f"[gpu] {torch.cuda.get_device_name(0)} (sm_{cap[0]}{cap[1]}) -> "
-              f"{'bf16' if bf16_ok else 'fp16'} compute")
+        prec = "bf16" if bf16_ok else ("fp32(no-amp)" if amp_off else "fp16")
+        print(f"[gpu] {torch.cuda.get_device_name(0)} (sm_{cap[0]}{cap[1]}) | "
+              f"preset={label} | compute={prec}")
+    if preset:
+        print(f"[preset] {label} -> {preset}")
+
     use_4bit = (not args.no_4bit) and cfg["quant"]["load_in_4bit"] and has_cuda
     if cfg["quant"]["load_in_4bit"] and not has_cuda and not args.no_4bit:
         print("[warn] 4-bit requested but no CUDA — falling back to full-precision "
@@ -62,19 +121,18 @@ def build_and_train(args) -> str:
         model_kwargs["quantization_config"] = build_bnb_config(cfg, compute_dtype=compute_dtype)
     model = AutoModelForCausalLM.from_pretrained(base_id, **model_kwargs)
 
-    if compute_dtype == torch.float16:
-        # On T4 the fp16 GradScaler can't unscale bf16 grads. transformers may load
-        # the non-quantized weights in the model's default bf16 regardless of dtype=,
-        # so LoRA inherits bf16. Force any remaining bf16 params to fp16 and pin the
-        # config dtype so PEFT doesn't recreate bf16 adapters from the stale config.
-        model.config.torch_dtype = torch.float16
+    if has_cuda and compute_dtype in (torch.float16, torch.float32):
+        # transformers may keep non-quantized weights in the model's default bf16
+        # even with dtype=; a bf16 grad breaks the fp16 GradScaler and mixes dtypes.
+        # Force any bf16 params to the chosen compute dtype and pin the config.
+        model.config.torch_dtype = compute_dtype
         cast = 0
         for p in model.parameters():
             if p.dtype == torch.bfloat16:
-                p.data = p.data.to(torch.float16)
+                p.data = p.data.to(compute_dtype)
                 cast += 1
         if cast:
-            print(f"[dtype] cast {cast} bf16 param tensors -> fp16 (T4 AMP compatibility)")
+            print(f"[dtype] cast {cast} bf16 param tensors -> {compute_dtype}")
 
     data_files = {"train": train_file}
     if val_file and Path(val_file).exists():
@@ -84,42 +142,51 @@ def build_and_train(args) -> str:
     eval_ds = ds.get("validation")
 
     t = cfg["train"]
-    # 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")
+
+    def pick(cli_val, key, cfg_default):
+        """Precedence: explicit CLI flag > GPU preset > yaml config."""
+        if cli_val is not None:
+            return cli_val
+        if key in preset:
+            return preset[key]
+        return cfg_default
+
+    max_ctx = cfg["model"].get("max_context", 4096)
+    epochs = getattr(args, "epochs", None) or t.get("epochs", 3)
+    batch = pick(getattr(args, "batch_size", None), "batch_size", t.get("per_device_train_batch_size", 2))
+    accum = pick(getattr(args, "grad_accum", None), "grad_accum", t.get("gradient_accumulation_steps", 8))
+    maxlen = min(pick(getattr(args, "max_length", None), "max_length", t.get("max_length", max_ctx)), max_ctx)
+
+    gc_cli = getattr(args, "grad_ckpt", "auto")
+    gc_mode = gc_cli if gc_cli != "auto" else preset.get("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"
+    print(f"[train] epochs={epochs} batch={batch} accum={accum} (eff={batch * accum}) "
+          f"max_length={maxlen} grad_ckpt={grad_ckpt}")
 
     sft_args = SFTConfig(
         output_dir=output_dir,
-        num_train_epochs=t.get("epochs", 3),
-        per_device_train_batch_size=t.get("per_device_train_batch_size", 2),
-        gradient_accumulation_steps=t.get("gradient_accumulation_steps", 8),
+        num_train_epochs=epochs,
+        per_device_train_batch_size=batch,
+        gradient_accumulation_steps=accum,
         learning_rate=float(t.get("learning_rate", 2e-4)),
         lr_scheduler_type=t.get("lr_scheduler_type", "cosine"),
         warmup_ratio=t.get("warmup_ratio", 0.03),
         logging_steps=t.get("logging_steps", 10),
         save_steps=t.get("save_steps", 50),
-        max_length=t.get("max_length", cfg["model"].get("max_context", 2048)),
+        max_length=maxlen,
         completion_only_loss=t.get("completion_only_loss", True),
         packing=t.get("packing", False),
         gradient_checkpointing=grad_ckpt,
         optim=optim,
-        bf16=bf16_ok and not getattr(args, "no_amp", False),
-        fp16=has_cuda and not bf16_ok and not getattr(args, "no_amp", False),
+        bf16=bf16_ok and not amp_off,
+        fp16=has_cuda and not bf16_ok and not amp_off,
         seed=t.get("seed", 0),
         report_to=(["wandb"] if args.wandb else "none"),
         max_steps=args.max_steps if args.max_steps else -1,
@@ -164,6 +231,8 @@ def main() -> None:
     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("--gpu-preset", choices=["auto", "none", "t4", "l4", "a100", "h100"],
+                    default="auto", help="GPU-tuned defaults (auto-detects by default)")
     ap.add_argument("--no-amp", action="store_true",
                     help="disable fp16/bf16 mixed precision (T4 fallback if the scaler errors)")
     ap.add_argument("--wandb", action="store_true")