MoonEP support

megatron/core/models/common/utils.py

@@ -77,6 +77,10 @@ def should_free_input(name, is_moe, config, num_local_experts):
         config.moe_token_dispatcher_type == "flex"
         and config.moe_flex_dispatcher_backend == "ncclep"
     )
+    enable_moonep = (
+        config.moe_token_dispatcher_type == "flex"
+        and config.moe_flex_dispatcher_backend == "moonep"
+    )
     # Define which nodes should free input memory.
     # Since we split the computing graph into multiple nodes, we can manually control
     # when and how to free the input memory.
@@ -87,22 +91,22 @@ def should_free_input(name, is_moe, config, num_local_experts):
     # original bf16 tensors are safe to be freed.
     free_mlp = config.fp8 is not None or config.fp4 is not None
     if not free_mlp:
-        # AlltoAll dispatcher with local_num_experts=1, HybridEP, and NCCL EP all use
+        # AlltoAll dispatcher with local_num_experts=1, HybridEP, MoonEP, and NCCL EP all use
         # identity operation for `dispatch_postprocess`, hence the mlp inputs will be
         # directly passed to GroupedGemm and should be saved for backward pass.
         free_mlp = num_local_experts > 1 or config.moe_token_dispatcher_type != "alltoall"
-        free_mlp = free_mlp and not (enable_hybridep or enable_ncclep)
+        free_mlp = free_mlp and not (enable_hybridep or enable_moonep or enable_ncclep)
 
     free_input_nodes = {
         "mlp": free_mlp,
         "moe_combine": True,
-        # For non-DeepEP/HybridEP/NCCL-EP dispatcher mode, the input is the un-dispatched
+        # For non-DeepEP/HybridEP/MoonEP/NCCL-EP dispatcher mode, the input is the un-dispatched
         # tokens and probs before dispatch A2A and it's not needed anymore after the
-        # forward pass. For DeepEP, HybridEP, and NCCL EP dispatcher mode, they are all
+        # forward pass. For DeepEP, HybridEP, MoonEP, and NCCL EP dispatcher mode, they are all
         # needed in backward pass and cannot be freed.
         # If moe_preprocess is in cuda graph scope, tokens and probs are fixed size
         # tensors, so they cannot be freed.
-        "moe_dispatch": not (enable_deepep or enable_hybridep or enable_ncclep)
+        "moe_dispatch": not (enable_deepep or enable_hybridep or enable_moonep or enable_ncclep)
         and (CudaGraphModule.moe_preprocess not in config.cuda_graph_modules),
     }
 

megatron/core/models/gpt/fine_grained_callables.py

@@ -63,6 +63,10 @@ def build_transformer_layer_callables(layer: TransformerLayer):
         layer.config.moe_token_dispatcher_type == "flex"
         and layer.config.moe_flex_dispatcher_backend == "ncclep"
     )
+    enable_moonep = (
+        layer.config.moe_token_dispatcher_type == "flex"
+        and layer.config.moe_flex_dispatcher_backend == "moonep"
+    )
 
     def submodule_pre_dispatch_forward(node: ScheduleNode, hidden_states: torch.Tensor):
         """
@@ -160,7 +164,7 @@ def submodule_dispatch_forward(
         Dispatches tokens to the experts based on the router output.
         """
         token_dispatcher = layer.mlp.token_dispatcher
-        if enable_deepep or enable_hybridep or enable_ncclep:
+        if enable_deepep or enable_hybridep or enable_moonep or enable_ncclep:
             # update token_probs to be the detached version, prevents
             # backward graph from connecting to pre_dispatch_computation submodule
             token_dispatcher._comm_manager.token_probs = probs
