Recalibrate batch-size estimate for the post-Phase-2 model size

The estimate_batch_size(training=True) calibration point was measured on
the pre-Phase-2 architecture (hidden_dim=512). Re-measured against the
current hidden_dim=1024 stack (Stage-2 secondary decoder + n_sec head
included): ~29696 batch size at ~7683 MiB VRAM.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
This commit is contained in:
2026-07-08 11:42:03 +02:00
parent 980b6ae7da
commit 436d9fa4d4
+9 -8
View File
@@ -53,14 +53,15 @@ def auto_device() -> torch.device:
return torch.device("cpu")
# 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
# Calibration point for estimate_batch_size(training=True): hidden_dim=1024,
# n_blocks=6, batch_size=29696 measured at ~7683 MiB VRAM (post-Phase-2
# architecture, including the Stage-2 secondary decoder and n_sec head).
# 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 = 7683 * 1024**2
_REF_BATCH_SIZE = 29696
_REF_HIDDEN_DIM = 1024
_REF_N_BLOCKS = 6
# Calibration point for estimate_batch_size(training=False): inference has no