Calibrate auto batch size separately for inference vs training

Inference has no backward graph or optimizer state, so it has a much
lower per-sample memory footprint than training. estimate_batch_size
now takes a training flag selecting between two calibration points;
predict uses the inference one (hidden_dim=1024, n_blocks=8,
batch_size=65536 measured at ~2037 MiB).

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

giant/cli.py

@@ -304,7 +304,10 @@ def predict(
     if batch_size_auto:
         try:
             batch_size_value = gconfig.estimate_batch_size(
-                model_cfg["hidden_dim"], model_cfg["n_blocks"], _device
+                model_cfg["hidden_dim"],
+                model_cfg["n_blocks"],
+                _device,
+                training=False,
             )
         except ValueError as exc:
             typer.echo(f"error: {exc}", err=True)

giant/config.py

@@ -51,15 +51,25 @@ 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.
+# Calibration point for estimate_batch_size(training=True): 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
 
+# Calibration point for estimate_batch_size(training=False): inference has no
+# backward graph or optimizer state, so its memory footprint is much smaller
+# per sample. hidden_dim=1024, n_blocks=8, batch_size=65536 measured at ~2037
+# MiB VRAM.
+_REF_BYTES_PREDICT = 2037 * 1024**2
+_REF_BATCH_SIZE_PREDICT = 65536
+_REF_HIDDEN_DIM_PREDICT = 1024
+_REF_N_BLOCKS_PREDICT = 8
+
 
 def estimate_batch_size(
     hidden_dim: int,
@@ -67,11 +77,14 @@ def estimate_batch_size(
     device: torch.device,
     safety_factor: float = 0.8,
     min_batch_size: int = 1024,
+    training: bool = True,
 ) -> 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.
+    other backends (cpu, mps) raise ValueError. Pass `training=False` for
+    inference (e.g. `predict`), which uses a much lower per-sample memory
+    calibration since there's no backward graph or optimizer state.
     """
     if device.type != "cuda":
         raise ValueError(
@@ -81,7 +94,21 @@ def estimate_batch_size(
         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)
+    if training:
+        ref_bytes, ref_batch_size, ref_hidden_dim, ref_n_blocks = (
+            _REF_BYTES,
+            _REF_BATCH_SIZE,
+            _REF_HIDDEN_DIM,
+            _REF_N_BLOCKS,
+        )
+    else:
+        ref_bytes, ref_batch_size, ref_hidden_dim, ref_n_blocks = (
+            _REF_BYTES_PREDICT,
+            _REF_BATCH_SIZE_PREDICT,
+            _REF_HIDDEN_DIM_PREDICT,
+            _REF_N_BLOCKS_PREDICT,
+        )
+    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)