Conversation
Replace the padded (E, max_tokens, dim) expert batches with one ragged torch._grouped_mm call per projection, driven by the int32 prefix sum of the per-expert token counts. Outputs and gradients are bit-identical to the padded path, and activation memory and time no longer scale with the busiest expert. Token counts stay on the device, and the per-expert gradient-scaling loops in MoEMLP and ep_dispatch_and_compute are vectorised.
Codecov Report✅ All modified and coverable lines are covered by tests.
🚀 New features to boost your workflow:
|
Run the same ragged grouped GEMM on CPU: with the gradient hook in place torch's CPU kernel accepts expanded upstream gradients, so the per-group loop is not needed. Add single-rank CPU tests for ep_dispatch_and_compute and for the bf16 grouped path inside MoEMLP.
Naeemkh
approved these changes
Sep 14, 2026
Naeemkh
left a comment
Naeemkh
left a comment
Member
There was a problem hiding this comment.
Please make sure this test runs successfully, if it does, feel free to merge:
tests/integration/test_moe_grouped_gemm.py
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
(E, max_tokens, dim)expert batches ingrouped_expert_forward/grouped_expert_forward_packedwith one raggedtorch._grouped_mm(..., offs=)call per projection, driven by the int32 prefix sum of the per-expert token counts. Activation memory and compute now scale with the routed tokens, not with the busiest expert..tolist()sync). The per-expert gradient-scaling loops inMoEMLP._local_forwardandep_dispatch_and_computeare vectorised (scale_by_expert_load); the EP path keeps one zero-token guard for FSDP2.tests/integration/test_moe_grouped_gemm.py); 40 training steps of a small MoE log identical losses under both dispatches.docs/moe/capacity-and-dispatch.md,docs/distributed/expert-parallelism.md; changelog entry.One MoE layer, forward + backward, 8192 tokens, bf16, H200:
Every row: forward output and all gradients
torch.equalto the padded path.Testing
uv run ruff check kempnerforge/ tests/passesuv run ruff format --check kempnerforge/ tests/ scripts/passesuv run pyright kempnerforge/passes (0 errors)uv run pytest tests/unit/ -v --timeout=60passes (1819 passed;TestHuggingFaceDataset::test_lentimes out on the login node without network, unrelated)uv run pytest tests/integration/ -v— 92 passed (1 H200)uv run torchrun --nproc_per_node=4 -m pytest tests/distributed/ -v— 102 passed, 2 pre-existing skips (4 H200)uv run pytest tests/e2e/ --e2e -v— 27 passed, 1 skipped (slow); the 4 failures (test_checkpoint_save_and_resume,test_pp_checkpoint_save_and_resume,test_sigterm_triggers_emergency_checkpoint,test_moe_checkpoint_resume) predate this branch: they assert the string"Checkpoint saved"while the log line isSync checkpoint saved, and three of them are dense-onlyuv run make -C docs strictCloses #200