필사 모드: Mixture of Experts（MoE）架构论文综述与实战扩展 2026

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

class Expert(nn.Module):
    """单个专家 FFN 模块"""
    def __init__(self, d_model: int, d_ff: int, dropout: float = 0.1):
        super().__init__()
        self.w1 = nn.Linear(d_model, d_ff)
        self.w2 = nn.Linear(d_ff, d_model)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.w2(self.dropout(F.gelu(self.w1(x))))

class MoELayer(nn.Module):
    """Mixture of Experts 层"""
    def __init__(
        self,
        d_model: int,
        d_ff: int,
        num_experts: int,
        top_k: int = 2,
        dropout: float = 0.1,
    ):
        super().__init__()
        self.num_experts = num_experts
        self.top_k = top_k
        self.experts = nn.ModuleList([
            Expert(d_model, d_ff, dropout) for _ in range(num_experts)
        ])
        # 门控网络: 输入 -> 专家分数
        self.gate = nn.Linear(d_model, num_experts, bias=False)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # x: (batch, seq_len, d_model)
        batch_size, seq_len, d_model = x.shape
        x_flat = x.view(-1, d_model)  # (B*S, d_model)

        # 计算门控分数
        gate_logits = self.gate(x_flat)  # (B*S, num_experts)
        gate_probs = F.softmax(gate_logits, dim=-1)

        # 选择 Top-K 专家
        top_k_probs, top_k_indices = torch.topk(
            gate_probs, self.top_k, dim=-1
        )  # (B*S, top_k)
        # 对选中的专家权重重新归一化
        top_k_probs = top_k_probs / top_k_probs.sum(dim=-1, keepdim=True)

        # 汇总各专家输出
        output = torch.zeros_like(x_flat)
        for k in range(self.top_k):
            expert_idx = top_k_indices[:, k]  # (B*S,)
            weight = top_k_probs[:, k].unsqueeze(-1)  # (B*S, 1)
            for i in range(self.num_experts):
                mask = (expert_idx == i)
                if mask.any():
                    expert_input = x_flat[mask]
                    expert_output = self.experts[i](expert_input)
                    output[mask] += weight[mask] * expert_output

        return output.view(batch_size, seq_len, d_model)

import torch
import torch.nn as nn
import torch.nn.functional as F
from typing import Tuple

