手撕一个极简 MoE-FFN

上一篇讲了 MoE 的原理——用多个 Expert 替代单一 FFN，每次只激活部分 Expert。

这一篇我们用 PyTorch 从零实现一个极简的 MoE-FFN，80 行代码看清楚 Router 怎么路由、Expert 怎么计算、Load Balancing Loss 怎么加。

一、标准 FFN 的实现（对比参照）

先回顾标准 Transformer 的 FFN：

import torch
import torch.nn as nn

class FeedForward(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)
        self.relu = nn.ReLU()
    
    def forward(self, x):
        # x: (batch_size, seq_len, d_model)
        return self.w2(self.relu(self.w1(x)))

关键点：

• 输入 x 的每个 Token 都过同一个 w1 和 w2
• 参数利用率 100%（所有参数都参与计算）

MoE 要做的是：把 w1 和 w2 拆成 N 个 Expert，每个 Token 只过其中 K 个。

二、MoE-FFN 的整体结构

MoE-FFN 由三部分组成：

1. Router：给每个 Token 打分，决定它走哪些 Expert
2. N 个 Expert：每个 Expert 是一个独立的 FFN
3. 加权求和：把选中的 Expert 输出加权求和

数据流：

输入 x (batch_size, seq_len, d_model)
  ↓
Router 打分 → (batch_size, seq_len, num_experts)
  ↓
Top-K 选择 → 每个 Token 选 K 个 Expert
  ↓
Expert 前向计算 → 每个 Token 只过选中的 Expert
  ↓
加权求和 → (batch_size, seq_len, d_model)

三、实现 Router：打分 + Top-K 选择

Router 是一个线性层 + Softmax，输出每个 Token 对每个 Expert 的得分。

class Router(nn.Module):
    def __init__(self, d_model, num_experts):
        super().__init__()
        self.linear = nn.Linear(d_model, num_experts)
    
    def forward(self, x):
        # x: (batch_size, seq_len, d_model)
        # 输出: (batch_size, seq_len, num_experts)
        logits = self.linear(x)  # (batch_size, seq_len, num_experts)
        scores = torch.softmax(logits, dim=-1)  # 归一化到 [0, 1]
        return scores

Top-K 选择：

def top_k_routing(scores, k):
    """
    scores: (batch_size, seq_len, num_experts)
    k: 每个 Token 选几个 Expert
    
    返回:
    - top_k_indices: (batch_size, seq_len, k) 选中的 Expert 索引
    - top_k_weights: (batch_size, seq_len, k) 对应的权重
    """
    # 取得分最高的 K 个 Expert
    top_k_weights, top_k_indices = torch.topk(scores, k, dim=-1)
    
    # 重新归一化（让选中的 K 个权重加起来等于 1）
    top_k_weights = top_k_weights / top_k_weights.sum(dim=-1, keepdim=True)
    
    return top_k_indices, top_k_weights

例子：

scores = torch.tensor([[[0.1, 0.5, 0.3, 0.1]]])  # 1 个 Token，4 个 Expert
indices, weights = top_k_routing(scores, k=2)
print(indices)  # [[1, 2]]  # Expert-1 和 Expert-2
print(weights)  # [[0.625, 0.375]]  # 归一化后的权重（0.5/(0.5+0.3), 0.3/(0.5+0.3)）

四、实现 Expert：独立的 FFN

每个 Expert 就是一个标准的 FFN，结构和前面的 FeedForward 一样。

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)
        self.relu = nn.ReLU()
    
    def forward(self, x):
        # x: (batch_size, seq_len, d_model)
        return self.w2(self.relu(self.w1(x)))

关键点：每个 Expert 有自己的 w1 和 w2，参数不共享。

五、组装 MoE-FFN：Router + Expert + 加权求和

