Sparse Mixture-of-Experts (MoE): Routing Math, Load Balancing & Expert Parallelism

import torch
import torch.nn as nn
import torch.nn.functional as F

class ClassicalTopKRouter(nn.Module):
    """Switch / Mixtral Top-K Softmax Router with Auxiliary Load-Balancing Loss and Router z-Loss."""
    def __init__(self, d_model: int, num_experts: int = 8, top_k: int = 2, alpha_aux: float = 0.01, c_z: float = 1e-3):
        super().__init__()
        self.num_experts = num_experts
        self.top_k = top_k
        self.alpha_aux = alpha_aux
        self.c_z = c_z
        self.gate = nn.Linear(d_model, num_experts, bias=False)

    def forward(self, h: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
        # h: [T, d_model]
        logits = self.gate(h.float())                           # [T, E] in FP32 for stability
        probs = F.softmax(logits, dim=-1)                       # [T, E]
        topk_probs, topk_idx = torch.topk(probs, self.top_k, dim=-1)
        gates = topk_probs / topk_probs.sum(dim=-1, keepdim=True).clamp(min=1e-12)

        # 1. Switch Auxiliary Load-Balancing Loss: E * sum(f_i * P_i)
        one_hot = F.one_hot(topk_idx, num_classes=self.num_experts).float()  # [T, K, E]
        f_i = one_hot.mean(dim=(0, 1)) * self.num_experts                    # [E], mean is 1.0
        P_i = probs.mean(dim=0)                                              # [E], sum is 1.0
        aux_loss = self.alpha_aux * (f_i * P_i).sum()

        # 2. Router z-Loss: c_z * mean(logsumexp(logits)^2)
        z_loss = self.c_z * (torch.logsumexp(logits, dim=-1) ** 2).mean()
        return gates.to(h.dtype), topk_idx, aux_loss + z_loss, f_i.detach()

class AuxLossFreeRouter(nn.Module):
    """DeepSeek-V3 Auxiliary-Loss-Free Router with Sigmoid Affinity & Dynamic Bias Update."""
    def __init__(self, d_model: int, num_experts: int = 16, top_k: int = 4, gamma_bias: float = 0.01):
        super().__init__()
        self.num_experts = num_experts
        self.top_k = top_k
        self.gamma_bias = gamma_bias
        self.centroids = nn.Linear(d_model, num_experts, bias=False)
        self.register_buffer("expert_bias", torch.zeros(num_experts))

    def forward(self, h: torch.Tensor, update_bias: bool = True) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        # 1. Unbiased sigmoid affinity s_{t,i}
        scores = torch.sigmoid(self.centroids(h.float()))       # [T, E]

        # 2. Biased Top-K selection: argtopk(s_{t,i} + b_i)
        biased_scores = scores + self.expert_bias.unsqueeze(0)
        _, topk_idx = torch.topk(biased_scores, self.top_k, dim=-1)  # [T, K]

        # 3. Gather UNBIASED scores for forward gating weights!
        unbiased_topk = torch.gather(scores, dim=-1, index=topk_idx)
        gates = unbiased_topk / unbiased_topk.sum(dim=-1, keepdim=True).clamp(min=1e-12)

        # 4. Compute token load per expert and update bias without gradients
        counts = torch.bincount(topk_idx.view(-1), minlength=self.num_experts).float()
        mean_count = counts.mean()
        if update_bias and self.training:
            with torch.no_grad():
                self.expert_bias.add_(self.gamma_bias * torch.sign(mean_count - counts))

        load_ratio = counts / mean_count.clamp(min=1e-6)
        return gates.to(h.dtype), topk_idx, load_ratio.detach()

import math
import torch
import torch.nn as nn
import torch.nn.functional as F

class SwiGLUExpert(nn.Module):
    def __init__(self, d_model: int, d_ffn: int):
        super().__init__()
        self.w_gate = nn.Linear(d_model, d_ffn, bias=False)
        self.w_up = nn.Linear(d_model, d_ffn, bias=False)
        self.w_down = nn.Linear(d_ffn, d_model, bias=False)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # x: [num_tokens, d_model] -> [num_tokens, d_model]
        return self.w_down(F.silu(self.w_gate(x)) * self.w_up(x))

class DeepSeekStyleMoELayer(nn.Module):
    """
    Combines always-active Shared Experts with Fine-Grained Routed Experts,
    Aux-Loss-Free dynamic bias routing, and optional Capacity Factor token dropping.
    """
    def __init__(
        self,
        d_model: int = 128,
        d_expert: int = 64,
        num_shared: int = 1,
        num_routed: int = 8,
        top_k: int = 2,
        capacity_factor: float | None = None
    ):
        super().__init__()
        self.num_routed = num_routed
        self.top_k = top_k
        self.capacity_factor = capacity_factor

        self.shared_experts = nn.ModuleList([SwiGLUExpert(d_model, d_expert * 2) for _ in range(num_shared)])
        self.routed_experts = nn.ModuleList([SwiGLUExpert(d_model, d_expert) for _ in range(num_routed)])
        self.router = AuxLossFreeRouter(d_model, num_experts=num_routed, top_k=top_k)

    def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
        B, S, D = x.shape
        h = x.view(-1, D)  # [T, D] where T = B * S
        T = h.size(0)

        # 1. Shared Experts (always active on all T tokens)
        y_shared = sum(exp(h) for exp in self.shared_experts)

        # 2. Routed Experts with Aux-Loss-Free dynamic bias
        gates, topk_idx, load_ratio = self.router(h)  # gates, topk_idx: [T, K]
        y_routed = torch.zeros_like(h)

        # Optional capacity cap C = ceil(CF * T * K / E)
        max_cap = math.ceil(self.capacity_factor * T * self.top_k / self.num_routed) if self.capacity_factor else T

        for e_idx, expert in enumerate(self.routed_experts):
            token_indices, k_positions = torch.where(topk_idx == e_idx)
            if token_indices.numel() == 0:
                continue
            if token_indices.numel() > max_cap:
                token_indices = token_indices[:max_cap]
                k_positions = k_positions[:max_cap]

            expert_in = h[token_indices]
            expert_out = expert(expert_in)
            weights = gates[token_indices, k_positions].unsqueeze(-1)
            y_routed.index_add_(0, token_indices, expert_out * weights)

        out = (h + y_shared + y_routed).view(B, S, D)
        return out, load_ratio

import torch

def verify_moe_stack():
    torch.manual_seed(42)
    T, d_model, E, K = 512, 64, 8, 2

    # 1. Skew router centroids so Expert 0 & 1 dominate initially
    router = AuxLossFreeRouter(d_model=d_model, num_experts=E, top_k=K, gamma_bias=0.02).train()
    with torch.no_grad():
        router.centroids.weight.data[0] += 0.75
        router.centroids.weight.data[1] += 0.50

    h = torch.randn(T, d_model)
    _, _, initial_load = router(h, update_bias=False)
    initial_max_vio = initial_load.max().item()

    # Run 60 steps of Aux-Loss-Free dynamic bias updates (zero weight gradients!)
    for _ in range(60):
        _, _, final_load = router(torch.randn(T, d_model), update_bias=True)
    final_max_vio = final_load.max().item()
    assert final_max_vio < 1.30 and final_max_vio < initial_max_vio, "Bias routing failed to balance load!"

    # 2. Verify Classical Router Aux + z-Loss vs Full DeepSeekStyleMoELayer backward pass
    classical = ClassicalTopKRouter(d_model=d_model, num_experts=E, top_k=K)
    gates_c, idx_c, reg_loss, _ = classical(h)
    assert gates_c.shape == (T, K) and reg_loss.item() > 0.0

    moe_layer = DeepSeekStyleMoELayer(d_model=d_model, d_expert=32, num_shared=1, num_routed=E, top_k=K, capacity_factor=1.25)
    x_batch = torch.randn(4, 32, d_model, requires_grad=True)
    out, load_ratio = moe_layer(x_batch)
    out.sum().backward()
    assert out.shape == x_batch.shape and x_batch.grad is not None

    print(f"[PASS] Skewed Initial Max Load Ratio   : {initial_max_vio:.2f}x mean load (Imbalanced)")
    print(f"[PASS] Post-Bias Update Max Load Ratio : {final_max_vio:.2f}x mean load (Balanced with ZERO aux grad!)")
    print(f"[PASS] Full DeepSeek MoE Layer Grad    : Verified shape {tuple(out.shape)}, input grad norm {x_batch.grad.norm().item():.2f}")

if __name__ == "__main__":
    verify_moe_stack()