大模型原理--混合专家模型

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


class Expert(nn.Module):
    """单个专家：前馈网络"""
    def __init__(self, d_model, d_ff):
        super().__init__()
        self.w1 = nn.Linear(d_model, d_ff)
        self.w2 = nn.Linear(d_ff, d_model)

    def forward(self, x):
        return self.w2(F.gelu(self.w1(x)))


class MoE_Layer_Efficient(nn.Module):
    """
    高效 MoE 前馈层（稀疏激活版本）
    Args:
        d_model: 输入/输出维度
        d_ff: 专家内部隐藏维度
        num_experts: 专家总数
        top_k: 每 token 激活的专家数
    """
    def __init__(self, d_model, d_ff, num_experts, top_k=2):
        super().__init__()
        assert top_k <= num_experts, f"top_k ({top_k}) 不能超过 num_experts ({num_experts})"
        self.num_experts = num_experts
        self.top_k = top_k

        # 路由器（无偏置）
        self.router = nn.Linear(d_model, num_experts, bias=False)
        # 专家模块
        self.experts = nn.ModuleList([Expert(d_model, d_ff) for _ in range(num_experts)])

    def forward(self, x, return_aux_loss=False):
        """
        x: [batch_size, seq_len, d_model]
        返回:
            out: 专家加权求和输出 [batch_size, seq_len, d_model]
            aux_loss: 负载均衡辅助损失 (若 return_aux_loss=True，否则仅返回 out)
        """
        batch_size, seq_len, d_model = x.shape
        N = batch_size * seq_len                    # 展平后的 token 数
        flat_x = x.reshape(N, d_model)              # [N, d_model]

        # 1. 路由得分
        router_logits = self.router(x)              # [batch, seq, num_experts]

        flat_logits = router_logits.reshape(N, self.num_experts)  # [N, E]
        # 2. Top‑K 选择
        top_k_logits, top_k_indices = torch.topk(flat_logits, self.top_k, dim=-1)  # [N, K]
        top_k_weights = F.softmax(top_k_logits, dim=-1)                            # [N, K]

        # 3. 构建稀疏权重矩阵（仅存储权重，不计算所有专家！）
        #    形状 [N, num_experts]，大部分位置为 0
        sparse_weights = torch.zeros_like(flat_logits)           # [N, E]
        sparse_weights.scatter_(1, top_k_indices, top_k_weights)  # 将 Top‑K 权重填入对应位置

        # 4. 仅对每个专家计算被其选中的 token
        out = torch.zeros_like(flat_x)                          # [N, d_model]
        for expert_idx in range(self.num_experts):
            # 找出所有使用当前专家的 token（该列权重非零）
            col = sparse_weights[:, expert_idx]                 # [N]
            mask = col != 0                                     # 哪些 token 用了这个专家
            if not mask.any():
                continue
            token_indices = mask.nonzero(as_tuple=True)[0]      # 这些 token 在 flat_x 中的索引

            x_selected = flat_x[token_indices]                  # [M, d_model]
            expert_out = self.experts[expert_idx](x_selected)   # [M, d_model]
            weights = col[token_indices].unsqueeze(-1)          # [M, 1]
            # 将加权结果累加到对应 token 的输出中（index_add 保证同一专家多次出现时累积）
            out.index_add_(0, token_indices, expert_out * weights)

        out = out.reshape(batch_size, seq_len, d_model)         # 恢复形状

        if not return_aux_loss:
            return out

        # ---------- 计算负载均衡辅助损失 ----------
        # 路由概率（对所有 token 做 softmax）
        router_probs = F.softmax(flat_logits, dim=-1)           # [N, E]
        # 每个专家的平均路由概率 P_i
        P_i = router_probs.mean(dim=0)     # [E]

        # 每个专家被选中的 token 数量（权重非零即算选中）
        counts = (sparse_weights != 0).sum(dim=0).float()       # [E]
        # 分派比例 f_i （注意：所有 f_i 的和等于 top_k）
        f_i = counts / N
        # 标准辅助损失（Switch Transformer 形式）
        aux_loss = self.num_experts * (f_i * P_i).sum()
        return out, aux_loss