Add --batch-size auto to estimate batch size from free GPU memory

Calibrated against a measured reference point (hidden_dim=512,
n_blocks=6, batch_size=131072 -> ~8 GiB VRAM), assuming activation
memory scales linearly with batch_size * hidden_dim * n_blocks.
CUDA-only for now since it relies on torch.cuda.mem_get_info.

giant/cli.py

@@ -68,7 +68,15 @@ def train(
         typer.Option("--mode", "-m", help="Generative model: flow matching or DDPM"),
     ] = None,
     epochs: Annotated[Optional[int], typer.Option("--epochs", "-e")] = None,
-    batch_size: Annotated[Optional[int], typer.Option("--batch-size", "-b")] = None,
+    batch_size: Annotated[
+        Optional[str],
+        typer.Option(
+            "--batch-size",
+            "-b",
+            help="Integer, or 'auto' to estimate from free GPU memory "
+            "(cuda devices only)",
+        ),
+    ] = None,
     lr: Annotated[Optional[float], typer.Option("--lr", "-l")] = None,
     warmup_epochs: Annotated[
         Optional[int], typer.Option("--warmup-epochs", "-w")
@@ -129,12 +137,28 @@ def train(
     ] = None,
 ) -> None:
     """Train the GIANT surrogate model."""
+    batch_size_auto = False
+    batch_size_value: Optional[int] = None
+    if batch_size is not None:
+        if batch_size.strip().lower() == "auto":
+            batch_size_auto = True
+        else:
+            try:
+                batch_size_value = int(batch_size)
+            except ValueError:
+                typer.echo(
+                    f"error: --batch-size must be an integer or 'auto', "
+                    f"got {batch_size!r}",
+                    err=True,
+                )
+                raise typer.Exit(1)
+
     cli_train = {
         k: v
         for k, v in {
             "mode": mode.value if mode is not None else None,
             "epochs": epochs,
-            "batch_size": batch_size,
+            "batch_size": batch_size_value,
             "lr": lr,
             "warmup_epochs": warmup_epochs,
             "val_fraction": val_fraction,
@@ -161,6 +185,19 @@ def train(
     t, m = cfg["train"], cfg["model"]
 
     _device = torch.device(device) if device else gconfig.auto_device()
+
+    if batch_size_auto:
+        try:
+            t["batch_size"] = gconfig.estimate_batch_size(
+                m["hidden_dim"], m["n_blocks"], _device
+            )
+        except ValueError as exc:
+            typer.echo(f"error: {exc}", err=True)
+            raise typer.Exit(1)
+        typer.echo(
+            f"batch_size: {t['batch_size']} (auto-estimated from free GPU memory)"
+        )
+
     out_dir = out or Path(
         f"checkpoints/{t['mode']}"
         f"_h{m['hidden_dim']}"

giant/config.py

@@ -51,6 +51,43 @@ def auto_device() -> torch.device:
     return torch.device("cpu")
 
 
+# Calibration point for estimate_batch_size: hidden_dim=512, n_blocks=6,
+# batch_size=131072 measured at ~8 GiB VRAM. Activation memory is assumed to
+# scale linearly with batch_size * hidden_dim * n_blocks (the ResBlock stack
+# dominates), so this is a rough estimate rather than a guaranteed bound.
+_REF_BYTES = 8 * 1024**3
+_REF_BATCH_SIZE = 131072
+_REF_HIDDEN_DIM = 512
+_REF_N_BLOCKS = 6
+
+
+def estimate_batch_size(
+    hidden_dim: int,
+    n_blocks: int,
+    device: torch.device,
+    safety_factor: float = 0.8,
+    min_batch_size: int = 1024,
+) -> int:
+    """Estimate a batch size that fits in the free memory on `device`.
+
+    Only supported on CUDA devices, which expose a free/total memory query;
+    other backends (cpu, mps) raise ValueError.
+    """
+    if device.type != "cuda":
+        raise ValueError(
+            f"--batch-size auto is only supported on cuda devices, got {device.type!r}"
+        )
+    device_index = (
+        device.index if device.index is not None else torch.cuda.current_device()
+    )
+    free_bytes, _total_bytes = torch.cuda.mem_get_info(device_index)
+    bytes_per_unit = _REF_BYTES / (_REF_BATCH_SIZE * _REF_HIDDEN_DIM * _REF_N_BLOCKS)
+    bytes_per_sample = bytes_per_unit * hidden_dim * n_blocks
+    batch_size = int(free_bytes * safety_factor / bytes_per_sample)
+    batch_size = max(min_batch_size, (batch_size // 1024) * 1024)
+    return batch_size
+
+
 def load_toml(path: Path) -> dict:
     with open(path, "rb") as f:
         return tomllib.load(f)