Sparse Mixture of Experts（MoE）架构深度解析：从设计原理到 DeepSeek-V3、Qwen3

y = sum_{i=1}^{N} G(x)_i * E_i(x)

G(x)_i = softmax(W_g * x + noise)_i  (if i in TopK)
G(x)_i = 0                           (otherwise)

import torch
import torch.nn as nn
import torch.nn.functional as F
from dataclasses import dataclass

@dataclass
class MoEConfig:
    d_model: int = 1024
    d_ff: int = 4096
    num_experts: int = 8
    top_k: int = 2
    dropout: float = 0.1
    aux_loss_weight: float = 0.01

class SwiGLUExpert(nn.Module):
    """使用 SwiGLU 激活函数的专家 FFN。
    LLaMA、Mistral 等最新模型采用的标准结构。"""

    def __init__(self, d_model: int, d_ff: int, dropout: float = 0.1):
        super().__init__()
        self.w_gate = nn.Linear(d_model, d_ff, bias=False)
        self.w_up = nn.Linear(d_model, d_ff, bias=False)
        self.w_down = nn.Linear(d_ff, d_model, bias=False)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        gate = F.silu(self.w_gate(x))
        up = self.w_up(x)
        return self.w_down(self.dropout(gate * up))

class TopKGating(nn.Module):
    """Top-K 门控网络。
    实现 Noisy Top-K Gating（Shazeer et al., 2017）。"""

    def __init__(self, d_model: int, num_experts: int, top_k: int = 2):
        super().__init__()
        self.num_experts = num_experts
        self.top_k = top_k
        self.gate = nn.Linear(d_model, num_experts, bias=False)
        self.noise_linear = nn.Linear(d_model, num_experts, bias=False)

    def forward(self, x: torch.Tensor):
        # x: (batch * seq_len, d_model)
        logits = self.gate(x)

        if self.training:
            noise = F.softplus(self.noise_linear(x))
            logits = logits + noise * torch.randn_like(logits)

        probs = F.softmax(logits, dim=-1)
        top_k_probs, top_k_indices = torch.topk(probs, self.top_k, dim=-1)
        # 重新归一化：使被选中专家的权重之和为 1
        top_k_probs = top_k_probs / top_k_probs.sum(dim=-1, keepdim=True)

        return top_k_probs, top_k_indices, probs

class ExpertChoiceGating(nn.Module):
    """实现 Expert Choice Routing。
    每个专家直接选择自己要处理的 token，
    从结构上保证负载均衡。"""

    def __init__(
        self,
        d_model: int,
        num_experts: int,
        capacity_factor: float = 1.0,
    ):
        super().__init__()
        self.num_experts = num_experts
        self.capacity_factor = capacity_factor
        self.gate = nn.Linear(d_model, num_experts, bias=False)

    def forward(self, x: torch.Tensor):
        # x: (num_tokens, d_model)
        num_tokens = x.shape[0]
        expert_capacity = int(
            num_tokens * self.capacity_factor / self.num_experts
        )

        # 门控分数：(num_tokens, num_experts)
        gate_logits = self.gate(x)
        # 从专家视角计算分数：(num_experts, num_tokens)
        gate_scores = F.softmax(gate_logits.T, dim=-1)

        # 每个专家选择排名前 capacity 的 token
        top_k_scores, top_k_indices = torch.topk(
            gate_scores, expert_capacity, dim=-1
        )  # (num_experts, capacity)

        # 生成调度掩码
        dispatch_mask = torch.zeros(
            self.num_experts, num_tokens,
            device=x.device, dtype=x.dtype,
        )
        dispatch_mask.scatter_(1, top_k_indices, top_k_scores)

        return dispatch_mask, top_k_indices, top_k_scores

def compute_load_balancing_loss(
    gate_probs: torch.Tensor,
    top_k_indices: torch.Tensor,
    num_experts: int,
    top_k: int,
) -> torch.Tensor:
    """计算 Switch Transformer 风格的负载均衡辅助损失。

    Args:
        gate_probs: 门控概率 (num_tokens, num_experts)
        top_k_indices: 被选中的专家索引 (num_tokens, top_k)
        num_experts: 专家数量
        top_k: 被选中的专家数量

    Returns:
        辅助损失的标量值
    """
    num_tokens = gate_probs.shape[0]

    # f_i：分配给专家 i 的 token 比例
    expert_mask = F.one_hot(top_k_indices, num_experts).float()
    # (num_tokens, top_k, num_experts) -> (num_tokens, num_experts)
    expert_mask = expert_mask.sum(dim=1)
    tokens_per_expert = expert_mask.sum(dim=0)  # (num_experts,)
    f = tokens_per_expert / (num_tokens * top_k)

    # P_i：专家 i 的平均门控概率
    P = gate_probs.mean(dim=0)  # (num_experts,)

    # 辅助损失：N * sum(f_i * P_i)
    # 均匀分配时取得最小值
    aux_loss = num_experts * (f * P).sum()

    return aux_loss

class MoELayerWithAuxLoss(nn.Module):
    """包含辅助损失的完整 MoE 层实现。"""

    def __init__(self, config: MoEConfig):
        super().__init__()
        self.config = config
        self.experts = nn.ModuleList([
            SwiGLUExpert(config.d_model, config.d_ff, config.dropout)
            for _ in range(config.num_experts)
        ])
        self.gating = TopKGating(
            config.d_model, config.num_experts, config.top_k
        )
        self.aux_loss_weight = config.aux_loss_weight

    def forward(self, x: torch.Tensor):
        batch_size, seq_len, d_model = x.shape
        x_flat = x.view(-1, d_model)

        top_k_probs, top_k_indices, gate_probs = self.gating(x_flat)

        # 计算辅助损失
        aux_loss = self.aux_loss_weight * compute_load_balancing_loss(
            gate_probs, top_k_indices,
            self.config.num_experts, self.config.top_k,
        )

        # 计算专家输出
        output = torch.zeros_like(x_flat)
        for k in range(self.config.top_k):
            expert_indices = top_k_indices[:, k]  # (num_tokens,)
            expert_weights = top_k_probs[:, k]    # (num_tokens,)

            for i in range(self.config.num_experts):
                mask = (expert_indices == i)
                if mask.any():
                    expert_input = x_flat[mask]
                    expert_output = self.experts[i](expert_input)
                    output[mask] += expert_weights[mask].unsqueeze(-1) * expert_output

        output = output.view(batch_size, seq_len, d_model)
        return output, aux_loss