class MoE_FFN(nn.Module):
    def __init__(self, d_model, d_ff, num_experts, top_k):
        super().__init__()
        self.num_experts = num_experts
        self.top_k = top_k
        
        # Router
        self.router = Router(d_model, num_experts)
        
        # N 个 Expert
        self.experts = nn.ModuleList([
            Expert(d_model, d_ff) for _ in range(num_experts)
        ])
    
    def forward(self, x):
        # x: (batch_size, seq_len, d_model)
        batch_size, seq_len, d_model = x.shape
        
        # 1. Router 打分
        scores = self.router(x)  # (batch_size, seq_len, num_experts)
        
        # 2. Top-K 选择
        top_k_indices, top_k_weights = top_k_routing(scores, self.top_k)
        # top_k_indices: (batch_size, seq_len, top_k)
        # top_k_weights: (batch_size, seq_len, top_k)
        
        # 3. 初始化输出
        output = torch.zeros_like(x)  # (batch_size, seq_len, d_model)
        
        # 4. 对每个选中的 Expert，计算输出并加权
        for i in range(self.top_k):
            # 当前这一列选中的 Expert 索引
            expert_idx = top_k_indices[:, :, i]  # (batch_size, seq_len)
            
            # 当前这一列的权重
            weight = top_k_weights[:, :, i].unsqueeze(-1)  # (batch_size, seq_len, 1)
            
            # 遍历每个 Expert，找出哪些 Token 选了它
            for expert_id in range(self.num_experts):
                # mask: 哪些 Token 选了这个 Expert
                mask = (expert_idx == expert_id)  # (batch_size, seq_len)
                
                if mask.any():
                    # 提取选了这个 Expert 的 Token
                    selected_x = x[mask]  # (num_selected, d_model)
                    
                    # 过这个 Expert
                    expert_out = self.experts[expert_id](selected_x)  # (num_selected, d_model)
                    
                    # 加权后累加到 output
                    output[mask] += weight[mask] * expert_out
        
        return output

数据流详解：

1. Router 打分：每个 Token 对每个 Expert 都有一个得分
2. Top-K 选择：每个 Token 选得分最高的 K 个 Expert
3. 遍历每个 Top-K 列：第 1 列、第 2 列...第 K 列
4. 遍历每个 Expert：找出哪些 Token 在这一列选了这个 Expert
5. 批量计算：把选了同一个 Expert 的 Token 打包，过这个 Expert
6. 加权累加：把输出乘以权重，累加到 output

六、优化：批量计算提高效率

上面的实现有个问题：两层 for 循环，效率低。

优化方法：把选了同一个 Expert 的 Token 打包成一个 batch，一次过 Expert。

class MoE_FFN_Optimized(nn.Module):
    def __init__(self, d_model, d_ff, num_experts, top_k):
        super().__init__()
        self.num_experts = num_experts
        self.top_k = top_k
        self.router = Router(d_model, num_experts)
        self.experts = nn.ModuleList([
            Expert(d_model, d_ff) for _ in range(num_experts)
        ])
    
    def forward(self, x):
        batch_size, seq_len, d_model = x.shape
        
        # 1. Router 打分
        scores = self.router(x)  # (batch_size, seq_len, num_experts)
        
        # 2. Top-K 选择
        top_k_indices, top_k_weights = top_k_routing(scores, self.top_k)
        
        # 3. 把 x 展平成 (batch_size * seq_len, d_model)
        x_flat = x.view(-1, d_model)  # (batch_size * seq_len, d_model)
        top_k_indices_flat = top_k_indices.view(-1, self.top_k)  # (batch_size * seq_len, top_k)
        top_k_weights_flat = top_k_weights.view(-1, self.top_k)  # (batch_size * seq_len, top_k)
        
        # 4. 初始化输出
        output = torch.zeros_like(x_flat)  # (batch_size * seq_len, d_model)
        
        # 5. 对每个 Expert，找出哪些 Token 选了它
        for expert_id in range(self.num_experts):
            # 找出哪些 Token 在 Top-K 里选了这个 Expert
            mask = (top_k_indices_flat == expert_id).any(dim=-1)  # (batch_size * seq_len,)
            
            if not mask.any():
                continue  # 这个 Expert 没被任何 Token 选中，跳过
            
            # 提取选了这个 Expert 的 Token
            selected_x = x_flat[mask]  # (num_selected, d_model)
            
            # 过这个 Expert
            expert_out = self.experts[expert_id](selected_x)  # (num_selected, d_model)
            
            # 找出这些 Token 对这个 Expert 的权重
            # mask.nonzero(as_tuple=False).squeeze(-1) 得到 Token 的索引
            token_indices = mask.nonzero(as_tuple=False).squeeze(-1)
            
            for i, token_idx in enumerate(token_indices):
                # 找出这个 Token 在 Top-K 里的哪一列选了这个 Expert
                col = (top_k_indices_flat[token_idx] == expert_id).nonzero(as_tuple=False)[0].item()
                weight = top_k_weights_flat[token_idx, col]
                output[token_idx] += weight * expert_out[i]
        
        # 6. 恢复形状
        output = output.view(batch_size, seq_len, d_model)
        return output, scores  # 返回 scores 用于计算 Load Balancing Loss

