From 9d91b1108da6df41e722b82b18742110f2a77d81 Mon Sep 17 00:00:00 2001 From: HaochenYuan Date: Sun, 13 Sep 2026 18:36:40 -0700 Subject: [PATCH 1/2] fix attributes discard in mxfp8 Signed-off-by: HaochenYuan --- .../transformer/moe/megakernel/mok/backend.py | 35 +++++++++++++++++++ 1 file changed, 35 insertions(+) diff --git a/megatron/core/transformer/moe/megakernel/mok/backend.py b/megatron/core/transformer/moe/megakernel/mok/backend.py index 65746205385..46e6a552985 100644 --- a/megatron/core/transformer/moe/megakernel/mok/backend.py +++ b/megatron/core/transformer/moe/megakernel/mok/backend.py @@ -10,6 +10,8 @@ from typing import TYPE_CHECKING +from types import MethodType + import torch from torch import nn from torch.distributed import ProcessGroup @@ -308,6 +310,39 @@ def _load_from_state_dict( self._split_main_grad_descriptor_cache = None self.is_first_microbatch = True + def _apply(self, *args, **kwargs): + """Preserve externally attached parameter attributes across device/dtype moves. + + The routed/shared weights registered on this module are ALIASES of parameters owned by TE + modules (``TEGroupedLinear`` etc.). ``torch.nn.Module._apply`` moves traceable wrapper + subclasses (TE quantized tensors such as ``MXFP8Tensor``) with ``torch.utils.swap_tensors``, + which exchanges the parameter object's whole ``__dict__``. TE's own modules re-attach the + externally attached state afterwards (``TransformerEngineBaseModule._apply``), but this plain + module's visit of the *same* object (e.g. during ``model.cuda()``) swaps it again and nothing + restores it, so mcore's ``allreduce`` / ``partition_dim`` / ``tensor_model_parallel`` stamps + and TE's ``_high_precision_init_val`` are lost. DDP / the layer-wise optimizer then treat + the routed experts as data-parallel replicas (grads all-reduced across EP ranks, master + weights sharded over dp_cp, identical experts across the EP group). Mirror TE's fix here. + + """ + snapshots = { + name: (param, dict(param.__dict__)) + for name, param in self._parameters.items() + if param is not None + } + out = super()._apply(*args, **kwargs) + for name, (old_param, attrs) in snapshots.items(): + new_param = self._parameters.get(name) + if new_param is None: + continue + for key, value in attrs.items(): + if key in new_param.__dict__: + continue + if isinstance(value, MethodType) and value.__self__ is old_param: + value = MethodType(value.__func__, new_param) + setattr(new_param, key, value) + return out + def forward( self, hidden_states: torch.Tensor, probs: torch.Tensor, routing_map: torch.Tensor ) -> torch.Tensor: From 4ead3ad5f419ab153a3de135204eaf1ea7e0adb6 Mon Sep 17 00:00:00 2001 From: HaochenYuan Date: Mon, 14 Sep 2026 01:27:02 -0700 Subject: [PATCH 2/2] fix linting Signed-off-by: HaochenYuan --- .../transformer/moe/megakernel/mok/backend.py | 31 ++++++++++++------- 1 file changed, 20 insertions(+), 11 deletions(-) diff --git a/megatron/core/transformer/moe/megakernel/mok/backend.py b/megatron/core/transformer/moe/megakernel/mok/backend.py index 46e6a552985..64b5231bc17 100644 --- a/megatron/core/transformer/moe/megakernel/mok/backend.py +++ b/megatron/core/transformer/moe/megakernel/mok/backend.py @@ -8,9 +8,8 @@ from __future__ import annotations -from typing import TYPE_CHECKING - from types import MethodType +from typing import TYPE_CHECKING import torch from torch import nn @@ -315,15 +314,15 @@ def _apply(self, *args, **kwargs): The routed/shared weights registered on this module are ALIASES of parameters owned by TE modules (``TEGroupedLinear`` etc.). ``torch.nn.Module._apply`` moves traceable wrapper - subclasses (TE quantized tensors such as ``MXFP8Tensor``) with ``torch.utils.swap_tensors``, - which exchanges the parameter object's whole ``__dict__``. TE's own modules re-attach the - externally attached state afterwards (``TransformerEngineBaseModule._apply``), but this plain - module's visit of the *same* object (e.g. during ``model.cuda()``) swaps it again and nothing - restores it, so mcore's ``allreduce`` / ``partition_dim`` / ``tensor_model_parallel`` stamps - and TE's ``_high_precision_init_val`` are lost. DDP / the layer-wise optimizer then treat - the routed experts as data-parallel replicas (grads all-reduced across EP ranks, master - weights sharded over dp_cp, identical experts across the EP group). Mirror TE's fix here. - + subclasses (TE quantized tensors such as ``MXFP8Tensor``) with + ``torch.utils.swap_tensors``, which exchanges the parameter object's whole ``__dict__``. + TE's own modules re-attach the externally attached state afterwards + (``TransformerEngineBaseModule._apply``), but this plain module's visit of the *same* + object (e.g. during ``model.cuda()``) swaps it again and nothing restores it, so mcore's + ``allreduce`` / ``partition_dim`` / ``tensor_model_parallel`` stamps and TE's + ``_high_precision_init_val`` are lost. DDP / the layer-wise optimizer then treat the routed + experts as data-parallel replicas (grads all-reduced across EP ranks, master weights + sharded over dp_cp, identical experts across the EP group). Mirror TE's fix here. """ snapshots = { name: (param, dict(param.__dict__)) @@ -346,6 +345,16 @@ def _apply(self, *args, **kwargs): def forward( self, hidden_states: torch.Tensor, probs: torch.Tensor, routing_map: torch.Tensor ) -> torch.Tensor: + """Run the fused MoE megakernel (dispatch, routed/shared experts, combine). + + Args: + hidden_states: Input activations of shape ``[..., hidden_size]``. + probs: Router probabilities, flattened to ``[num_tokens, num_experts]``. + routing_map: Token-to-expert routing map matching ``probs``. + + Returns: + MoE output with the same shape as ``hidden_states``. + """ original_shape = hidden_states.shape x = hidden_states.reshape(-1, original_shape[-1]).contiguous() probs = probs.reshape(x.shape[0], -1)