import logging
from collections import defaultdict

logger = logging.getLogger(__name__)

class ExpertUtilizationMonitor:
    """专家利用率监控与崩溃检测工具。

    在训练过程中追踪每个专家的利用率，
    并及早发现崩溃迹象。
    """

    def __init__(
        self,
        num_experts: int,
        collapse_threshold: float = 0.01,
        window_size: int = 100,
    ):
        self.num_experts = num_experts
        self.collapse_threshold = collapse_threshold
        self.window_size = window_size
        self.history: list[dict[int, float]] = []

    def record(self, expert_counts: dict[int, int], total_tokens: int):
        """记录每个批次的专家利用率。"""
        utilization = {
            i: expert_counts.get(i, 0) / max(total_tokens, 1)
            for i in range(self.num_experts)
        }
        self.history.append(utilization)

        if len(self.history) > self.window_size:
            self.history = self.history[-self.window_size:]

    def detect_collapse(self) -> list[int]:
        """检测专家崩溃。返回利用率低于阈值（threshold）的专家。"""
        if len(self.history) < self.window_size // 2:
            return []

        collapsed = []
        for expert_id in range(self.num_experts):
            recent_util = [
                h[expert_id] for h in self.history[-self.window_size:]
            ]
            avg_util = sum(recent_util) / len(recent_util)
            if avg_util < self.collapse_threshold:
                collapsed.append(expert_id)

        if collapsed:
            logger.warning(
                f"Expert collapse detected! "
                f"Experts {collapsed} have utilization below "
                f"{self.collapse_threshold:.2%}. "
                f"Consider increasing aux_loss_weight or "
                f"reinitializing collapsed experts."
            )

        return collapsed

    def get_load_imbalance_ratio(self) -> float:
        """计算负载不均衡比率。
        1.0 表示完美均衡，数值越大越不均衡。"""
        if not self.history:
            return 0.0

        latest = self.history[-1]
        utils = list(latest.values())
        max_util = max(utils) if utils else 0
        min_util = min(utils) if utils else 0
        avg_util = sum(utils) / len(utils) if utils else 0

        if avg_util == 0:
            return float("inf")

        return max_util / avg_util

import torch
import torch.distributed as dist
from typing import Optional

class ExpertParallelRouter:
    """实现 Expert Parallelism 所需的 token 分发/收集。

    每个 GPU 负责一部分专家，
    通过 All-to-All 通信路由 token。
    """

    def __init__(
        self,
        num_experts: int,
        ep_group: Optional[dist.ProcessGroup] = None,
    ):
        self.num_experts = num_experts
        self.ep_group = ep_group
        self.ep_size = dist.get_world_size(ep_group) if ep_group else 1
        self.ep_rank = dist.get_rank(ep_group) if ep_group else 0
        self.experts_per_rank = num_experts // self.ep_size

    def dispatch(
        self,
        tokens: torch.Tensor,
        expert_indices: torch.Tensor,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        """把 token 分发到负责的 GPU。

        Args:
            tokens: (num_tokens, d_model) 输入 token
            expert_indices: (num_tokens,) 每个 token 的目标专家索引

        Returns:
            dispatched_tokens: 本 GPU 需要处理的 token
            recv_counts: 从每个 GPU 收到的 token 数
        """
        # 计算发往每个 GPU 的 token 数
        send_counts = torch.zeros(
            self.ep_size, dtype=torch.long, device=tokens.device
        )
        for rank in range(self.ep_size):
            start_expert = rank * self.experts_per_rank
            end_expert = start_expert + self.experts_per_rank
            mask = (expert_indices >= start_expert) & (
                expert_indices < end_expert
            )
            send_counts[rank] = mask.sum()

        # 通过 All-to-All 交换接收计数
        recv_counts = torch.zeros_like(send_counts)
        dist.all_to_all_single(
            recv_counts, send_counts, group=self.ep_group
        )

        # token 排序及 All-to-All 发送
        sorted_indices = torch.argsort(expert_indices)
        sorted_tokens = tokens[sorted_indices]

        send_splits = send_counts.tolist()
        recv_splits = recv_counts.tolist()

        dispatched_tokens = torch.zeros(
            int(recv_counts.sum()), tokens.shape[1],
            dtype=tokens.dtype, device=tokens.device,
        )
        dist.all_to_all_single(
            dispatched_tokens, sorted_tokens,
            output_split_sizes=recv_splits,
            input_split_sizes=send_splits,
            group=self.ep_group,
        )

        return dispatched_tokens, recv_counts

def router_z_loss(gate_logits: torch.Tensor) -> torch.Tensor:
    """ST-MoE 风格的 Router z-loss。
    通过约束门控 logit 的大小来提高数值稳定性。

    Args:
        gate_logits: (num_tokens, num_experts) 门控 logit

    Returns:
        z_loss 标量
    """
    log_z = torch.logsumexp(gate_logits, dim=-1)  # (num_tokens,)
    z_loss = (log_z ** 2).mean()
    return z_loss