class TopKRouter(nn.Module):
    """带噪声注入的 Top-K 路由"""
    def __init__(
        self,
        d_model: int,
        num_experts: int,
        top_k: int = 2,
        noise_std: float = 1.0,
    ):
        super().__init__()
        self.top_k = top_k
        self.num_experts = num_experts
        self.noise_std = noise_std
        self.gate = nn.Linear(d_model, num_experts, bias=False)

    def forward(
        self, x: torch.Tensor
    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        """
        Args:
            x: (batch * seq_len, d_model)
        Returns:
            top_k_indices: (batch * seq_len, top_k)
            top_k_weights: (batch * seq_len, top_k)
            gate_logits: (batch * seq_len, num_experts)
        """
        gate_logits = self.gate(x)

        # 训练时注入噪声以促进探索
        if self.training:
            noise = torch.randn_like(gate_logits) * self.noise_std
            gate_logits = gate_logits + noise

        gate_probs = F.softmax(gate_logits, dim=-1)
        top_k_probs, top_k_indices = torch.topk(
            gate_probs, self.top_k, dim=-1
        )
        # 重新归一化
        top_k_weights = top_k_probs / top_k_probs.sum(dim=-1, keepdim=True)

        return top_k_indices, top_k_weights, gate_logits

def load_balancing_loss(
    gate_logits: torch.Tensor,
    top_k_indices: torch.Tensor,
    num_experts: int,
    top_k: int = 2,
) -> torch.Tensor:
    """
    Switch Transformer 风格的负载均衡损失计算

    最小化每个专家分配到的 token 比例(f_i)
    与门控概率均值(p_i)的内积。
    在均匀分布时取得最小值。

    Args:
        gate_logits: (num_tokens, num_experts)
        top_k_indices: (num_tokens, top_k)
        num_experts: 专家数
        top_k: 活跃专家数
    Returns:
        标量负载均衡损失
    """
    num_tokens = gate_logits.shape[0]
    gate_probs = F.softmax(gate_logits, dim=-1)

    # f_i: 每个专家分配到的 token 比例
    # 用 one-hot 编码统计分配数
    expert_mask = F.one_hot(
        top_k_indices, num_classes=num_experts
    ).float()  # (num_tokens, top_k, num_experts)
    expert_mask = expert_mask.sum(dim=1)  # (num_tokens, num_experts)
    tokens_per_expert = expert_mask.sum(dim=0)  # (num_experts,)
    f = tokens_per_expert / (num_tokens * top_k)

    # p_i: 每个专家的平均门控概率
    p = gate_probs.mean(dim=0)  # (num_experts,)

    # 均衡损失: num_experts * sum(f_i * p_i)
    balance_loss = num_experts * (f * p).sum()

    return balance_loss

def router_z_loss(gate_logits: torch.Tensor) -> torch.Tensor:
    """
    ST-MoE Router Z-Loss 实现

    最小化路由器 logit 的 log-sum-exp 值的平方均值。
    防止 logit 过度增大，从而确保训练稳定性。

    Args:
        gate_logits: (num_tokens, num_experts)
    Returns:
        标量 z-loss 值
    """
    # log(sum(exp(x))) 的平方
    log_z = torch.logsumexp(gate_logits, dim=-1)  # (num_tokens,)
    z_loss = (log_z ** 2).mean()
    return z_loss


# 在训练循环中的整合
def compute_moe_loss(
    task_loss: torch.Tensor,
    gate_logits: torch.Tensor,
    top_k_indices: torch.Tensor,
    num_experts: int,
    alpha_balance: float = 0.01,
    alpha_z: float = 0.001,
) -> torch.Tensor:
    """
    MoE 总损失 = 任务损失 + 均衡损失 + Z-Loss
    """
    bal_loss = load_balancing_loss(
        gate_logits, top_k_indices, num_experts
    )
    z_loss = router_z_loss(gate_logits)

    total_loss = task_loss + alpha_balance * bal_loss + alpha_z * z_loss
    return total_loss

class AuxLossFreeRouter(nn.Module):
    """DeepSeek-V3 风格的 Auxiliary-Loss-Free 路由器"""
    def __init__(
        self,
        d_model: int,
        num_experts: int,
        top_k: int = 8,
        bias_update_speed: float = 0.001,
    ):
        super().__init__()
        self.num_experts = num_experts
        self.top_k = top_k
        self.gamma = bias_update_speed
        self.gate = nn.Linear(d_model, num_experts, bias=False)
        # 偏置不通过梯度更新，而是基于规则调整
        self.register_buffer(
            "expert_bias", torch.zeros(num_experts)
        )
        self.register_buffer(
            "target_load", torch.ones(num_experts) / num_experts
        )

    def forward(self, x: torch.Tensor):
        gate_logits = self.gate(x)
        # 加上偏置以影响路由决策
        biased_logits = gate_logits + self.expert_bias.unsqueeze(0)
        gate_probs = F.softmax(gate_logits, dim=-1)

        # Top-K 选择使用带偏置的 logit
        _, top_k_indices = torch.topk(
            biased_logits, self.top_k, dim=-1
        )
        # 权重从原始门控概率中提取
        top_k_probs = torch.gather(gate_probs, 1, top_k_indices)
        top_k_weights = top_k_probs / top_k_probs.sum(dim=-1, keepdim=True)

        # 训练时更新偏置（不需要梯度）
        if self.training:
            with torch.no_grad():
                expert_load = F.one_hot(
                    top_k_indices.view(-1), self.num_experts
                ).float().sum(0)
                expert_load = expert_load / expert_load.sum()
                # 过载专家降低偏置，负载不足的专家提高偏置
                self.expert_bias += self.gamma * (
                    self.target_load - expert_load
                )

        return top_k_indices, top_k_weights, gate_logits

from dataclasses import dataclass
from typing import Dict, Optional
import torch

@dataclass
class ExpertCacheConfig:
    """专家缓存配置"""
    gpu_cache_size: int = 4   # GPU 上常驻的专家数
    prefetch_count: int = 2    # 预取的专家数
    device: str = "cuda:0"

class ExpertOffloadManager:
    """
    专家卸载管理器。
    基于 LRU 管理 GPU 缓存，
    并预取下一层需要用到的专家。
    """
    def __init__(
        self,
        experts: Dict[int, torch.nn.Module],
        config: ExpertCacheConfig,
    ):
        self.experts = experts
        self.config = config
        self.gpu_cache: Dict[int, torch.nn.Module] = {}
        self.access_order: list = []

    def get_expert(self, expert_id: int) -> torch.nn.Module:
        """把专家加载到 GPU（LRU 缓存）"""
        if expert_id in self.gpu_cache:
            self.access_order.remove(expert_id)
            self.access_order.append(expert_id)
            return self.gpu_cache[expert_id]

        # GPU 缓存已满时淘汰 LRU 专家
        if len(self.gpu_cache) >= self.config.gpu_cache_size:
            evict_id = self.access_order.pop(0)
            self.gpu_cache[evict_id].cpu()
            del self.gpu_cache[evict_id]

        # 把专家移动到 GPU
        expert = self.experts[expert_id].to(self.config.device)
        self.gpu_cache[expert_id] = expert
        self.access_order.append(expert_id)
        return expert

    def prefetch(self, predicted_expert_ids: list):
        """异步预取下一层将使用的专家"""
        for eid in predicted_expert_ids[:self.config.prefetch_count]:
            if eid not in self.gpu_cache:
                self.get_expert(eid)