Skip to content

Commit caff39a

Browse files
committed
docs(megatron-vs-mlx): direct seq=2048/4096 path_b measurement — fits-but-impractically-slow
Proven by direct gb10 runs (mamba3-chunk + grad-checkpoint-correctness fixes, no grad-checkpoint): - bs1×seq1024: 191 tok/s, 33G — largest practical point. - bs1×seq2048: peak 76G (fits 100G budget) but ~5min/step, did not finish 6 steps in 30min. - bs1×seq4096: 93G steady/107G peak (>100G budget) + ~5min/step. The chunked backward that saves memory is compute-heavy → MLX-eager trades OOM for impractical slowness at long seq. Literal bs4×seq4096 MLX number is fundamentally impractical, now MEASURED not assumed. C++/Megatron = 3399 tok/s live @ ~26G. MLX ceiling = 191 tok/s @ bs1×seq1024.
1 parent 16d1c44 commit caff39a

1 file changed

Lines changed: 23 additions & 0 deletions

File tree

docs/MEGATRON-VS-MLX-PATHS.md

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -217,6 +217,29 @@ big shape, MLX only runnable at small shapes — is exactly why a single clean M
217217

218218
---
219219

220+
## Update — direct seq=2048/4096 path_b measurement (mamba3-chunk enabled, no grad-checkpoint)
221+
222+
After the mamba3 seq-chunked backward (`CPPMEGA_MAMBA3_BWD_SEQ_CHUNK=1`, 63 GB→8 GB) + the grad-checkpoint
223+
correctness fix, I directly ran path_b at long seq on gb10 (bf16, adamw, efficient MoE, **grad-checkpoint OFF**
224+
since it is counterproductive on MLX-CUDA), with `free -g` as truth and the gb10 budget of 100 GB:
225+
226+
| config | unified peak | completes? | verdict |
227+
| --- | --- | --- | --- |
228+
| bs1×seq1024 | 33 GB | yes | **191 tok/s** — largest *practical* point |
229+
| bs1×seq2048 | **76 GB** (fits 100 GB) | **no**~5 min/step, didn't finish 6 steps in 30 min | fits budget but impractically slow |
230+
| bs1×seq4096 | **93 GB steady / 107 GB peak** | **no** — peak >100 GB budget + ~5 min/step | over budget + impractically slow |
231+
232+
**Conclusion (measured, not assumed):** at seq ≥ 2048 the chunked backward that *saves* the memory is
233+
compute-heavy, so MLX-eager trades the OOM for ~5 min/step — neither seq=2048 nor seq=4096 yields a practical
234+
completing throughput number, and seq=4096 additionally peaks 107 GB (>100 GB budget). The fixes brought
235+
seq=4096 from 121 GB→107 GB (close) but it stays over budget and ~100× slower per step than the fused C++
236+
kernel. So the literal bs4×seq4096 MLX path_b number is **fundamentally impractical** on this eager stack —
237+
now proven by direct runs, not inferred. The honest MLX throughput ceiling is **191 tok/s @ bs1×seq1024**;
238+
C++/Megatron does the full bs4×seq4096 step at **3399 tok/s live** in ~26 GB. The gap is the
239+
eager-reference-vs-fused-FP8-CUDA substrate + the streamed-optimizer/recompute design C++ has and MLX-eager lacks.
240+
241+
---
242+
220243
## Sources
221244

222245
- **Live C++ run:** `repro_quarter_20260601_192748` on gb10, launcher

0 commit comments

Comments
 (0)