Contract the mamba2 chunk scan with einsum instead of broadcast-then-sum - #48978
Conversation
|
run-slow: bamba, falcon_h1, granitemoehybrid, mamba2, nemotron_h, zamba2 |
The fallback `mamba2_chunk_scan` built each contraction as a broadcast product followed by a sum, materialising the un-summed tensor first. G's is (batch, chunks, chunk, chunk, heads, state) in float32 — 4 GiB per sequence at bamba-9B's shapes, so BambaModelIntegrationTest OOMs asking for 8 GiB on a 22.3 GiB runner for a 10-token generate, with 18.2 GiB of that card holding the model setUpClass loaded. Largest intermediate drops 4.000 GiB -> 0.031 GiB.
|
The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update. |
a4d08a1 to
b078515
Compare
|
run-slow: bamba, falcon_h1, granitemoehybrid, mamba2, nemotron_h, zamba2 |
|
run-slow: bamba, falcon_h1, granitemoehybrid, mamba2, nemotron_h, zamba2 |
|
@vasqu should |
|
The bamba fairures are post-OOM, not this patch. |
molbap
left a comment
molbap
left a comment
There was a problem hiding this comment.
cc @Rocketknight1 who I think made also a long investigation on the matter of einsum vs matmuls in many cases (here it might remove some intermediate instances that consume too much memory but we have to be sure it's worth losing the readbility)
vasqu
left a comment
vasqu
left a comment
There was a problem hiding this comment.
Just small comments but overall just 1 thing to keep in mind -> explain the shapes a bit more so non einsum enthusiasts can also understand
(unsure about that one contiguous call which might be unnecessary?)
| output = Y_diag + Y_off | ||
| # Add output of intra-chunk and inter-chunk terms (diagonal and off-diagonal blocks). | ||
| # `contiguous` because einsum may return a permuted view and callers `.view()` this. | ||
| output = (Y_diag + Y_off).contiguous() |
There was a problem hiding this comment.
dont understand why we use the contigous here because reshape below may handle non contigous as well no?
There was a problem hiding this comment.
You're right about the reshape, but he breakage is one level up.
FalconH1Mixer calls .view(batch_size, seq_len, -1) on this function's return value, and without the contiguous()
that breaks 12 falcons tests. (that call is cheap btw)
There was a problem hiding this comment.
lets shorten it to the original comment tho, the contigous is usually self-explanatory
|
run-slow: bamba, falcon_h1 |
This reverts commit 9bf8239.
The einsum contractions reassociate a 128-term reduction, which is bit-exact in float64 but shifts the last bits in float32, and greedy decoding turns that into a different continuation from token ~100 on. Text captured from the A10G runners on this PR; single-gpu and multi-gpu produced it byte-identically.
|
[For maintainers] Suggested jobs to run (before merge) run-slow: bamba, falcon_h1, granitemoehybrid, mamba2, nemotron_h, zamba2 |
|
run-slow: bamba, falcon_h1, granitemoehybrid, mamba2, nemotron_h, zamba2 |
Nvidia CIThis comment contains models: ["models/bamba", "models/falcon_h1", "models/granitemoehybrid", "models/mamba2", "models/nemotron_h", "models/zamba2"] |
AMD CIThis comment contains models: ["models/bamba", "models/falcon_h1", "models/granitemoehybrid", "models/mamba2", "models/nemotron_h", "models/zamba2"] |
CI recapDashboard: View test results in Grafana |
CI Results (Nvidia)Commit Info
Model CI Report❌ 2 new failed tests from this PR 😭
|
4 participants
What does this PR do?
The fallback
mamba2_chunk_scanbuilt each contraction as a broadcast product followed by a sum, materialising the un-summed tensor first. G's is (batch, chunks, chunk, chunk, heads, state) in float32 : 4 GiB per sequence at bamba-9B's shapes, soBambaModelIntegrationTestOOMs asking for 8 GiB on a 22.3 GiB runner for a 10-token generate, with 18.2 GiB of that card holding the model setUpClass loaded. Largest intermediate drops 4GiB -> 0.03 GiB.