Is this page helpful?

core.transformer.moe.megakernel.mok.backend#

Minimal Mixture-of-Kittens adapter for MCore MoE training experiments.

This module intentionally keeps the integration narrow: MCore owns routing, parameters, DDP, the optimizer, and the logical checkpoint format, while MoK replaces dispatch, routed/shared expert computation, and combine.

Module Contents#

Classes#

MoKMegakernel

Execute MOK using trainable parameters owned by native MCore modules.

API#

class core.transformer.moe.megakernel.mok.backend.MoKMegakernel(
config: megatron.core.transformer.transformer_config.TransformerConfig,
ep_group: torch.distributed.ProcessGroup,
routed_experts: torch.nn.Module,
shared_experts: torch.nn.Module,
num_local_experts: int,
)#

Bases: megatron.core.transformer.moe.megakernel.backend.MegakernelBackend

Execute MOK using trainable parameters owned by native MCore modules.

Initialization

property routed_fc1_parameters: tuple[torch.nn.Parameter, ...]#

Return the MCore-owned routed FC1 parameters.

property routed_fc2_parameters: tuple[torch.nn.Parameter, ...]#

Return the MCore-owned routed FC2 parameters.

property autograd_routed_parameters: tuple[torch.nn.Parameter, ...]#

Return routed parameters tracked by the MOK autograd bridge.

main_grad_arguments()#

Return MOK logical main grads and optional per-expert descriptors.

finish_routed_weight_gradients() → tuple[torch.Tensor, ...]#

Finalize and return routed weight gradients.

_register_shared_weights(
shared: torch.nn.Module,
*,
use_output_gate: bool = False,
) → None#

Validate and alias MCore-owned native BF16 shared weights.

quantized_routed_weights()#

Alias native payloads and refresh the converted scales after parameter sync.

With overlapped parameter gathering, DDP’s pre-hook on our aliases finishes all-gather before forward. Grad-buffer reuse requantizes into TE storage and clears the staging buffer; the MOK payload views reference TE storage, not that temporary buffer. Scale conversion still needs to run after each optimizer update, in place so descriptors keep their addresses. It does not requantize the weights.

_routed_weight_view(parameters, *, rows, columns, cached_view)#

Select the storage adapter using this backend’s layout and precision.

sharded_state_dict(prefix='', sharded_offsets=(), metadata=None)#

Emit no aliases; native expert modules own all checkpoint shards.

_save_to_state_dict(destination, prefix, keep_vars)#

Emit no aliases in regular state dicts either.

_load_from_state_dict(
state_dict,
prefix,
local_metadata,
strict,
missing_keys,
unexpected_keys,
error_msgs,
)#

Skip parameter aliases and invalidate state derived from native weights.

_apply(*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.

forward(
hidden_states: torch.Tensor,
probs: torch.Tensor,
routing_map: torch.Tensor,
) → torch.Tensor#

Run the fused MoE megakernel (dispatch, routed/shared experts, combine).

Parameters:

  • 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.