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']}"