@@ -186,16 +190,16 @@ def submodule_moe_forward(node: ScheduleNode, dispatched_tokens: torch.Tensor):
         """
         dispatched_probs = node.layer_state.dispatched_probs
         token_dispatcher = layer.mlp.token_dispatcher
-        if enable_deepep or enable_hybridep or enable_ncclep:
+        if enable_deepep or enable_hybridep or enable_moonep or enable_ncclep:
             # update dispatched_probs to be detached version, prevents
             # backward graph from connecting to dispatch submodule
             token_dispatcher._comm_manager.dispatched_probs = dispatched_probs
 
         expert_output, _ = layer.mlp.routed_experts_compute(dispatched_tokens, dispatched_probs)
 
-        # For HybridEP and NCCL EP, tokens_per_expert is generated on comm stream, as the
+        # For HybridEP, MoonEP, and NCCL EP, tokens_per_expert is generated by dispatch, as the
         # input to `routed_experts_compute`, a ref is needed to prevent it from being freed.
-        if enable_hybridep or enable_ncclep:
+        if enable_hybridep or enable_moonep or enable_ncclep:
             tokens_per_expert = token_dispatcher._comm_manager.get_number_of_tokens_per_expert()
             node.layer_state.tokens_per_expert = tokens_per_expert
 

megatron/core/parallel_state.py

@@ -2505,10 +2505,14 @@ def get_all_ranks():
 
 def destroy_model_parallel():
     """Set the groups to none."""
-    # Release the NCCL EP context (if the 'ncclep' flex dispatcher bootstrapped one) before the
-    # process group's communicator is torn down. TE registers an atexit ep_finalize that would
-    # otherwise run after dist.destroy_process_group() and hit a "corrupted comm object" at exit.
-    # Idempotent and a no-op when NCCL EP was never bootstrapped.
+    # Release flex-dispatcher contexts before their process-group communicators are torn down.
+    # The finalizers are idempotent no-ops when their corresponding backends were never used.
+    try:
+        from megatron.core.transformer.moe.fused_a2a import moonep_finalize
+
+        moonep_finalize()
+    except Exception:  # finalize must never block teardown
+        pass
     try:
         from megatron.core.transformer.moe.fused_a2a import nccl_ep_finalize
 

megatron/core/transformer/moe/README.md

@@ -309,8 +309,53 @@ After routing, tokens are **dispatched** to the GPU hosting the assigned expert.
 | **alltoall** | NCCL-based All-to-All communication for token exchange | Standard EP > 1 setups | `--moe-token-dispatcher-type alltoall` |
 | **FlexDispatcher with [DeepEP](https://github.com/deepseek-ai/DeepEP) backend** | Removes redundant tokens during cross-node communication, fuses intra/inter-node communication into single kernel | Cross-node EP, fine-grained MoE (DeepSeek-V3) | `--moe-token-dispatcher-type flex --moe-flex-dispatcher-backend deepep` |
 | **FlexDispatcher with [HybridEP](https://github.com/deepseek-ai/DeepEP/tree/hybrid-ep) backend** | NVIDIA's optimized dispatcher using TMA and IBGDA, fewer SMs, native MNNVL support | GB200 NVL72, Multi-Node NVLink | `--moe-token-dispatcher-type flex --moe-flex-dispatcher-backend hybridep` |
+| **FlexDispatcher with MoonEP backend** | Prefetches redundant experts into NVLink-visible slots while retaining Megatron grouped parameters for optimization and checkpoints | Eager BF16 training on one NVLink-connected node | `--moe-token-dispatcher-type flex --moe-flex-dispatcher-backend moonep` |
 | **allgather** | Gathers all tokens to each GPU, no inter-GPU token movement | TP-only setups, small EP, large Top-K | `--moe-token-dispatcher-type allgather` |
 
+#### MoonEP
+
+MoonEP is an optional external dependency and is not installed by Megatron Core. Its initial
+integration supports eager BF16 training on one NVLink-connected host with expert tensor
+parallelism set to 1. Enable it with:
+
+```bash
+--moe-token-dispatcher-type flex
+--moe-flex-dispatcher-backend moonep
+--moe-grouped-gemm
+--moe-single-grouped-weight
+--use-transformer-engine-op-fuser
+--gradient-accumulation-fusion
+--disable-bias-linear
+--moe-router-dtype fp32
+```
+
+The supported expert activations are fused SwiGLU, quick-GeGLU, and weighted squared ReLU. Latent
+MoE is supported with `--moe-latent-size`; routing remains in the model hidden dimension while
+MoonEP dispatch, expert computation, and combine operate in the latent dimension. Packed and
+GLU-interleaved FC1 layouts remain one contiguous projection. The registered
+`moe_single_grouped_weight` parameters remain the optimizer and checkpoint source of truth;
+MoonEP's `E+B` runtime weights and FP32 gradient storage are mirrors, so the distributed optimizer
+and existing sharded checkpoint layout remain supported.
+
+The dispatched activation and probability tensors retain MoonEP's fixed `[NvS, H]` and `[NvS]`
+capacity through the fused expert MLP. Device-side `E+B` counts delimit the live rows; unused tail
+rows are ignored by grouped GEMM and combine. This is padding-only—the router's capacity-based
+drop-and-pad mode remains disabled, and MoonEP does not drop routed tokens. Steady-state dispatch,
+FC2 output/combine, combine-backward, and FC1 dgrad/dispatch-backward use symmetric token buffers
+directly, without full-activation boundary copies. Per-forward dispatch buffers are pooled until
+FC1 backward so multiple outstanding microbatches cannot overwrite saved activations; the pool
+grows only to the maximum in-flight depth and is then reused. Router weights remain in independent
+small tensors. Steady-state dispatch, expert compute, weight publication, and gradient reduction do
+not read routing counts back to the CPU or use device-wide synchronization.
+
+MoonEP requires fixed, equal local token counts across its ranks, top-k no greater than 32,
+128-element-aligned expert projection dimensions (including `moe_latent_size`, when set), and
+VMM-aligned per-rank expert chunks. CUDA graphs, FP8/FP4, delayed wgrad, expert
+communication/shared-expert overlap, and capacity-based dropping or padding are rejected.
+Megatron's
+`destroy_model_parallel()` tears down live MoonEP VMM buffers collectively before process-group
+teardown.
+
 ### Upcycling
 Use `--moe-use-upcycling` to enable upcycling, which loads the dense model from the `--load` directory, converts it to an MoE model at runtime, and starts training. The converted model is saved to the `--save` path before training begins. Upcycling is built on distributed checkpointing, supporting parallel modes different from existing dense checkpoints, such as arbitrary expert parallelism during upcycling.
 
@@ -549,6 +594,7 @@ For MoE models, certain configurations may prevent CUDA Graph capture of MoE lay
 | Argument | Description | Default |
 |----------|-------------|---------|
 | --moe-token-dispatcher-type | Dispatcher: allgather, alltoall, flex | allgather |
+| --moe-flex-dispatcher-backend | Flex backend: deepep, hybridep, moonep, ncclep | deepep |
 | --moe-enable-deepep | Enable DeepEP (with flex) | False |
 | --moe-expert-capacity-factor | Capacity factor | None |
 | --moe-pad-expert-input-to-capacity | Pad to capacity | False |