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:
+9
-8
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user