nemo_automodel.components.moe.config

MoE model configuration.

Module Contents

Classes

Name	Description
`MoEConfig`	Configuration for routed and shared MoE expert modules.
`MoEMetricsConfig`	Configuration for MoE load balance metrics logging.

API

class nemo_automodel.components.moe.config.MoEConfig(
    n_routed_experts: int,
    n_shared_experts: int,
    n_activated_experts: int,
    n_expert_groups: int,
    n_limited_groups: int,
    train_gate: bool,
    gate_bias_update_factor: float,
    aux_loss_coeff: float,
    score_func: str,
    route_scale: float,
    dim: int,
    inter_dim: int,
    moe_inter_dim: int,
    norm_topk_prob: bool,
    router_bias: bool = False,
    expert_bias: bool = False,
    expert_activation: typing.Literal['swiglu', 'swigluoai', 'quick_geglu', 'geglu', 'relu2'] = 'swiglu',
    apply_router_weight_after_down: bool = False,
    activation_alpha: float = 1.702,
    activation_limit: float = 7.0,
    swiglu_limit: float = 0.0,
    softmax_before_topk: bool = False,
    router_weights_fp32: bool = False,
    router_weight_uses_score_correction_bias: bool = False,
    dtype: str | torch.dtype = torch.bfloat16,
    gate_dtype: str | torch.dtype | None = None,
    shared_expert_gate: bool = False,
    shared_expert_inter_dim: int | None = None,
    shared_expert_activation: str = 'swiglu',
    force_e_score_correction_bias: bool = False,
    moe_latent_size: int | None = None,
    enable_routing_replay: bool = False
)

Dataclass

Configuration for routed and shared MoE expert modules.

activation_alpha

float = 1.702

activation_limit

float = 7.0

apply_router_weight_after_down

bool = False

aux_loss_coeff

float

dim

int

dtype

str | dtype = torch.bfloat16

enable_routing_replay

bool = False

expert_activation

Literal['swiglu', 'swigluoai', 'quick_geglu', 'geglu', 'relu2'] = 'swiglu'

expert_bias

bool = False

expert_dim

int

Dimension used for expert projections (latent size when set, otherwise model dim).

force_e_score_correction_bias

bool = False

gate_bias_update_factor

float

gate_dtype

str | dtype | None = None

inter_dim

int

moe_inter_dim

int

moe_latent_size

int | None = None

n_activated_experts

int

n_expert_groups

int

n_limited_groups

int

n_routed_experts

int

n_shared_experts

int

norm_topk_prob

bool

route_scale

float

router_bias

bool = False

router_weight_uses_score_correction_bias

bool = False

router_weights_fp32

bool = False

score_func

str

shared_expert_activation

str = 'swiglu'

shared_expert_gate

bool = False

shared_expert_inter_dim

int | None = None

softmax_before_topk

bool = False

swiglu_limit

float = 0.0

train_gate

bool

nemo_automodel.components.moe.config.MoEConfig.__post_init__()

class nemo_automodel.components.moe.config.MoEMetricsConfig(
    enabled: bool = False,
    mode: str = 'brief',
    detailed_every_steps: int | None = None,
    top_k_experts: int = 0
)

Dataclass

Configuration for MoE load balance metrics logging.

detailed_every_steps

int | None = None

How often to log detailed metrics (only used when mode=“detailed”). None means every step.

enabled

bool = False

Whether to enable load balance metric tracking.

mode

str = 'brief'

Logging mode - “brief” for scalar line charts only, “detailed” adds per-layer breakdowns.

top_k_experts

int = 0

Number of top (highest) and bottom (lowest) utilization experts to emit per layer. Reduces wandb key count for models with many experts. Set to 0 to disable per-expert utilization logging entirely.