diff --git a/giant/cli.py b/giant/cli.py index d71e640..ab71be6 100644 --- a/giant/cli.py +++ b/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']}" diff --git a/giant/config.py b/giant/config.py index 5fbb127..0b3f180 100644 --- a/giant/config.py +++ b/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)