Batch 48 Memory and Risk Check | Generated by AI
Good thought. Let me analyze without killing the current run.
batch=48 estimate:
| Component | batch=32 (current) | batch=48 (estimate) |
|---|---|---|
| Microbatch tokens | 32×2048 = 65,536 | 48×2048 = 98,304 |
Logit tensor (batch×seq, vocab) |
(65536, 32768) = 8 GB | (98304, 32768) = 12 GB |
| Activations (embed + attn + MLP) | ~30 GB | ~45 GB |
| Model weights (fp8) | ~1.4 GB | ~1.4 GB (same) |
| Optimizer states (fp32) | ~16.5 GB (1.38B × 12 bytes for Muon+AdamW) | ~16.5 GB (same) |
| Gradients (fp8) | ~1.4 GB | ~1.4 GB (same) |
| Compiled graph memory | ~35 GB (torch.compile caches) | ~35 GB (similar) |
| Total estimate | ~93 GB | ~112 GB |
| H200 capacity | 140 GB | 140 GB |
| Headroom | 47 GB free | ~28 GB free (probably fine) |
Key question: Does torch.compile recompile the graph for the new shape, using more memory temporarily? Yes — it allocates extra buffers during compilation that get freed after. During the first step (compilation), peak memory can spike 10-15 GB above steady state.
So real risk profile:
- Steady state: ~112 GB → 28 GB free → likely safe ✅
- Compilation spike: ~112 GB + ~15 GB = 127 GB → 13 GB free → should still fit ✅
- Grad accum: batch=48 with total-batch=1M means grad_accum = 10-11, which is fine
Where it could fail: if the compiled kernel for the attention forward pass allocates extra scratch space for FA3 with the (48, 2048, 12, 128) shape. That’s a hot path allocation that isn’t cached.
My call: 48 would probably work, but tight. The real question is whether it’s worth the risk — you gain only ~$15\%$ throughput (bigger matmuls) but lose the current 22 minutes of training progress. The MFU gain from 32→48 would be marginal (maybe 65% → 66-67%).
Want me to try 48 after this run finishes (since we can’t save checkpoints mid-run)? Or stay at 32 which is already running fine with 47 GB headroom?
Generated by AI. Curating and sharing still takes effort.
If you find it useful, feel free to
donate.
WeChat: @lzwjava
·
X: @lzwjava ·
Say hi 👋