diff --git a/megatron/core/transformer/moe/moe_layer.py b/megatron/core/transformer/moe/moe_layer.py index 80d5a04b0f0..c27d05171f3 100644 --- a/megatron/core/transformer/moe/moe_layer.py +++ b/megatron/core/transformer/moe/moe_layer.py @@ -28,11 +28,20 @@ from megatron.core.transformer.transformer_config import TransformerConfig from megatron.core.typed_torch import apply_module from megatron.core.utils import internal_api +from miles_megatron_plugins.true_on_policy.moe_layer_ext import ( + forward_compacted_true_on_policy_padding, + run_direct_sglang_ep_forward, + uses_true_on_policy_moe_kernel, +) try: import transformer_engine as te # pylint: disable=unused-import - from megatron.core.extensions.transformer_engine import TELinear, te_checkpoint + from megatron.core.extensions.transformer_engine import ( + TELinear, + set_save_original_input, + te_checkpoint, + ) HAVE_TE = True except ImportError: @@ -389,10 +398,18 @@ def forward( if padding_mask is not None: padding_mask = padding_mask.transpose(0, 1).bool() + use_true_on_policy_moe_kernel = uses_true_on_policy_moe_kernel(self) + # MoE forward: route -> dispatch -> compute -> combine def custom_forward(hidden_states, intermediate_tensors, padding_mask=None): + shared_expert_output = None try: if "route" in self.fwd_execution_map: + if use_true_on_policy_moe_kernel: + return run_direct_sglang_ep_forward( + self, hidden_states, padding_mask, intermediate_tensors + ) + shared_expert_output = self.shared_experts_compute(hidden_states) probs, routing_map = self.route(hidden_states, padding_mask) hidden_states, probs = self.preprocess(hidden_states, probs, routing_map) @@ -433,7 +450,21 @@ def custom_forward(hidden_states, intermediate_tensors, padding_mask=None): return output, mlp_bias - if self.moe_layer_recompute: + use_compact = ( + use_true_on_policy_moe_kernel + and padding_mask is not None + and intermediate_tensors is None + and padding_mask.dtype == torch.bool + and bool(padding_mask.any().item()) + and not self.use_shared_expert + and not self.config.moe_latent_size + ) + + if use_compact: + outputs = forward_compacted_true_on_policy_padding( + hidden_states, padding_mask, custom_forward + ) + elif self.moe_layer_recompute: if self.config.fp8 or self.config.fp4: outputs = te_checkpoint( custom_forward, @@ -464,6 +495,4 @@ def set_for_recompute_pre_mlp_layernorm(self): # If shared_experts_recompute is used, nothing needs to be done because the checkpoint # function will save the original input tensors. if self.shared_experts is not None and not self.shared_experts_recompute: - from megatron.core.extensions.transformer_engine import set_save_original_input - set_save_original_input(self.shared_experts.linear_fc1) diff --git a/miles_megatron_plugins/true_on_policy/moe.py b/miles_megatron_plugins/true_on_policy/moe.py new file mode 100644 index 00000000000..e7e66d8bcf6 --- /dev/null +++ b/miles_megatron_plugins/true_on_policy/moe.py @@ -0,0 +1,18 @@ +"""True-on-policy MoE re-exports. + +The implementation is split across: + * :mod:`moe_reduce` -- ``sglang_moe_ep_tree_all_reduce`` deterministic EP combine. + * :mod:`moe_experts` -- ``SGLangGroupedMLP`` SGLang fused-expert helpers. + * :mod:`moe_layer_ext` -- EP-invariant local-masked forward and padding compaction. + +Importers should prefer the dedicated submodules; this module is kept as a +backward-compatible re-export point for older call sites. +""" + +from .moe_experts import SGLangGroupedMLP +from .moe_reduce import sglang_moe_ep_tree_all_reduce + +__all__ = [ + "SGLangGroupedMLP", + "sglang_moe_ep_tree_all_reduce", +] diff --git a/miles_megatron_plugins/true_on_policy/moe_reduce.py b/miles_megatron_plugins/true_on_policy/moe_reduce.py new file mode 100644 index 00000000000..301b273c175 --- /dev/null +++ b/miles_megatron_plugins/true_on_policy/moe_reduce.py @@ -0,0 +1,21 @@ +"""Deterministic EP combine for true-on-policy MoE. + +Wraps SGLang's ``tree_all_reduce_sum`` so SGLang rollout and Megatron training +both perform the post-EP sum in the same fixed-tree order. Using this in place +of NCCL's ring all-reduce is what keeps the EP combine bit-exact across the +two engines. +""" + +from __future__ import annotations + +import torch + + +def sglang_moe_ep_tree_all_reduce( + input_: torch.Tensor, + ep_group, +) -> torch.Tensor: + """Mirror SGLang's fixed-order EP sum for true-on-policy MoE combine.""" + from sglang.srt.tp_invariant_ops import tree_all_reduce_sum + + return tree_all_reduce_sum(input_, device_group=ep_group)