nemo_automodel.components.moe.layers

Module Contents

Classes

Name	Description
`FakeBalancedGate`	Load balanced gate implementation, spreads tokens uniformly across all experts.
`Gate`	Gating mechanism for routing inputs in a mixture-of-experts (MoE) model.
`MLP`	Multi-Layer Perceptron (MLP) used as a feed-forward layer.
`MoE`	Mixture-of-Experts (MoE) module.
`_GateRoutingCore`	Parameterless, fixed-shape portion of a learned MoE router.

Functions

Name	Description
`_clamped_swiglu_act`	fp32 clamped SwiGLU on separate gate / up projections, cast back to the input dtype.
`_init_weights`	-
`_shared_expert_stream`	Return the per-device side stream used for shared-expert overlap (created lazily, one per process).

Data

_SHARED_EXPERT_STREAMS

API

class nemo_automodel.components.moe.layers.FakeBalancedGate(
    config: nemo_automodel.components.moe.config.MoEConfig,
    skip_first_n_experts: int = 0,
    noise: float = 0.0
)

Bases: Module

Load balanced gate implementation, spreads tokens uniformly across all experts. The rationale for this class is to do performance experiments to understand how the load imbalance with real data is impacting end-to-end performance.

When noise > 0, random perturbation is added to mimic realistic routing imbalance. A noise value of 0.0 gives perfectly balanced assignment, while 1.0 gives fully random expert selection and non-uniform weights.

bias_update_factor

= 0.0

n_activated_experts

= config.n_activated_experts

n_routed_experts

= config.n_routed_experts

nemo_automodel.components.moe.layers.FakeBalancedGate.forward(
    x: torch.Tensor,
    token_mask: torch.Tensor,
    cp_mesh: torch.distributed.device_mesh.DeviceMesh | None
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]

Forward pass for the gating mechanism.

Parameters:

torch.Tensor

Input tensor.

token_mask

torch.Tensor

Boolean mask indicating valid tokens.

cp_mesh

Optional[DeviceMesh]

Device mesh for context parallel computation.

Returns: torch.Tensor

Routing weights for the selected experts.

nemo_automodel.components.moe.layers.FakeBalancedGate.init_weights(
    buffer_device: torch.device,
    init_std: float = 0.02
) -> None

nemo_automodel.components.moe.layers.FakeBalancedGate.update_bias() -> None

class nemo_automodel.components.moe.layers.Gate(
    config: nemo_automodel.components.moe.config.MoEConfig,
    gate_precision: torch.dtype | None = None
)

Bases: Module

Gating mechanism for routing inputs in a mixture-of-experts (MoE) model.

_cumulative_expert_load

Tensor | None = None

_last_aux_loss

Tensor | None = None

_last_expert_load

Tensor | None = None

_track_load_balance

bool = False

aux_loss_coeff

= config.aux_loss_coeff

bias

Optional[torch.nn.Parameter]

Optional bias term for the gate.

bias_update_factor

= config.gate_bias_update_factor

dim

int = config.dim

Dimensionality of input features.

n_experts

= config.n_routed_experts

n_groups

int = config.n_expert_groups

Number of groups for routing.

norm_topk_prob

= config.norm_topk_prob

route_scale

float = config.route_scale

Scaling factor for routing weights.

router_replay

RouterReplay | None

router_weight_uses_score_correction_bias

= config.router_weight_uses_score_correction_bias

router_weights_fp32

= config.router_weights_fp32

routing_core

= _GateRoutingCore()

score_dtype

score_func

str = config.score_func

Scoring function (‘softmax’, ‘sigmoid’, ‘sigmoid_with_bias’, ‘softmax_with_bias’, or ‘sqrtsoftplus’).

softmax_before_topk

= config.softmax_before_topk

topk

int = config.n_activated_experts

Number of top experts activated for each input.

topk_groups

int = config.n_limited_groups

Number of groups to route inputs to.

train_gate

= config.train_gate

weight

torch.nn.Parameter

Learnable weights for the gate.

nemo_automodel.components.moe.layers.Gate._compute_aux_loss(
    original_scores: torch.Tensor,
    expert_load: torch.Tensor,
    token_mask: torch.Tensor,
    cp_mesh: torch.distributed.device_mesh.DeviceMesh | None
) -> torch.Tensor

Computes the auxiliary loss for load balancing.

Warning: Assumes batch size = 1, if batch size > 1, the aux_loss will be computed across multiple sequences.

Parameters:

original_scores

torch.Tensor

Original scores from the gating mechanism. Shape is [num_tokens, num_experts].

expert_load

torch.Tensor

