st-moe-pytorch

The MoE class implements the core mixture-of-experts logic. It allows you to increase the number of parameters in your model without increasing computation by routing tokens to a subset of experts.

Parameters

• dim: Dimension of the input.
• num_experts: Number of experts to create.
• gating_top_n: Number of experts to route to (e.g., 2 for top-2 gating).
• threshold_train: Threshold to accept a token for the second expert and beyond during training.
• threshold_eval: Threshold for routing during evaluation.
• capacity_factor_train: Fixed capacity per batch for experts during training (should be $\ge 1$).
• capacity_factor_eval: Fixed capacity per batch for experts during evaluation (should be $\ge 1$).
• balance_loss_coef: Multiplier for the auxiliary expert balancing loss.
• router_z_loss_coef: Weight for the router z-loss.

Returns

The call returns the processed tensor and three auxiliary loss scalars:

1. total_aux_loss: The sum of all auxiliary losses (must be added to your main loss).
2. balance_loss: Unweighted breakdown for logging.
3. router_z_loss: Unweighted breakdown for logging.

import torch
from st_moe_pytorch import MoE

moe = MoE(
    dim = 512,
    num_experts = 16,
    gating_top_n = 2,
    threshold_train = 0.2,
    threshold_eval = 0.2,
    capacity_factor_train = 1.25,
    capacity_factor_eval = 2.,
    balance_loss_coef = 1e-2,
    router_z_loss_coef = 1e-3,
)

inputs = torch.randn(4, 1024, 512)
out, total_aux_loss, balance_loss, router_z_loss = moe(inputs)

The SparseMoEBlock provides a complete mixture-of-experts block designed for use within a transformer architecture. It includes RMSNorm and residual connections, and allows you to specify whether to add feedforward (FF) layers before or after the MoE component for improved stability.

Parameters

• moe: An instance of the MoE class.
• add_ff_before: Boolean indicating if a feedforward layer should be added before the MoE.
• add_ff_after: Boolean indicating if a feedforward layer should be added after the MoE.

Returns

The block returns the processed tensor and the aggregated auxiliary loss:

• out: The output tensor.
• total_aux_loss: The sum of all auxiliary losses (must be added to your main loss).
• balance_loss: Unweighted breakdown for logging.
• router_z_loss: Unweighted breakdown for logging.

import torch
from st_moe_pytorch import MoE, SparseMoEBlock

moe = MoE(
    dim = 512,
    num_experts = 16,
    gating_top_n = 2,
    threshold_train = 0.2,
    threshold_eval = 0.2,
    capacity_factor_train = 1.25,
    capacity_factor_eval = 2.,
    balance_loss_coef = 1e-2,
    router_z_loss_coef = 1e-3,
)

inputs = torch.randn(4, 1024, 512)

moe_block = SparseMoEBlock(
    moe,
    add_ff_before = True,
    add_ff_after = True
)

out, total_aux_loss, balance_loss, router_z_loss = moe_block(inputs)