优化点：

• 只遍历 num_experts 次（不是 top_k × num_experts）
• 每个 Expert 只前向一次，处理选了它的所有 Token

七、Load Balancing Loss：防止 Expert 崩溃

如果不加限制，Router 可能只选一个 Expert，其他 Expert 闲置。

Load Balancing Loss 惩罚 Expert 使用不均衡：

$$L_{\text{balance}} = \alpha \cdot N \cdot \sum_{i=1}^N f_i \cdot P_i$$

• $f_i$：Expert-i 被选中的次数占比
• $P_i$：所有 Token 对 Expert-i 的平均得分
• $\alpha$：系数（通常取 0.01）

def load_balancing_loss(scores, top_k_indices, num_experts, alpha=0.01):
    """
    scores: (batch_size, seq_len, num_experts) Router 的输出
    top_k_indices: (batch_size, seq_len, top_k) Top-K 选中的 Expert 索引
    num_experts: Expert 总数
    alpha: 系数
    
    返回: Load Balancing Loss
    """
    batch_size, seq_len, _ = scores.shape
    total_tokens = batch_size * seq_len
    
    # 1. 计算每个 Expert 被选中的次数占比 f_i
    expert_counts = torch.zeros(num_experts, device=scores.device)
    for i in range(num_experts):
        expert_counts[i] = (top_k_indices == i).sum().float()
    f = expert_counts / (total_tokens * top_k_indices.size(-1))  # 归一化
    
    # 2. 计算每个 Expert 的平均得分 P_i
    P = scores.mean(dim=(0, 1))  # (num_experts,)
    
    # 3. 计算 Load Balancing Loss
    loss = alpha * num_experts * (f * P).sum()
    return loss

训练时的总 Loss：

# 前向计算
output, scores = moe_ffn(x)

# 主任务 Loss（如语言模型的 Cross-Entropy）
main_loss = criterion(output, target)

# Load Balancing Loss
balance_loss = load_balancing_loss(scores, top_k_indices, num_experts)

# 总 Loss
total_loss = main_loss + balance_loss

八、完整代码汇总

import torch
import torch.nn as nn

class Router(nn.Module):
    def __init__(self, d_model, num_experts):
        super().__init__()
        self.linear = nn.Linear(d_model, num_experts)
    
    def forward(self, x):
        logits = self.linear(x)
        return torch.softmax(logits, dim=-1)

def top_k_routing(scores, k):
    top_k_weights, top_k_indices = torch.topk(scores, k, dim=-1)
    top_k_weights = top_k_weights / top_k_weights.sum(dim=-1, keepdim=True)
    return top_k_indices, top_k_weights

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)
        self.relu = nn.ReLU()
    
    def forward(self, x):
        return self.w2(self.relu(self.w1(x)))