Load of each expert (number of tokens routed to each expert). Shape is [num_experts].

token_mask

torch.Tensor

Boolean mask indicating valid tokens. Shape is [num_tokens].

cp_mesh

Optional[DeviceMesh]

Device mesh for context parallel computation.

Returns: torch.Tensor

torch.Tensor: Auxiliary loss for load balancing. Shape is [].

nemo_automodel.components.moe.layers.Gate._compute_expert_load(
    indices: torch.Tensor,
    token_mask: torch.Tensor
) -> torch.Tensor

Computes the load of each expert based on the selected indices. Args: indices (torch.Tensor): Indices of the selected experts. Shape is [num_tokens, num_activated_experts]. token_mask (torch.Tensor): Boolean mask indicating valid tokens. Shape is [num_tokens].

Returns: torch.Tensor

torch.Tensor: Load of each expert (number of tokens routed to each expert). Shape is [num_local_experts].

nemo_automodel.components.moe.layers.Gate._load_from_state_dict(
    state_dict: dict[str, torch.Tensor],
    prefix: str,
    local_metadata: dict[str, object],
    strict: bool,
    missing_keys: list[str],
    unexpected_keys: list[str],
    error_msgs: list[str]
) -> None

Load gate state while preserving the routing-bias DTensor layout.

Parameters:

state_dict

dict[str, torch.Tensor]

State mapping whose correction bias is either a full tensor of shape [experts] or an already distributed tensor.

prefix

str

Prefix applied to this module’s state-dict keys.

local_metadata

dict[str, object]

Metadata saved for this module.

strict

bool

Whether state-dict keys must match exactly.

missing_keys

list[str]

Missing state-dict keys collected during loading.

unexpected_keys

list[str]

Unexpected state-dict keys collected during loading.

error_msgs

list[str]

State-dict loading errors collected during loading.

Raises:

RuntimeError: If a full correction-bias tensor targets an incompatible shape or a DTensor that is not replicated on every mesh dimension.

nemo_automodel.components.moe.layers.Gate._local_score_correction_bias() -> torch.Tensor | None

Return the local tensor used to adjust this rank’s routing scores.

Returns: torch.Tensor | None

Tensor of shape [experts] replicated across the gate’s DP/CP mesh, or

nemo_automodel.components.moe.layers.Gate._route_scores(
    scores: torch.Tensor
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]

Apply fixed-shape expert selection and probability math to router logits.

Parameters:

scores

torch.Tensor

Tensor of shape [tokens, experts] containing local router logits.

Returns: torch.Tensor

Tuple containing routing weights of shape [tokens, activated_experts],

nemo_automodel.components.moe.layers.Gate.forward(
    x: torch.Tensor,
    token_mask: torch.Tensor,
    cp_mesh: torch.distributed.device_mesh.DeviceMesh | None
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]

Forward pass for the gating mechanism.

Parameters:

torch.Tensor

Input tensor.

token_mask

torch.Tensor

Boolean mask indicating valid tokens.

cp_mesh

Optional[DeviceMesh]

Device mesh for context parallel computation.

Returns: torch.Tensor

Routing weights for the selected experts.

nemo_automodel.components.moe.layers.Gate.init_weights(
    buffer_device: torch.device,
    init_std: float = 0.02
) -> None

nemo_automodel.components.moe.layers.Gate.update_bias() -> None

Updates the correction bias used in the gate based on the popularity of experts. This function is a NoOp if the gate is not trained.

To avoid routing collapse, and to promote better load balance of experts, DeepSeek-V3 uses a correction mechanism to adjust the scores of experts using a learned bias parameter. The bias parameter is updated based on the popularity of experts, i.e., the number of tokens routed to each expert. If an expert is more popular than the average, its bias term is decreased, and vice versa. This encourages the model to route tokens to less popular experts, promoting better load balance.

class nemo_automodel.components.moe.layers.MLP(
    dim: int,
    inter_dim: int,
    backend: str,
    dtype: torch.dtype = torch.bfloat16,
    activation: str = 'swiglu',
    bias: bool = False,
    swiglu_limit: float = 0.0
)

Bases: Module

Multi-Layer Perceptron (MLP) used as a feed-forward layer.

Supports both gated activations (SwiGLU) and simple activations (ReLU²).

down_proj

nn.Module

Linear layer for hidden-to-output transformation.

gate_proj

nn.Module

Linear layer for gate in gated activations (or up_proj for simple).

is_gated

= is_gated_activation(activation)

swiglu_limit

= float(swiglu_limit)

up_proj

nn.Module

Additional linear layer for gated activations (None for simple).

nemo_automodel.components.moe.layers.MLP.forward(
    x: torch.Tensor
) -> torch.Tensor

