Conversation
nil0x9
reviewed
Sep 14, 2026
| remainder_strategy: Annotated[ | ||
| Literal["agrs", "pad_all2all"], | ||
| Parameter(help="Communication strategy for Muon parameter batches smaller than the FSDP group size"), | ||
| ] = "agrs" |
Collaborator
There was a problem hiding this comment.
LGTM. It's good to have another option for flexibility.
| if sharded_tensor_dim is not None and non_fsdp_shard_factor > 1: | ||
| full_shard_dim_size = params[0].size(sharded_tensor_dim) | ||
| assert full_shard_dim_size % non_fsdp_shard_factor == 0 | ||
| global_shard_dim_size = full_shard_dim_size // non_fsdp_shard_factor |
_create_muon_tasks mixed group-, mesh-, and batch-level logic in one generator body. Extract _resolve_mesh_comm_plan (shape-independent sharding topology and communication strategy, resolved once per device_mesh/placements group) and _build_muon_task (shape-dependent per-batch metadata and AsyncTask construction), leaving _create_muon_tasks as a plain linear loop over the three levels. Also replace the EP-mesh-dim loop with a direct lookup by elimination, since the multi-shard case is validated to be exactly FSDP + EP.
- muon_update_batch_async's agrs/all_to_all branches recomputed global_shard_dim_size from X[0].size(shard_dim) whenever the caller passed None. The optimizer's only caller, _build_muon_task, is the sole owner of this FSDP-visible size and always derives it the same way, so the fallback was dead defensive code; replace it with an assert that the caller supplied it, matching the existing shard_dim/ process_group assertions for these strategies. - _build_muon_task correspondingly always computes global_shard_dim_size when sharded_tensor_dim is set (dividing by a non_fsdp_shard_factor of 1 is a no-op), instead of special-casing factor > 1 and leaving it None otherwise. - Replace the EP-mesh-dim loop in _resolve_mesh_comm_plan with a direct lookup by elimination, since exactly 2 shard placements are already validated and one is known to be the FSDP dim. - Add step-by-step comments to _resolve_mesh_comm_plan and _build_muon_task describing each stage of the mesh-plan and batch-task resolution.
Handle the unsharded mesh case with an early local plan so sharded_mesh_dim and sharded_tensor_dim stay non-optional, and replace the flag-then-dispatch strategy selection with one early return per communication strategy.
The sub-group path was gated on `len(params) < fsdp_size`, so a deep MoE model whose expert params fill a batch (GLM-5.2: 76 fused expert weights over a 64-rank FSDP mesh) fell back to the full-group all-to-all and its AGRS remainder, which reconstruct the whole multi-expert matrix on every rank and blow up memory. Drop the guard so any group whose experts span a sub-group uses the sub-group all-gather, and record the measured per-strategy memory cost next to the choice.
AGRS all-gathers every matrix of the batch onto every rank while each rank
orthogonalizes only R/W of them, so both its time and its peak memory grow
linearly with the batch size, while padded all-to-all is flat in R.
Measured at world size 8 on 96MB (8192x6144, bf16) matrices, one remainder
batch, mean of 10 steps:
R agrs pad_all2all
1 14.90ms 624MB 15.15ms 876MB
3 16.68ms 1104MB 15.38ms 804MB
5 18.34ms 1488MB 15.29ms 732MB
7 20.07ms 2016MB 15.21ms 660MB
The gap widens with the FSDP group size: GLM-5.2 on a 64-rank mesh leaves a
12-param remainder per shape group, where AGRS moves ~12x the bytes.
Without all-to-all AGRS is the only remainder path, so enable_all2all=False now
coerces the strategy instead of rejecting the default combination.
…o-all
A remainder batch of R matrices used to be padded up to the FSDP group size, so
every rank received a full matrix and W-R of them orthogonalized zeros. Sending
the batch with all_to_all_single split sizes instead keeps the assembling ranks'
work identical and leaves the rest idle: no zero shards on the wire, no zero
Newton-Schulz, and no full-matrix buffers on ranks without a matrix.
Measured at world size 8 on 96MB (8192x6144, bf16) matrices, R=2, mean of 10
steps, peak allocation during the step:
strategy ms/step busy rank idle rank
agrs 16.03 912 MB 576 MB
padded a2a 14.86 804 MB 804 MB
ragged a2a 14.82 648 MB 48 MB
Step time is unchanged at this scale — the collective is ~0.3ms against ~14ms of
Newton-Schulz, and a task's latency is gated by the ranks that do assemble a
matrix. The gain is the freed memory plus the freed GPU time on idle ranks, which
AsyncRuntime can spend on the other batches in flight; both grow with the group
size, and GLM-5.2 on a 64-rank mesh leaves 52 of 64 ranks idle per remainder.
Uneven shards still pad, since split sizes require equal-sized shards.
The ragged remainder path was gated on the sharded dimension dividing the FSDP group size, which nothing guarantees: DTensor chunks a dimension into ceil(size / world_size) pieces, so trailing ranks can hold a short shard or none, and both AGRS and the padded all-to-all carry explicit uneven branches for it. Split sizes are per rank, so they express an uneven shard directly. Sending the sharded dimension first (a no-op view under FSDP, which shards dim 0) lets one code path cover both cases and drops the gate, so every remainder batch now takes the ragged exchange. The DTensor split rule moves into _shard_sizes, shared with the uneven all-to-all branch.
…tegy The ragged exchange was applied to every remainder batch, which left no way to ask for the padded one and made remainder_strategy='pad_all2all' a no-op. Add 'ragged_all_to_all' as a third value so the knob names the batch strategy directly, and keep 'pad_all2all' as the default. Tested across all three values: the FSDP parity test gains a ragged case, and the EP test asserts that padding hands every rank a matrix while AGRS and the ragged exchange leave the batch to the rank that assembles it.
These were scratch artifacts from diagnosing the AGRS EP/FSDP OOM and are not needed now that the fix and its regression test are in place. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
jayhenry
force-pushed
the
fix_muon_oom
branch
from
September 16, 2026 11:34
6afd17c to
2b0c3ed
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Fixes an EP+FSDP AGRS out-of-memory in the Muon optimizer's remainder-batch path, then replaces the
implicit remainder handling with an explicit, benchmarked
remainder_strategychoice(
agrs/pad_all2all/ragged_all_to_all) and a companion sub-group all-gather fix.Root Cause
For GLM-5.2 with 512 ranks and EP8, the expert FSDP group has 64 ranks. The 76 same-shaped
fused_w2parameters form a remainder batch of 12. AGRS used the full DTensor dimension(256 * 6144)instead of the EP-local, FSDP-visible dimension(256 * 6144) / 8, producing a72 GiB communication buffer and a second 72 GiB
flatten()copy.AGRS itself also all-gathers every matrix in a batch onto every rank, so its time and peak memory
scale linearly with batch size — worst case is exactly a large remainder batch, which is where
GLM-5.2 hit this OOM.
Fix
shape-homogeneous batch instead of once per mesh group (the previous scope let a later
differently-shaped batch reuse an earlier batch's dimension).
Muon._create_muon_taskswassplit into
_resolve_mesh_comm_plan(mesh-topology / communication-strategy, resolved once per(device_mesh, placements)group) and_build_muon_task(per-batch, shape-dependent metadata),so the split makes this scoping explicit instead of implicit in one long function.
sub-group) was gated on
len(params) < fsdp_size, so a model whose expert params happen tofill a batch (GLM-5.2: 76 fused expert weights over 64 ranks) fell back to full-group
all-to-all / AGRS reconstructing the whole multi-expert matrix on every rank. The guard is
dropped so any group whose experts span a sub-group always takes the cheaper path.
remainder_strategyoption, benchmarked at world size 8 on 96 MB (8192x6144, bf16)matrices:
"agrs"— all-gather + reduce-scatter; time/memory grow linearly with remainder size(14.9 ms / 624 MB at R=1 → 20.1 ms / 2016 MB at R=7).
"pad_all2all"(default) — zero-pads the remainder to world size and uses a uniformall-to-all; flat ~15.2 ms regardless of R, but every rank (including idle ones) allocates a
full matrix.
"ragged_all_to_all"— sends only the real matrices via all-to-all split sizes, so idleranks allocate nothing (648 MB busy / 48 MB idle at R=2, vs. 804 MB / 804 MB for padded
all-to-all). The same split-size path also covers uneven shards from DTensor's
ceil(size / world_size)chunking, so it needs no separate padding branch.enable_all2all=Falsenow coerces the strategy to"agrs"instead of rejecting thecombination, since AGRS is the only remainder path that doesn't use all-to-all.
regression test.
Test Plan
PYTHONPATH=. python -m pytest tests/optim/test_muon.py -q— 20 passed (run under the localGPU lock), covering:
where Newton-Schulz sees the un-divided global shape instead of the EP-local one).
remainder_strategyvalidation and default-value behavior, including theenable_all2all=Falsecoercion.ruff check/ruff format --checkon the changed files.