Add MORI as a Token Dispatcher#135
Merged
Merged
Conversation
- Introduced `train_qwen2.sh` and `train_qwen3.sh` scripts for pretraining Qwen2 and Qwen3 models, respectively, with detailed configuration options. - Added `readme.md` files for both Qwen2 and Qwen3, outlining dataset preparation, model training processes, and environment setup. - Included instructions for single-node and multi-node training setups, along with customizable parameters for model size and training iterations.
…node runs - Updated `train_qwen3.sh` to include FP8 recipe options for improved precision handling during training. - Added a new script `1P1G_AAI_proxy_mi355x_qwen3_235B_A22B.sh` for smaller Qwen3-235B-style MoE proxy configurations, allowing for single-node and single-GPU smoke runs. - Improved environment variable handling for model size, layers, and experts in the training scripts to facilitate easier customization. - Adjusted default values and added comments for clarity on tunable parameters.
…ation - Modified `train_deepseekv3.sh` to disable legacy grouped GEMM when using FP8, ensuring compatibility with TE grouped GEMM requirements. - Updated `1P1G_AAI_proxy_mi355x_qwen3_235B_A22B.sh` to set `MOE_USE_LEGACY_GROUPED_GEMM` to false for improved performance. - Enhanced `train_qwen3.sh` to include automatic handling of `--moe-router-padding-for-fp8` for MXFP8 configurations, streamlining the training process for FP8 models. - Added informative logging to clarify changes in grouped GEMM settings based on FP8 configurations.
…bgroup
MORI's symmetric-heap allocation (hipExtMallocWithFlags, 32 GB+) collides
with TransformerEngine's HIP runtime init and aborts with glibc
`free(): invalid size`. Initialize MORI shmem in pretrain_gpt.py as the
first statement of the entry point, before any megatron/TE imports, against
a dedicated gloo subgroup registered as "mori" (mirrors vLLM's
MoriAll2AllManager). World group stays pure NCCL.
- pretrain_gpt.py: inline early bootstrap (no-op without --moe-enable-mori)
- fused_a2a.py: drop now-dead init_mori_shmem and its lazy call site
- training/initialize.py: pass device_id to init_process_group for eager
NCCL init (required by MORI)
- examples/{qwen3,deepseek_v3}/train_*.sh: ulimit -l unlimited (libibverbs
pin), force MORI_SHMEM_HEAP_SIZE=32G, --moe-mori-max-tokens-per-rank=32768
- MoriDispatch.forward: rebase op.dispatch's global expert ids to local ids with -1 sentinels (DeepEP contract) so _indices_to_multihot stops writing OOB and triggering HSA_STATUS_ERROR_EXCEPTION 0x1016. Fix tokens_per_expert via the is_local mask. - _MoriManager.combine and MoriDispatch.backward: pass receiver-side dispatched_probs ([total_recv, topk]) to op.combine instead of sender-side token_probs ([N, topk]). The IntraNode kernel indexes weightsBuf by totalRecvTokenNum, so the sender buffer caused a ~1.1 MiB OOB read and a delayed Memory access fault. - Slice op.combine's [max_num_inp_token_per_rank, hidden_dim] output down to num_tokens in MoriCombine.forward and MoriDispatch.backward so view(hidden_shape) and dx match the original input. Also drop the double reset_mori_op() (call_reset=True already resets), remove triage prints, and lower qwen3's --moe-mori-max-tokens-per-rank from 32768 to 8192.
- Introduced a process-wide CUDA stream for MORI dispatch/combine operations to allow concurrent execution with non-dependent work. - Implemented `_get_mori_comm_stream` to lazily allocate and cache the communication stream with high priority. - Updated `_run_mori_op_on_stream` to execute MORI operations on the dedicated stream when specified, enhancing performance by overlapping communication and computation. - Modified `MoriDispatch` and `MoriCombine` classes to support new parameters for asynchronous execution and stream allocation, ensuring compatibility with existing functionality. - Adjusted `_MoriManager` to pass the new flags for async operations during token dispatching and combining, improving overall efficiency in the MORI framework.
and auto calculate -moe-mori-max-tokens-per-rank
…side global indices are used for backward operations. Updated `MoriDispatch` and `MoriCombine` to save and pass the correct indices, preventing incorrect summation in the kernel. Improved documentation for clarity on input expectations and added receiver-side caches in `_MoriManager` for better integration with MORI's operations.
…f bug and adding a helper function to ensure non-empty weights for MORI operations. This change prevents potential issues with zero-row tensors during dispatch and improves the robustness of the MORI dispatch and combine functionality.
…reset_mori_op` functions. The `get_mori_op` now ensures a resolved kernel type is used and clarifies the operator's lifecycle in the documentation. The `reset_mori_op` function has been updated to clear the global MORI operator reference, preventing stale handles during tests. Additionally, the test suite has been modified to call `reset_mori_op` in the teardown method, ensuring a fresh state for each test case.
Drop the `recv_num_token[0].item()` host sync in `MoriDispatch.forward` and the matching slice in `MoriCombine.backward` in favor of working on the full `[max_num_tokens_per_rank, ...]` padded buffers and masking the tail with `masked_fill` (0/`-1`/0), threading the live `recv_num_token` tensor through `_MoriManager` into `MoriCombine` so both backwards can re-mask without re-syncing. Replace `torch.bincount` (internal `.max().item()`) and the `recv_token_indices[is_local]` boolean gather with a fixed-size `scatter_add_` + `torch.where`. Remove the now-unused `_ensure_combine_weights_non_empty` workaround. Clamp fused-permute's `num_out_tokens` to `>= 1` to keep TE's unpermute backward from raising `AttributeError` on cold ranks. Full `tests/.../test_forward_backward[mori-*]` matrix passes.
Enhance the _MoriManager class by introducing a sync-free caching mechanism for tokens per expert. The new implementation allows for efficient retrieval of token counts without GPU-CPU synchronization, leveraging CPU and lazy-allocated GPU tensors. This change improves performance during expert layer operations and reduces unnecessary host syncs, particularly benefiting the Triton GroupedGEMM path. Additionally, new methods for computing fixed target per expert and padding routing maps have been added to streamline the token management process.
Introduce the MOE_PERMUTE_FUSION environment variable to control the fused token permute/unpermute behavior in the training script. Update the MoE options section to conditionally include the --moe-permute-fusion flag based on the new variable. This enhancement allows for more flexible configuration of the MoE training process, improving usability and performance tuning.
This reverts commit 4aa2272.
…zation lazy to call it just before combine or dispatch ops
…ver distinctions. Updated argument naming conventions in `MoriCombine` to include `sender_` and `recv_` prefixes, enhancing code readability and preventing potential bugs related to argument misalignment. Improved documentation to detail the expected shapes and roles of input tensors, addressing issues that led to backward gradient errors. Additionally, ensured correct handling of probabilities and indices for both forward and backward passes in the MORI operations.
This commit introduces the MORI installation steps in the Dockerfile. It also updates the comments for argument validation logic.
There was a problem hiding this comment.
Pull request overview
This PR integrates ROCm/MORI as an additional backend for Megatron-Core’s flex MoE token dispatcher, enabling MORI-based expert-parallel dispatch/combine alongside existing DeepEP and HybridEP paths. It also adds padding-aware routing behavior, argument/config plumbing for MORI sizing, ROCm container install steps, and expands unit tests and example scripts to exercise the new backend.
Changes:
- Add MORI dispatch/combine implementation + lifecycle helpers (op caching, shmem init/finalize, comm stream) and wire it into the flex token dispatcher.
- Propagate/handle padding masks in routing and fine-grained GPT execution, and adjust test tolerances/coverage for the new communication path.
- Add MORI-related config/args fields and update ROCm container + example training scripts to enable MORI.
Reviewed changes
Copilot reviewed 16 out of 16 changed files in this pull request and generated 8 comments.
Show a summary per file
| File | Description |
|---|---|
megatron/core/transformer/moe/fused_a2a.py |
Adds MORI EP dispatch/combine autograd functions, op caching, shmem lifecycle, comm-stream execution helpers. |
megatron/core/transformer/moe/token_dispatcher.py |
Introduces _MoriManager and wires MORI as a flex dispatcher backend. |
megatron/core/transformer/moe/router.py |
Masks router outputs for padding tokens (so padding doesn’t participate in routing/dispatch). |
megatron/core/transformer/transformer_config.py |
Adds MORI backend option + MORI-specific config fields and validation. |
megatron/training/arguments.py |
Auto-derives --moe-mori-max-tokens-per-rank when MORI backend is selected. |
megatron/core/models/gpt/fine_grained_callables.py |
Passes padding mask into MoE routing in fine-grained execution path. |
tests/unit_tests/transformer/moe/test_token_dispatcher.py |
Parametrizes flex tests over MORI + adds MORI shared-op stress tests and teardown hooks. |
tests/unit_tests/transformer/moe/conftest.py |
Adds MORI shmem teardown safety net between tests. |
tests/unit_tests/transformer/moe/test_moe_layer_discrepancy.py |
Relaxes allgather vs alltoall tolerance (fp32 accumulation-order noise). |
tests/unit_tests/models/test_mamba_moe_model.py |
Updates expected config keys to include MORI fields. |
run_unit_tests.sh |
Adjusts marker expression (now includes internal tests by default). |
Dockerfile_rocm.ci |
Installs MORI + build deps into ROCm CI container image. |
examples/qwen3/train_qwen3.sh |
Adds ENABLE_MORI toggle + MORI flex-dispatcher flags. |
examples/qwen3/proxy_mi355x_qwen3_235B_A22B.sh |
Enables MORI by default for this proxy script. |
examples/deepseek_v3/train_deepseekv3.sh |
Adds ENABLE_MORI toggle + MORI flex-dispatcher flags. |
examples/deepseek_v3/proxy_mi355x_deepseekv3_671B.sh |
Enables MORI by default for this proxy script. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
- Added validation to ensure `moe_mori_max_tokens_per_rank` is set when using the MORI EP backend, raising a ValueError if not. - Updated process group registration methods in `fused_a2a.py` to use the correct namespace for improved compatibility. - Introduced a new function `_resolve_mori_kernel_type` to determine the MORI kernel type based on world size or user input, enhancing flexibility in kernel selection. - Adjusted the `dispatch` method in `_MoriManager` to utilize the new `kernel_type` parameter for better control over dispatch operations. - Refined the routing logic in `router.py` to use bitwise operations for masking, improving performance and clarity. - Cleaned up imports in the test file `test_token_dispatcher.py` by removing unnecessary lines.
Collaborator
|
does EpDispatchCombineOp double-buffer its staging region per outstanding routing handle, or is there one shared region? |
- Introduced `enable_mori` flag in `fine_grained_callables.py` to manage memory freeing logic for the MORI backend. - Updated `MoriDispatch` class in `fused_a2a.py` to improve gradient handling and memory management during backward passes. - Added new test cases in `test_cuda_graphed_schedule_chunk_1f1b.py` and `test_schedule_chunk_1f1b.py` to validate MORI functionality, including proper teardown procedures to prevent state leakage between tests. - Enhanced comparison utilities to accommodate floating-point tolerance for MORI, ensuring accurate test results across different backends.
- Updated `MoriDispatch` in `fused_a2a.py` to eliminate unnecessary defensive cloning of dispatch output, enhancing memory management during backward passes. - Introduced `_permute_by_experts` method in `_MoriManager` to fuse the expert permute operation within the dispatch node, streamlining the token processing flow. - Adjusted floating-point comparison tolerances in `utils.py` to accommodate accumulation noise, ensuring accurate results across different backends.
- Updated `MoriCombine.backward` in `fused_a2a.py` to eliminate unnecessary cloning of the dispatch output, allowing for direct use of the gradient view without defensive copies. - Enhanced `_MoriManager` in `token_dispatcher.py` to fuse the expert unpermute operation into the combine node, ensuring serialized execution on the same communication stream and improving efficiency during backward passes.
… functionality - Updated test cases in `test_schedule_chunk_1f1b.py` and `test_schedule_layer_1f1b.py` to utilize the new `get_valid_token_dispatcher_types()` function, streamlining dispatcher type handling. - Enhanced `reinitialize_model_parallel_for_mori()` in `utils.py` to allow for independent tensor model parallel size configuration, improving flexibility in test setups. - Removed deprecated dispatcher configuration functions to simplify the codebase and reduce redundancy.
…n floating-point evaluations. Adjusted `MORI_COMPARE_ATOL` to 1.5e-1 and `MORI_COMPARE_RTOL` to 5e-2, reflecting improved handling of accumulation noise across different backends.
Collaborator
Author
One shared staging region per op, so we cannot have 2 dispatch/combine calls from the same op in flight at the same time. |
- Introduced `MORI_EP_LAUNCH_CONFIG_MODE` environment variable to control the launch configuration mode for MORI, defaulting to "AUTO".
- Removed FP8 dispatch integration and updated the dispatch configuration to use the standard data type and a scale dimension of zero
wenchenvincent
approved these changes
Jul 8, 2026
Collaborator
Author
|
While CI is having issues, all the tests passed locally. |
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
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.
Add MORI as a Token Dispatcher
Summary
This PR integrates MORI EP as a new MoE token-dispatch/combine backend for the flex token dispatcher, giving AMD/ROCm GPUs an optimized expert-parallel all-to-all path alongside the existing DeepEP and HybridEP backends. It also adds the supporting padding-aware routing fixes, argument plumbing, example training scripts, container build steps, and unit-test coverage needed to run DeepSeek-V3 and Qwen3 MoE models with MORI.
Enable it with:
--moe-token-dispatcher-type flex --moe-flex-dispatcher-backend moriMotivation
The flex dispatcher previously only supported
deepepandhybridep. MORI provides a ROCm-native dispatch/combine kernel set. This PR wires MORI into Megatron's flex-dispatcher abstraction so MoE training on AMD hardware gets efficient, overlap-capable expert parallelism without changing the model code.What's included
New MORI backend (
megatron/core/transformer/moe/fused_a2a.py)MoriDispatch/MoriCombineautograd.Functions implementing the forward/backward all-to-all via MORI'sEpDispatchCombineOp.get_mori_op), shmem lifecycle management (init_mori_shmem/finalize_mori_shmem/reset_mori_op), and a dedicated high-priority CUDA comm stream so dispatch/combine can overlap non-dependent compute (mirrors DeepEP'sasync_finish/allocate_on_comm_streamcontract).IntraNodefor ≤8 ranks,InterNodeV1otherwise) and optional FP8 dispatch support.-1sentinels), with sync-free per-expert token counting.Dispatcher manager (
token_dispatcher.py)_MoriManager(_DispatchManager)mirroring the DeepEP workflow (setup_metadata→dispatch→ permute → combine), including multihot conversion (withpermute_fusionsupport), FP8 quantization padding, and a deferred-sync prefetch of per-expert counts to keep host reads off the critical path.MoEFlexTokenDispatchernow constructs the MORI manager whenmoe_flex_dispatcher_backend == "mori".Config & args
transformer_config.py: addmoritomoe_flex_dispatcher_backend, plus newmoe_mori_kernel_typeandmoe_mori_max_tokens_per_rankfields, with validation (MORI requires the flex dispatcher and is mutually exclusive with DeepEP).arguments.py: auto-derive--moe-mori-max-tokens-per-rankinvalidate_argsfrommicro_batch_size,seq_length, TP/SP, and CP when not explicitly set.Examples & build
Dockerfile_rocm.ci: install MORI (pinned commit) and itslibpci-dev/libibverbs-devdependencies.ENABLE_MORItoggle and the correspondingmoe_options.Tests
test_token_dispatcher.py: parametrize flex dispatcher tests over["deepep", "hybridep", "mori"](skipping when MORI isn't available) and add aTestMoriSharedOpsuite stressing process-wide op reuse across layers/steps.conftest.py: proper MORI shmem/op teardown between parametrized cases to isolate EP state.test_moe_layer_discrepancy.py: relax tolerances to account for floating-point accumulation noise.How to test
ENABLE_MORI=trueusing the provided example scripts.Notes / follow-ups
moe_mori_max_tokens_per_rankis auto-derived; pass an explicit value only to over-provision the symmetric SHMEM heap for variable-batch scenarios.