Forward pass for the MLP layer.

Parameters:

torch.Tensor

Input tensor.

Returns: torch.Tensor

torch.Tensor: Output tensor after MLP computation.

nemo_automodel.components.moe.layers.MLP.init_weights(
    buffer_device: torch.device,
    init_std: float = 0.02
) -> None

class nemo_automodel.components.moe.layers.MoE(
    config: nemo_automodel.components.moe.config.MoEConfig,
    backend: nemo_automodel.components.models.common.BackendConfig
)

Bases: Module

Mixture-of-Experts (MoE) module.

cp_mesh

DeviceMesh | None = None

dim

int = config.dim

Dimensionality of input features.

experts

nn.ModuleList = GroupedExpertsMoK(config, backend)

List of expert modules.

fc1_latent_proj

fc2_latent_proj

gate

nn.Module

Gating mechanism to route inputs to experts.

n_activated_experts

int = config.n_activated_experts

Number of experts activated for each input.

n_local_experts

int

Number of experts handled locally in distributed systems.

n_routed_experts

int = config.n_routed_experts

Total number of experts in the model.

shared_expert_gate

shared_experts

nn.Module

Shared experts applied to all inputs.

nemo_automodel.components.moe.layers.MoE._join_shared_experts(
    z: torch.Tensor,
    stream: torch.cuda.Stream
) -> None

staticmethod

Make the current stream wait for the side-stream shared-expert output z.

nemo_automodel.components.moe.layers.MoE._shared_expert_output(
    x: torch.Tensor
) -> torch.Tensor

Shared-expert output for x, gated by shared_expert_gate when the model has one.

nemo_automodel.components.moe.layers.MoE._shared_experts_on_side_stream(
    x: torch.Tensor
) -> tuple[torch.Tensor, torch.cuda.Stream]

Run the shared experts for x on the per-device side stream.

Returns the shared-expert output and the stream it was computed on; the caller runs the routed path on the current stream and then joins with _join_shared_experts. Subclasses with their own forward (Kimi K3) reuse both helpers so the stream choreography lives in one place.

nemo_automodel.components.moe.layers.MoE.forward(
    x: torch.Tensor,
    padding_mask: torch.Tensor | None = None,
    cp_mesh: torch.distributed.device_mesh.DeviceMesh | None = None
) -> torch.Tensor

Route tokens through shared and routed experts.

Parameters:

torch.Tensor

Input tensor of shape [..., hidden].

padding_mask

torch.Tensor | NoneDefaults to None

Boolean tensor matching x.shape[:-1] where true entries are padding.

cp_mesh

DeviceMesh | NoneDefaults to None

Optional context-parallel mesh used by the router.

Returns: torch.Tensor

Tensor with the same shape and dtype as x.

nemo_automodel.components.moe.layers.MoE.init_weights(
    buffer_device: torch.device,
    init_std: float = 0.02
) -> None

class nemo_automodel.components.moe.layers._GateRoutingCore()

Bases: Module

Parameterless, fixed-shape portion of a learned MoE router.

The gate projection remains eager so FSDP can unshard its DTensor parameters. The owning gate is passed as a non-tensor control so this child remains an independent, parameterless CUDA graph boundary and survives model copies.

nemo_automodel.components.moe.layers._GateRoutingCore.forward(
    scores: torch.Tensor,
    gate: nemo_automodel.components.moe.layers.Gate
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]

Select experts and compute routing probabilities from projected scores.

nemo_automodel.components.moe.layers._clamped_swiglu_act(
    gate: torch.Tensor,
    up: torch.Tensor,
    limit: float
) -> torch.Tensor

fp32 clamped SwiGLU on separate gate / up projections, cast back to the input dtype.

Same arithmetic as the eager branch of MLP.forward (the DeepSeek-V4 reference Expert.forward); inductor fuses it into one kernel per direction. Dynamo’s default shape policy: the first token count compiles a static graph, the second recompiles once with a symbolic row count and every later shape reuses it — dynamic=False with fullgraph=True would instead raise FailOnRecompileLimitHit after recompile_limit distinct token counts (variable-length batches). CPU inputs keep the eager chain.

nemo_automodel.components.moe.layers._init_weights(
    module,
    buffer_device: torch.device,
    init_std: float = 0.02
)

nemo_automodel.components.moe.layers._shared_expert_stream(
    device: torch.device
) -> torch.cuda.Stream

Return the per-device side stream used for shared-expert overlap (created lazily, one per process).

nemo_automodel.components.moe.layers._SHARED_EXPERT_STREAMS: dict[int, Stream] = {}