class MoE_FFN(nn.Module):
    def __init__(self, d_model, d_ff, num_experts, top_k):
        super().__init__()
        self.num_experts = num_experts
        self.top_k = top_k
        self.router = Router(d_model, num_experts)
        self.experts = nn.ModuleList([
            Expert(d_model, d_ff) for _ in range(num_experts)
        ])
    
    def forward(self, x):
        batch_size, seq_len, d_model = x.shape
        scores = self.router(x)
        top_k_indices, top_k_weights = top_k_routing(scores, self.top_k)
        
        output = torch.zeros_like(x)
        for expert_id in range(self.num_experts):
            mask = (top_k_indices == expert_id).any(dim=-1)
            if not mask.any():
                continue
            selected_x = x[mask]
            expert_out = self.experts[expert_id](selected_x)
            token_indices = mask.nonzero(as_tuple=False).squeeze(-1)
            for i, token_idx in enumerate(token_indices):
                cols = (top_k_indices[token_idx // seq_len, token_idx % seq_len] == expert_id).nonzero(as_tuple=False).squeeze(-1)
                for col in cols:
                    weight = top_k_weights[token_idx // seq_len, token_idx % seq_len, col]
                    output[token_idx // seq_len, token_idx % seq_len] += weight * expert_out[i]
        
        return output, scores

def load_balancing_loss(scores, top_k_indices, num_experts, alpha=0.01):
    batch_size, seq_len, _ = scores.shape
    total_tokens = batch_size * seq_len
    expert_counts = torch.zeros(num_experts, device=scores.device)
    for i in range(num_experts):
        expert_counts[i] = (top_k_indices == i).sum().float()
    f = expert_counts / (total_tokens * top_k_indices.size(-1))
    P = scores.mean(dim=(0, 1))
    loss = alpha * num_experts * (f * P).sum()
    return loss

# 测试
if __name__ == "__main__":
    batch_size, seq_len, d_model, d_ff = 2, 4, 512, 2048
    num_experts, top_k = 8, 2
    
    x = torch.randn(batch_size, seq_len, d_model)
    moe = MoE_FFN(d_model, d_ff, num_experts, top_k)
    output, scores = moe(x)
    
    print(f"输入形状: {x.shape}")
    print(f"输出形状: {output.shape}")
    print(f"Router 得分形状: {scores.shape}")
    
    # 计算 Load Balancing Loss
    top_k_indices, _ = top_k_routing(scores, top_k)
    balance_loss = load_balancing_loss(scores, top_k_indices, num_experts)
    print(f"Load Balancing Loss: {balance_loss.item():.4f}")

九、测试：看看 Router 学到了什么

训练一个玩具任务，看看 Router 是否学会了路由。

任务：输入 8 个随机向量，前 4 个属于类别 A（目标输出全 1），后 4 个属于类别 B（目标输出全 0）。

# 生成玩具数据
x = torch.randn(1, 8, 512)
target = torch.cat([torch.ones(1, 4, 512), torch.zeros(1, 4, 512)], dim=1)

# 训练
moe = MoE_FFN(512, 2048, num_experts=4, top_k=1)
optimizer = torch.optim.Adam(moe.parameters(), lr=0.001)

for epoch in range(100):
    output, scores = moe(x)
    main_loss = ((output - target) ** 2).mean()
    top_k_indices, _ = top_k_routing(scores, 1)
    balance_loss = load_balancing_loss(scores, top_k_indices, 4)
    total_loss = main_loss + balance_loss
    
    optimizer.zero_grad()
    total_loss.backward()
    optimizer.step()
    
    if epoch % 20 == 0:
        print(f"Epoch {epoch}, Loss: {total_loss.item():.4f}")

# 查看 Router 的路由结果
with torch.no_grad():
    output, scores = moe(x)
    top_k_indices, _ = top_k_routing(scores, 1)
    print("Router 的路由结果（每个 Token 选了哪个 Expert）:")
    print(top_k_indices.squeeze())

预期结果：

• 前 4 个 Token 选 Expert-0 或 Expert-1
• 后 4 个 Token 选 Expert-2 或 Expert-3
• Router 自动学会了把不同类别的 Token 路由到不同的 Expert

十、总结

MoE-FFN 的核心组件：

1. Router：线性层 + Softmax，给每个 Token 打分
2. Top-K 选择：每个 Token 选得分最高的 K 个 Expert
3. Expert：独立的 FFN，参数不共享
4. 加权求和：把选中的 Expert 输出加权求和
5. Load Balancing Loss：惩罚 Expert 使用不均衡

关键代码：

• Router(x) → scores
• top_k_routing(scores, k) → indices, weights
• experts[i](selected_x) → expert_out
• output += weight * expert_out

MoE 的核心思想：

• 每个 Token 走不同的路径，Router 自动学会路由
• 参数多≠计算多，激活参数只有 K/N
• 需要 Load Balancing Loss，否则 Router 只选一个 Expert

下一章我们进入 Llama 架构——从标准 Transformer 进化到现代大模型，学会读技术报告。

---

相关文章：

• MoE 原理（理解稀疏激活和 Load Balancing）
• FFN 和 LayerNorm（理解标准 FFN）
• 手撕 Transformer FFN（对比标准 FFN 实现）