core.transformer.moe.training_matches_inference

The training MoE forward of moe_training_matches_inference.

The layer’s output comes from a value pass through training’s all-to-all dispatcher with generation’s FlashInfer expert kernel and its finalize arithmetic. The backward comes from the layer’s unchanged training forward, which the activation checkpoint recomputes.

Module Contents

Functions

training_matches_inference_forward	Forward values from the value pass, gradients from the layer’s training forward.
_value_pass	The layer’s training forward with generation’s routes, expert kernel and finalize.
_route	Route as generation does, returning [tokens, topk] probs and expert ids.
_postprocess	layer.postprocess on the unpermuted output, for the all-to-all dispatcher.
_drop_value_pass_weights
release_inference_value_pass_weights	Free the value pass’s cached kernel weights in model.
expert_value_pass	Generation’s FlashInfer kernel, on the training parameters, for each dispatched row.
_expert_weights
_source_versions	The address and version of every tensor the kernel weights are built from.
_finish_param_gather	Run the expert linears’ pre-forward hooks, which finish DDP’s parameter all-gather.
_build_value_pass_weights	The value pass’s (MXFP8, NVFP4) kernel weights, one of them None, from the parameters.
expert_params_use_te_mxfp8	Whether the expert parameters hold TE MXFP8 storage (training’s MXFP8 form).
mxfp8_training_stack	Requantize the expert parameters into generation’s stacked MXFP8 layout.

Data

_value_pass_weights
_hooked_experts
_optimizer_step_hook

API

core.transformer.moe.training_matches_inference.training_matches_inference_forward(

layer: megatron.core.transformer.moe.moe_layer.MoELayer,

training_forward: Callable,

hidden_states: torch.Tensor,

intermediate_tensors=None,

padding_mask: Optional[torch.Tensor] = None,

input_ids: Optional[torch.Tensor] = None,

)

Forward values from the value pass, gradients from the layer’s training forward.

Without grad, as in the log-prob forward, only the value pass runs. With grad, the activation checkpoint runs the value pass for the output, and its recompute runs training_forward, which builds the backward.

core.transformer.moe.training_matches_inference._value_pass(

layer: megatron.core.transformer.moe.moe_layer.MoELayer,

hidden_states: torch.Tensor,

padding_mask,

input_ids,

)

The layer’s training forward with generation’s routes, expert kernel and finalize.

Training’s all-to-all dispatcher sends each route as its own row and the FlashInfer kernel computes every row. The batch-invariant unpermute then adds a token’s rows on each EP rank as the kernel’s finalize does in generation, and adds the ranks as generation’s combine does.

core.transformer.moe.training_matches_inference._route(

layer: megatron.core.transformer.moe.moe_layer.MoELayer,

hidden_states: torch.Tensor,

input_ids,

)

Route as generation does, returning [tokens, topk] probs and expert ids.

The FlashInfer finalize sums each token’s routes in the order listed, so the value pass needs generation’s top-k order, which training’s routing map does not keep. Generation’s router path also leaves the training router’s expert-bias token counts to the recompute. It takes [tokens, 1, hidden].

core.transformer.moe.training_matches_inference._postprocess(

layer: megatron.core.transformer.moe.moe_layer.MoELayer,

output: torch.Tensor,

shared_expert_output,

)

layer.postprocess on the unpermuted output, for the all-to-all dispatcher.

core.transformer.moe.training_matches_inference._value_pass_weights: weakref.WeakKeyDictionary

‘WeakKeyDictionary(…)’

core.transformer.moe.training_matches_inference._hooked_experts: weakref.WeakSet

‘WeakSet(…)’

core.transformer.moe.training_matches_inference._optimizer_step_hook

None

core.transformer.moe.training_matches_inference._drop_value_pass_weights(module, *_)

core.transformer.moe.training_matches_inference.release_inference_value_pass_weights(model: torch.nn.Module) → None

Free the value pass’s cached kernel weights in model.

For a caller that needs the memory, for example before colocated generation. The next value pass rebuilds them from the parameters.

core.transformer.moe.training_matches_inference.expert_value_pass(

experts: megatron.core.transformer.moe.experts.InferenceGroupedMLP,

hidden_states,

tokens_per_expert,

)

Generation’s FlashInfer kernel, on the training parameters, for each dispatched row.

Each row is one route, sorted by local expert. The kernel gets it as a top-1 token routed to that expert with probability 1.0, so its finalize returns the expert’s output row unchanged. The routed MXFP8 kernel has no tile config for a few rows per local expert, so its launches are padded with unrouted rows to more than 128 per local expert, where its tile is its largest; the row values do not depend on the tile.

The kernel weights are built from the parameters and reused until an in-place write to them, an optimizer step, a state_dict() or checkpoint load, or release_inference_value_pass_weights. TE writes into MXFP8 parameters, as in an optimizer step or a copy into a state_dict() tensor, bypass the version counter, so a process-wide optimizer step post-hook and state_dict and load_state_dict post-hooks on the experts, all registered here, drop the weights; the source tensors’ versions catch other in-place writes. A reference-model swap that copies into state_dict() tensors, as NeMo-RL’s does, calls state_dict() first. Neither the training pipeline nor the experts module keeps any value-pass state.

core.transformer.moe.training_matches_inference._expert_weights(

experts: megatron.core.transformer.moe.experts.InferenceGroupedMLP,

)

core.transformer.moe.training_matches_inference._source_versions(

experts: megatron.core.transformer.moe.experts.InferenceGroupedMLP,

) → tuple

The address and version of every tensor the kernel weights are built from.

A TE MXFP8 parameter is represented by the FP8 data and scales it dequantizes from, since state_dict() hands out a new TE tensor over them rather than a view of the parameter.

core.transformer.moe.training_matches_inference._finish_param_gather(

experts: megatron.core.transformer.moe.experts.InferenceGroupedMLP,

) → None

Run the expert linears’ pre-forward hooks, which finish DDP’s parameter all-gather.

The value pass reads the weights without running linear_fc1 or linear_fc2. Only Megatron DDP’s hooks are known to need nothing from that forward, so any other pre-forward hook on them is rejected rather than run out of context.

core.transformer.moe.training_matches_inference._build_value_pass_weights(

experts: megatron.core.transformer.moe.experts.InferenceGroupedMLP,

)

The value pass’s (MXFP8, NVFP4) kernel weights, one of them None, from the parameters.

core.transformer.moe.training_matches_inference.expert_params_use_te_mxfp8(

experts: megatron.core.transformer.moe.experts.InferenceGroupedMLP,

) → bool

Whether the expert parameters hold TE MXFP8 storage (training’s MXFP8 form).

core.transformer.moe.training_matches_inference.mxfp8_training_stack(

experts,

linear_name: str,

backend: str,

) → megatron.core.inference.quantization.mxfp8_tensor.MXFP8Tensor

Requantize the expert parameters into generation’s stacked MXFP8 layout.

Same dequantize-then-from_bf16 route, per expert, as quantize_model_to_mxfp8, so the bits match generation’s.