diff --git a/giant/cli.py b/giant/cli.py index cf846ec..2fedc4b 100644 --- a/giant/cli.py +++ b/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) diff --git a/giant/config.py b/giant/config.py index 0b3f180..96d2c74 100644 --- a/giant/config.py +++ b/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)