混合专家模型（MoE）

到目前为止的 GPT，每一层只有一个 FFN，所有 token 共享。参数和计算是绑在一起的——参数翻倍，每个 token 的计算也翻倍。

这一节介绍 MoE。它把一个大 FFN 拆成多个专家，每个 token 只激活少数几个。总参数可以很大，计算量不怎么变。我们从零写一个 MoE 层，重点看路由器怎么选专家、负载怎么均衡。

MoE 的思路：一个 FFN 拆成多个小的专家 FFN，前面加一个路由器决定每个 token 走哪几个。比如 8 个专家，每次只走 2 个——总参数 8 倍，每个 token 的计算 2 倍。参数多，能装的知识多；计算变化不大，推理速度不受影响。

训一段时间后，路由通常会出现某种结构化分工：Mixtral 论文观察到相邻 token 往往分到相同专家，也有句法相关的路由模式；DeepSeekMoE 则用“细粒度专家 + 共享专家”推动更专门化的分工。但要注意，这不等于每个专家都能被简单命名成“语法专家”或“数字专家”。一些解释性研究提醒，专家更可能学到细粒度语言操作、局部语义模式，或者跨领域都常用的核心能力。参考：Mixtral、DeepSeekMoE、MoE Interpretability、Core Experts。

MoE 还有一个麻烦：路由器可能偷懒，把大部分 token 都发给少数几个专家，其余专家闲着。怎么让负载均衡，是 MoE 训练最核心的问题。

1. 普通 Transformer 的 FFN 层

Dense 模型和 MoE 模型的根本区别在于参数和计算量的关系：

Dense 模型:
  所有 token → 同一个 FFN → 输出
  参数量 = 每次推理的计算量（全部参数都参与计算）
  想让模型更强 → 扩大 FFN → 计算量同步增长

MoE 模型:
  每个 token → 路由器选 top-k 个专家 → 加权输出
  总参数量 = N × 单个专家参数量（可以很大）
  每次计算量 = k × 单个专家参数量（和 Dense 差不多）
  参数和计算量不再同步增长

Dense 模型中，想让模型拥有更大的知识容量，就要把 FFN 的维度成倍扩大。但 FFN 扩大后，每个 token 都要经过这个更大的 FFN，推理成本同步增长。MoE 的思路是把知识分散到多个专家里，每个 token 只激活少数几个——总参数可以很大，但每次推理的计算量基本不变。

下面先回顾标准 FFN 的结构，再在这个基础上改造成 MoE。

每个 Transformer Block 里有一个 FFN（前馈网络），结构很简单——两个线性变换，中间夹一个激活函数：

输入 x (d_model=512)
  ↓
Linear(512 → 2048)  ← 升维，给模型更大的「思考空间」
  ↓
ReLU / GELU          ← 非线性，让模型能学复杂模式
  ↓
Linear(2048 → 512)  ← 降维，回到原来的维度
  ↓
输出

参数量：512 × 2048 + 2048 × 512 ≈ 2M（两个矩阵，各约 1M）。

这个 FFN 是 Transformer 能力的核心来源之一。Attention 负责「从哪些位置获取信息」，FFN 负责「对获取的信息做什么处理」。Attention 告诉你「这个词和哪个词相关」，FFN 告诉你「知道了相关性之后，该把这个词变成什么」。

但现在的设计有一个隐含假设：所有 token 共享同一个 FFN。不管输入是「的」还是「量子力学」，都经过同两个矩阵，做同样的变换。这在小模型里没问题，但模型变大后，用一个 FFN 同时处理所有类型的知识就越来越难——就像用同一套规则处理语法问题和数学问题，规则本身会变得臃肿。

1.1 为什么 MoE 通常替换 FFN，而不是替换 Attention

你可能会问：既然 Transformer Block 里有 Attention 和 FFN 两个大部件，为什么 MoE 通常拿 FFN 开刀，而不是把 Attention 拆成专家？

原因有三个。

第一，FFN 是逐 token 计算的。每个 token 过 FFN 时，本来就像一个独立样本：输入一个向量，输出一个向量。路由器很容易对每个 token 单独决定「找哪几个专家」。

第二，FFN 通常参数很多。Dense FFN 的两层或三层大矩阵占了 Block 里相当多的参数。把它拆成多个专家，能显著增加模型容量；而每个 token 只走 top-k 个专家，计算不会按总专家数线性爆炸。

第三，Attention 负责 token 之间的信息交换。如果把 Attention 也动态路由，系统会更复杂：不同 token 不只要选专家，还要互相看，通信和缓存都会变麻烦。MoE 先替换 FFN，是收益大、改动相对清楚的一步。

所以你可以把 MoE 理解成：保留 Attention 这条信息交换通道，把后面的 FFN 加工车间扩建成多个专家车间。

2. MoE 的核心思想

既然一个 FFN 处理所有 token 负担太重，那换一个思路：把一个大 FFN 拆成 N 个小的专家 FFN，让每个 token 只找其中少数几个专家来处理。

               ┌─────────────┐
               │   路由器     │  ← 决定每个 token 找哪几个专家
               │  (Gate)     │
               └──┬──┬──┬───┘
                  │  │  │
          ┌───────┘  │  └───────┐
          ↓          ↓          ↓
     ┌────────┐┌────────┐┌────────┐
     │ 专家 1 ││ 专家 2 ││ 专家 3 │  ... (共 8 个)
     │ (FFN)  ││ (FFN)  ││ (FFN)  │
     └────────┘└────────┘└────────┘
          ↓          ↓          ↓
          └──────────┴──────────┘
                    加权求和

路由器的输入是 token 的 hidden state（一个 d_model 维的向量），输出是 N 个分数（每个专家一个分）。分数高的专家被选中。

每个 token 只激活 top-k 个专家（通常 k=2），不是全部 8 个：

token "的"     → 路由器 → 专家 1, 5  (功能词，通用专家处理)
token "量子"   → 路由器 → 专家 3, 7  (物理类知识)
token "hello"  → 路由器 → 专家 2, 6  (英文类知识)

效果：

总参数量 = N × 一个专家的参数量（8 倍增长）
每次推理的计算量 = k × 一个专家的参数量（2 倍增长）
参数多但计算少 ← 参数和计算不再同步增长

路由器本身也是一个可训练的参数矩阵——nn.Linear(d_model, num_experts)。它和专家 FFN 一起训练，通过梯度学习如何为每个 token 选择最合适的专家。

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

class MoELayer(nn.Module):
    """
    MoE FFN 层

    参数:
        d_model: 隐藏维度
        num_experts: 专家数量
        top_k: 每个 token 激活几个专家
    """
    def __init__(self, d_model, num_experts=8, top_k=2, expert_dim=None):
        super().__init__()
        self.num_experts = num_experts
        self.top_k = top_k
        expert_dim = expert_dim or 4 * d_model

        # 路由器：输入 d_model，输出 num_experts 个分数
        self.gate = nn.Linear(d_model, num_experts, bias=False)

        # N 个专家，每个是一个小 FFN
        self.experts = nn.ModuleList([
            nn.Sequential(
                nn.Linear(d_model, expert_dim),
                nn.ReLU(),
                nn.Linear(expert_dim, d_model)
            )
            for _ in range(num_experts)
        ])

    def forward(self, x):
        """
        x: [batch, seq_len, d_model]

        流程:
        1. 路由器给每个 token 打分
        2. 选 top-k 个专家
        3. 只算这 k 个专家的输出
        4. 加权求和
        """
        batch_size, seq_len, d_model = x.shape

        # Step 1: 路由器打分
        gate_logits = self.gate(x)  # [batch, seq_len, num_experts]

        # Step 2: 选 top-k
        top_k_logits, top_k_indices = torch.topk(gate_logits, self.top_k, dim=-1)
        top_k_weights = F.softmax(top_k_logits, dim=-1)  # 归一化权重

        # Step 3 & 4: 对每个 token，算选中专家的输出并加权求和
        output = torch.zeros_like(x)

        for b in range(batch_size):
            for s in range(seq_len):
                token = x[b, s]  # [d_model]
                for k in range(self.top_k):
                    expert_idx = top_k_indices[b, s, k].item()
                    weight = top_k_weights[b, s, k]
                    expert_out = self.experts[expert_idx](token.unsqueeze(0)).squeeze(0)
                    output[b, s] += weight * expert_out

        return output

print("MoE 层定义完成！")
print(f"8 个专家，每个 token 只激活 2 个")

# 演示 MoE 的路由过程
import torch
import torch.nn.functional as F

torch.manual_seed(42)

moe = MoELayer(d_model=16, num_experts=8, top_k=2)

# 模拟 3 个 token
x = torch.randn(1, 3, 16)

# 看路由器给每个 token 的评分
with torch.no_grad():
    gate_scores = moe.gate(x).squeeze(0)  # [3, 8]
    top_k_vals, top_k_idx = torch.topk(gate_scores, 2, dim=-1)

print("=== 路由器为 3 个 token 选择的专家 ===")
print()
for i in range(3):
    experts = top_k_idx[i].tolist()
    weights = F.softmax(top_k_vals[i], dim=-1).tolist()
    print(f"Token {i}: 选中专家 {experts}, 权重 {[f'{w:.2f}' for w in weights]}")

print()
print("每个 token 只激活 2/8 = 25% 的专家")
print("总参数是 8 个专家的和，但计算量只有 2 个专家的量")

2.1 对比：普通 Transformer Block vs MoE Block

要看清 MoE 的结构，最好的办法不是只看公式，而是把模型结构直接打印出来。

普通 decoder block 的主线是：

x → Attention → FFN → output

MoE decoder block 的主线是：

x → Attention → Router → 多个 FFN experts 里选 top-k → output

也就是说，MoE 通常不是替换 Attention，而是把 Transformer Block 里的 FFN 换成“路由器 + 多个专家 FFN”。

# 用真实 nn.Module 打印：Dense FFN Block vs MoE Block
import torch.nn as nn

class TinyDenseBlock(nn.Module):
    """普通 Transformer decoder block 的骨架：Attention + 单个 FFN"""
    def __init__(self, d_model=16, num_heads=2, ffn_dim=64):
        super().__init__()
        self.self_attn = nn.MultiheadAttention(d_model, num_heads, batch_first=True)
        self.norm1 = nn.LayerNorm(d_model)
        self.ffn = nn.Sequential(
            nn.Linear(d_model, ffn_dim),
            nn.ReLU(),
            nn.Linear(ffn_dim, d_model),
        )
        self.norm2 = nn.LayerNorm(d_model)

    def forward(self, x):
        attn_out, _ = self.self_attn(x, x, x, need_weights=False)
        x = self.norm1(x + attn_out)
        ffn_out = self.ffn(x)
        x = self.norm2(x + ffn_out)
        return x

class TinyMoEBlock(nn.Module):
    """MoE decoder block 的骨架：Attention + Router + 多个 FFN experts"""
    def __init__(self, d_model=16, num_heads=2, num_experts=4, top_k=2, expert_dim=64):
        super().__init__()
        self.self_attn = nn.MultiheadAttention(d_model, num_heads, batch_first=True)
        self.norm1 = nn.LayerNorm(d_model)
        self.moe = MoELayer(d_model, num_experts=num_experts, top_k=top_k, expert_dim=expert_dim)
        self.norm2 = nn.LayerNorm(d_model)

    def forward(self, x):
        attn_out, _ = self.self_attn(x, x, x, need_weights=False)
        x = self.norm1(x + attn_out)
        moe_out = self.moe(x)
        x = self.norm2(x + moe_out)
        return x

dense_block = TinyDenseBlock()
moe_block = TinyMoEBlock(num_experts=4, top_k=2)

print("=== 普通 Transformer Block ===")
print(dense_block)
print()
print("=== MoE Transformer Block ===")
print(moe_block)

上面的打印结果要看两个位置：

普通 block 里只有一个 ffn。
MoE block 里 self_attn 后面接的是 moe，里面有 gate 和多个 experts。

这就是 MoE 的核心结构证据：Attention 还在，FFN 变成了多个专家。

# trace 一次 forward：看 shape 和路由发生在哪里
import torch

trace_x = torch.randn(1, 3, 16)

print("=== Dense Block Trace ===")
with torch.no_grad():
    dense_attn_out, _ = dense_block.self_attn(trace_x, trace_x, trace_x, need_weights=False)
    dense_after_attn = dense_block.norm1(trace_x + dense_attn_out)
    dense_ffn_out = dense_block.ffn(dense_after_attn)
    dense_out = dense_block.norm2(dense_after_attn + dense_ffn_out)

print(f"input:        {tuple(trace_x.shape)}")
print(f"attention:    {tuple(dense_attn_out.shape)}")
print(f"single FFN:   {tuple(dense_ffn_out.shape)}")
print(f"output:       {tuple(dense_out.shape)}")

print()
print("=== MoE Block Trace ===")
with torch.no_grad():
    moe_attn_out, _ = moe_block.self_attn(trace_x, trace_x, trace_x, need_weights=False)
    moe_after_attn = moe_block.norm1(trace_x + moe_attn_out)
    gate_logits = moe_block.moe.gate(moe_after_attn)
    top_vals, top_idx = torch.topk(gate_logits, moe_block.moe.top_k, dim=-1)
    moe_out_raw = moe_block.moe(moe_after_attn)
    moe_out = moe_block.norm2(moe_after_attn + moe_out_raw)

print(f"input:        {tuple(trace_x.shape)}")
print(f"attention:    {tuple(moe_attn_out.shape)}")
print(f"gate logits:  {tuple(gate_logits.shape)}  # 每个 token 对每个 expert 的分数")
print(f"top-k index:  {tuple(top_idx.shape)}      # 每个 token 选中的 expert 编号")
print(f"MoE FFN:      {tuple(moe_out_raw.shape)}")
print(f"output:       {tuple(moe_out.shape)}")
print()
print("每个 token 选中的 experts：")
for token_i, experts in enumerate(top_idx[0].tolist()):
    print(f"token {token_i}: experts {experts}")

2.2 打印 HuggingFace 里的真实 MoE 模型

小模型的骨架看懂之后，再看真实工程里的 decoder layer。这里不下载权重，只用 config 初始化一个很小的 Qwen2-MoE / Mixtral，目的只有一个：把真实源码里的模块顺序打印出来。

读打印结果时盯住三件事：

self_attn 仍然在前面。
普通 FFN 的位置，变成了 mlp / SparseMoeBlock。
MoE 里面有 gate 和多个 experts。

新版 transformers 会把多个 expert 的权重打包成大 tensor，打印时不一定显示成 expert_0, expert_1, ...。所以除了打印 layer，还要打印参数形状： experts.gate_up_proj 的第 0 维就是 expert 数量。

# 打印真实 HuggingFace MoE decoder layer，并 trace 一次 router 输出
import inspect
import warnings

import torch

warnings.filterwarnings("ignore", message="IProgress not found.*")

from transformers import Qwen2MoeConfig, Qwen2MoeForCausalLM
from transformers import MixtralConfig, MixtralForCausalLM

qwen_cfg = Qwen2MoeConfig(
    vocab_size=128,
    hidden_size=32,
    intermediate_size=64,
    moe_intermediate_size=64,
    shared_expert_intermediate_size=64,
    num_hidden_layers=1,
    num_attention_heads=4,
    num_key_value_heads=4,
    num_experts=4,
    num_experts_per_tok=2,
)
qwen_moe = Qwen2MoeForCausalLM(qwen_cfg)
qwen_layer = qwen_moe.model.layers[0]

print("=== Qwen2-MoE decoder layer ===")
print(qwen_layer)

print()
print("=== Qwen2-MoE MoE 参数形状 ===")
for name, param in qwen_layer.mlp.named_parameters():
    if name.startswith("gate") or name.startswith("experts") or name.startswith("shared"):
        print(f"{name:32s} {tuple(param.shape)}")

mixtral_cfg = MixtralConfig(
    vocab_size=128,
    hidden_size=32,
    intermediate_size=64,
    num_hidden_layers=1,
    num_attention_heads=4,
    num_key_value_heads=4,
    num_local_experts=4,
    num_experts_per_tok=2,
)
mixtral_moe = MixtralForCausalLM(mixtral_cfg)
mixtral_layer = mixtral_moe.model.layers[0]

print()
print("=== Mixtral decoder layer ===")
print(mixtral_layer)

print()
print("=== Mixtral MoE 参数形状 ===")
for name, param in mixtral_layer.mlp.named_parameters():
    if name.startswith("gate") or name.startswith("experts"):
        print(f"{name:32s} {tuple(param.shape)}")

print()
print("=== Qwen2-MoE layer.forward 源码摘录 ===")
source = inspect.getsource(type(qwen_layer).forward).splitlines()
for line in source:
    if "self_attn" in line or "mlp" in line or "layernorm" in line:
        print(line)

print()
print("=== Qwen2-MoE mlp.forward 源码摘录 ===")
source = inspect.getsource(type(qwen_layer.mlp).forward).splitlines()
for line in source:
    if "gate" in line or "expert" in line or "router" in line or "top" in line:
        print(line)

input_ids = torch.randint(0, 128, (1, 6))
with torch.no_grad():
    qwen_out = qwen_moe(input_ids=input_ids, output_router_logits=True)

router_logits = qwen_out.router_logits[0]
top_vals, top_idx = torch.topk(router_logits, qwen_cfg.num_experts_per_tok, dim=-1)

print()
print("=== Qwen2-MoE router trace ===")
print(f"input_ids:      {tuple(input_ids.shape)}")
print(f"logits:         {tuple(qwen_out.logits.shape)}")
print(f"router logits:  {tuple(router_logits.shape)}")
print(f"top-k experts:  {tuple(top_idx.shape)}")
print()
print("前 6 个 token 选中的 experts：")
for token_i, experts in enumerate(top_idx[:6].tolist()):
    print(f"token {token_i}: experts {experts}")

3. MoE 的参数 vs 计算量

这是 MoE 最核心的优势。用具体数字感受一下：

假设一个 Dense 模型的 FFN 有 2M 参数（d_model=512，d_ff=2048）：

普通 Dense 模型:
  FFN 参数: 2M
  每个 token 的计算: 2M 次参数运算
  参数和计算绑定 → 参数翻倍，计算也翻倍

MoE 模型 (8 专家, top-2):
  FFN 总参数: 8 × 2M = 16M  ← 参数翻了 8 倍
  每个 token 的计算: 2 × 2M = 4M  ← 计算只翻了 2 倍
  参数和计算解耦

一个 token 在 MoE 中的实际计算路径是这样的：

经过路由器：W_gate @ x，这是一个很小的矩阵乘法（d_model × num_experts）
经过 top-k 个专家的 FFN：每个专家内部是两次矩阵乘法（升维 + 降维）
加权求和：把 k 个专家的输出按路由权重加起来

路由器的计算量远小于 FFN（d_model × num_experts << d_model × d_ff），所以忽略不计。主要计算就是 k 个专家的 FFN 计算。

这就是为什么 Mixtral 8×7B 虽然总参数 47B，但推理速度和 7B 的 Dense 模型差不多——每次推理只激活约 13B 参数（2 个专家的 FFN + 共享的 Attention 层）。

4. MoE 的训练难题：负载均衡

MoE 有一个内生的工程问题：路由器可能偷懒，只把 token 发给少数几个专家。

为什么会这样？路由器的训练目标只有一个——让模型预测下一个 token 的 loss 尽可能低。如果路由器发现「把所有 token 都给专家 3 和专家 5 就能让 loss 很低」，它就没有动力去用其他专家。

坏情况（负载不均）:
  专家 1: ████████████████████  (被过度使用)
  专家 2: ████
  专家 3: █
  专家 4-8: (几乎空闲，参数白训了)

好情况（负载均衡）:
  专家 1: ██████
  专家 2: ██████
  专家 3: ██████
  ...
  专家 8: ██████

负载不均的后果很严重：被过度使用的专家成为瓶颈（计算慢），空闲专家的参数没学到东西（浪费容量），最终模型退化成「只有 2-3 个有效专家的 Dense 模型」。

传统解决方案：辅助 loss

在语言模型的主 loss 之外，加一个额外的 loss 项，鼓励每个专家处理大致相同数量的 token。

# 负载均衡 loss（简化版）
load_balance_loss = 0
for expert_i in range(num_experts):
    actual_load = count_tokens_assigned_to(expert_i)
    ideal_load = total_tokens * top_k / num_experts
    load_balance_loss += (actual_load - ideal_load) ** 2

total_loss = lm_loss + alpha * load_balance_loss

系数 α 需要手动调节。α 太大会让路由器被迫选不合适的专家（影响模型质量），α 太小则负载均衡无效。

改进：无辅助 Loss 的负载均衡（DeepSeek-V3）

辅助 loss 方案虽然有效，但有一个内在矛盾：辅助 loss 和语言模型的主 loss 是竞争关系。辅助 loss 鼓励路由均匀分布，主 loss 鼓励路由把 token 发给最合适的专家。两者之间的平衡需要仔细调节系数 α——α 太大会干扰模型收敛，α 太小则负载不均。

DeepSeek-V3 提出了一种更直接的方案：不给路由加 loss，而是直接调整路由的偏置。给每个专家维护一个偏置项 $b_{i}$ ，加在路由 logits 上——

正常路由: gate_logits = W_gate @ x
改进后:   gate_logits = W_gate @ x + b   （b 不参与梯度计算）

$b_{i}$ 的更新规则很简单：统计这一轮每个专家处理了多少 token，负载超过均值的专家偏置减小 γ（让它下次少收到 token），低于均值的偏置增大 γ（让它下次多收到 token）。γ 通常取很小的值（如 0.001），偏置在训练过程中自然收敛到平衡点。

# 每个 training step 结束后执行
expert_loads = count_tokens_per_expert(selected_experts)
mean_load = total_tokens * top_k / num_experts

for i in range(num_experts):
    if expert_loads[i] > mean_load:
        bias[i] -= gamma   # 超载 → 降偏置
    else:
        bias[i] += gamma   # 欠载 → 升偏置

和辅助 loss 相比，这个方案不干扰主 loss——偏置更新是手工规则，不产生梯度，不会和语言模型 loss 竞争。计算开销也几乎为零：每个 step 做一次加减，比辅助 loss 的反向传播便宜得多。

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

class MoELayerWithBias(nn.Module):
    """
    使用动态偏置的 MoE 层（DeepSeek-V3 风格）

    和普通 MoE 的区别：
    - 多了一个 expert_bias 参数，加在路由 logits 上
    - expert_bias 是 buffer（不参与梯度），由手工规则更新
    - 不再需要辅助 loss

    参数:
        d_model: 隐藏维度
        num_experts: 专家数量
        top_k: 每个 token 激活几个专家
        gamma: 偏置调整步长（默认 0.001）
    """
    def __init__(self, d_model, num_experts=8, top_k=2, gamma=0.001, expert_dim=None):
        super().__init__()
        self.num_experts = num_experts
        self.top_k = top_k
        self.gamma = gamma
        expert_dim = expert_dim or 4 * d_model

        self.gate = nn.Linear(d_model, num_experts, bias=False)
        # 关键：每个专家一个偏置，用 register_buffer 而非 Parameter
        # buffer 不会被 optimizer 更新，由 update_bias 手工调整
        self.register_buffer('expert_bias', torch.zeros(num_experts))

        self.experts = nn.ModuleList([
            nn.Sequential(
                nn.Linear(d_model, expert_dim),
                nn.ReLU(),
                nn.Linear(expert_dim, d_model)
            )
            for _ in range(num_experts)
        ])

    def forward(self, x):
        batch_size, seq_len, d_model = x.shape

        # 路由 logits = 线性变换 + 偏置（偏置不参与梯度）
        gate_logits = self.gate(x) + self.expert_bias

        top_k_logits, top_k_indices = torch.topk(gate_logits, self.top_k, dim=-1)
        top_k_weights = F.softmax(top_k_logits, dim=-1)

        output = torch.zeros_like(x)
        for b in range(batch_size):
            for s in range(seq_len):
                token = x[b, s]
                for k in range(self.top_k):
                    expert_idx = top_k_indices[b, s, k].item()
                    weight = top_k_weights[b, s, k]
                    expert_out = self.experts[expert_idx](token.unsqueeze(0)).squeeze(0)
                    output[b, s] += weight * expert_out

        return output, top_k_indices

    def update_bias(self, top_k_indices):
        """
        根据本轮路由结果更新偏置。
        在每次 optimizer.step() 之后调用。

        top_k_indices: [batch, seq_len, top_k]
        """
        total_tokens = top_k_indices.numel()  # batch * seq_len * top_k
        ideal_per_expert = total_tokens / self.num_experts

        # 统计每个专家被选中的次数
        expert_counts = torch.zeros(self.num_experts, device=self.expert_bias.device)
        for idx in top_k_indices.flatten():
            expert_counts[idx] += 1

        # 超载的降偏置，欠载的升偏置
        for i in range(self.num_experts):
            if expert_counts[i] > ideal_per_expert:
                self.expert_bias[i] -= self.gamma
            elif expert_counts[i] < ideal_per_expert:
                self.expert_bias[i] += self.gamma
            # 等于理想值：不变

print("含动态偏置的 MoE 层定义完成！")
print("关键差异：gate_logits = W_gate @ x + expert_bias")
print("expert_bias 是 buffer，不被 optimizer 更新，由 update_bias() 手工调整")

# 演示：动态偏置如何让负载从失衡走向均衡
import torch

torch.manual_seed(42)

moe_bias = MoELayerWithBias(d_model=16, num_experts=8, top_k=2, gamma=0.1)

# 构造输入：前一半 token 是后一半的放大版，诱导路由器集中选择某几个专家
x = torch.randn(1, 32, 16)
x[:, 16:, :] = x[:, :16, :] * 1.5

print("=== 偏置收敛过程 ===")
print(f"初始偏置: {[f'{b:.2f}' for b in moe_bias.expert_bias.tolist()]}")
print()

num_steps = 30
for step in range(num_steps):
    _, top_k_idx = moe_bias(x)
    moe_bias.update_bias(top_k_idx)

    # 每隔几步打印一次负载分布
    if step < 5 or step % 10 == 0 or step == num_steps - 1:
        expert_counts = torch.zeros(8)
        for idx in top_k_idx.flatten():
            expert_counts[idx] += 1
        counts = expert_counts.int().tolist()
        imbalance = max(counts) - min(counts)
        print(f"Step {step:2d}: 负载分布 {counts}  极差={imbalance}")

print()
print(f"最终偏置: {[f'{b:.2f}' for b in moe_bias.expert_bias.tolist()]}")
print()
print("关键观察：")
print("1. 初始时负载差异大——某些专家被过度使用，某些几乎空闲")
print("2. 每次 update_bias 后，过载专家的偏置下降，欠载专家的偏置上升")
print("3. 经过若干步，各专家负载趋于均衡")
print("4. 整个过程不使用辅助 loss，完全靠手工规则调整偏置")

4.1 MoE 训练最佳实践

训练 MoE 时，最容易犯的错误是只盯着 loss。loss 下降不代表 MoE 真的训好了：router 可能只用少数专家，模型表面在学习，实际上大部分专家没有收到足够梯度。

实践里建议每隔几个 step 打印四个指标：

指标	看什么	异常信号
`task_loss`	主任务有没有学会	loss 不降或剧烈抖动
`expert_usage`	每个专家被选中比例	某个专家长期占 80%+
`imbalance`	最忙专家和最闲专家差距	差距长期很大
`router_entropy`	router 是否过早变得偏执	entropy 很低，说明选择太尖

一句话经验：先让路由分布健康，再谈专家分工。 DeepSeek-V3 的 auxiliary-loss-free load balancing 也是围绕这个目标：不用辅助 loss 干扰主任务，而是用动态 expert bias 把负载拉回合理范围。

更工程化地说，MoE 训练日志里至少要同时看三条线：

task_loss 有没有稳定下降。
expert_usage 有没有长期塌缩到少数专家。
router_entropy 有没有太早变低。

如果只看第一条，就像只看考试总分，不看每道题是谁做的。MoE 的问题恰恰在于：总分可能还行，但很多专家根本没参与学习。

参考：DeepSeek-V3 Technical Report、NVIDIA NeMo DeepSeek-V3 文档、2025 年关于 expert load balancing 的后续工作，都继续把 routing collapse / load imbalance 当成 MoE 的核心问题。

import torch
import torch.nn.functional as F

def router_metrics(gate_logits, top_k_indices, num_experts):
    """打印 MoE 训练时最应该盯住的路由指标"""
    probs = F.softmax(gate_logits, dim=-1)
    entropy = -(probs * torch.log(probs + 1e-9)).sum(dim=-1).mean()

    counts = torch.bincount(top_k_indices.flatten(), minlength=num_experts).float()
    usage = counts / counts.sum()
    imbalance = counts.max() - counts.min()

    print("expert_usage:", [f"{u:.1%}" for u in usage.tolist()])
    print(f"imbalance: {imbalance.item():.0f} tokens")
    print(f"router_entropy: {entropy.item():.3f}")

    if usage.max() > 0.80:
        print("警告：router 可能塌缩了，某个专家吃掉了 80% 以上的路由。")
    elif imbalance > counts.mean():
        print("提示：负载仍然偏斜，可以继续观察 bias / balance 策略。")
    else:
        print("观察：专家使用率比较健康，可以继续看 task_loss 是否下降。")

torch.manual_seed(42)
monitor_moe = MoELayerWithBias(d_model=16, num_experts=8, top_k=2, gamma=0.1)
monitor_x = torch.randn(1, 32, 16)
monitor_x[:, 16:, :] = monitor_x[:, :16, :] * 1.5

print("=== 训练前：只看初始路由 ===")
with torch.no_grad():
    gate_logits = monitor_moe.gate(monitor_x) + monitor_moe.expert_bias
    _, top_k_idx = torch.topk(gate_logits, monitor_moe.top_k, dim=-1)
router_metrics(gate_logits, top_k_idx, monitor_moe.num_experts)

for _ in range(30):
    _, top_k_idx = monitor_moe(monitor_x)
    monitor_moe.update_bias(top_k_idx)

print("\n=== 动态 bias 调整后：再看路由 ===")
with torch.no_grad():
    gate_logits = monitor_moe.gate(monitor_x) + monitor_moe.expert_bias
    _, top_k_idx = torch.topk(gate_logits, monitor_moe.top_k, dim=-1)
router_metrics(gate_logits, top_k_idx, monitor_moe.num_experts)

print("\n关键观察：MoE 训练日志里应该长期保留 expert_usage 和 router_entropy。")
print("只看 task_loss，可能发现不了专家塌缩。")

5. MoE 的推理难题：所有专家都要加载

虽然每次只激活 k 个专家，但路由器是根据 token 内容动态选择的——同一个 batch 里不同 token 可能选中不同的专家组合。在推理开始之前，你不知道哪些专家会被选中。

这意味着所有专家的参数都要在显存里等着被调用：

Dense 7B:  显存 ≈ 14GB (FP16)
MoE 8×7B:  显存 ≈ 94GB (FP16) ← 接近 8 倍！

即使每个 token 只跑 2 个专家，但 8 个专家的权重全都占用显存。这就是 MoE 的 trade-off：计算量小，但显存大。

这也是为什么 MoE 模型虽然推理速度快（计算量小），但需要更多 GPU——不是因为计算不够，而是因为显存放不下。一个 MoE 8×7B 模型至少需要 2 张 A100（80GB）才能跑起来，而 Dense 7B 一张即可。

常见的缓解手段：

量化（INT4/INT8）压缩专家权重，把 94GB 降到 ~24GB
把不常用的专家 offload 到 CPU 内存，用到时再加载（增加延迟换取显存）
Expert Parallelism：不同专家分布在不同 GPU 上，token 通过 all-to-all 通信路由到对应 GPU

6. 著名的 MoE 模型

模型	总参数	激活参数	专家数	Top-K	备注
Mixtral 8×7B	47B	13B	8	2	2023，早期代表性开源 MoE
DeepSeek-V2	236B	21B	160	6	2024，MLA + Shared Expert
DeepSeek-V3	671B	37B	256	8	2024，auxiliary-loss-free 负载均衡/动态 expert bias
Qwen2.5-MoE	57B	14B	64	8	2024，共享专家 + 细粒度专家
GPT-4 (未公开)	未公开	未公开	未公开	未公开	闭源模型没有公开完整架构，不能把社区传闻写成事实

关键数字：激活参数 / 总参数取决于 top-k、shared experts、attention 和非专家参数，不能只用一个固定比例概括。MoE 的优势是总参数容量可以很大，但每个 token 只激活一部分专家；真实速度还要看 batch、通信、显存带宽和推理框架。

从公开模型可以观察到一个趋势：研究者越来越重视更细粒度专家、共享专家和负载均衡。但专家数量、top-k 和激活比例都要和通信开销、batch 形状、显存容量一起权衡。

7. 为什么 MoE 有效：参数和计算的分离

回头看 Dense 模型的 FFN。d_model=4096 时，FFN 升维到 16384，参数量约 2 × 4096 × 16384 ≈ 134M。每个 token 都要经过全部 134M 参数——不管这个 token 是简单的「的」还是复杂的「量子纠缠」。

MoE 把这件事拆开了。假设 8 个专家，每个专家参数量相同，top-k=2：

总参数量 = 8 × 134M ≈ 1B（翻了 8 倍）
每次推理计算量 = 2 × 134M ≈ 268M（只翻了 2 倍）

参数多了 8 倍，但计算只多了 2 倍。这多出来的 6 个专家的参数不是白费的——它们存储了更多样的知识，但推理时不需要全部激活。

更具体地说，如果模型在训练时发现某些知识组合（比如「法语语法」和「C++ 指针语法」）很少同时出现在同一个 token 上，那么这两个知识可以被分到不同的专家里。路由器学会根据 token 的特征把它们分发到正确的专家——法语 token 去专家 A，代码 token 去专家 B。这样每个专家的 FFN 可以学得更专、更深，而不需要在一个大 FFN 里同时处理所有类型的知识。

这是 MoE 和 Dense 最根本的区别：Dense 把参数和计算绑在一起，MoE 把它们解开了。

8. MoE 的进阶话题

DeepSeek-V2 的创新：Shared Expert + Routed Expert

DeepSeek-V2 在 MoE 结构上做了一个重要改进：除了 top-k 路由专家外，还增加了一个「共享专家」（Shared Expert）。

DeepSeek-V2 MoE:
  token → 共享专家（所有 token 必过） + 路由专家 top-k（按需选择）
  
  共享专家处理: 语法结构、常见词汇、常识推理（所有 token 都需要的通用能力）
  路由专家处理: 数学公式、代码语法、专业术语（特定领域知识）

为什么要这样设计？在标准 MoE 中，通用知识（比如「如何组织一个句子」）会被每个专家各自学一遍，浪费参数。把通用知识提取到一个共享专家里，让路由专家只专注各自的专长领域，参数利用效率更高。

Qwen2-MoE 也采用了类似的设计，在打印的模型结构中可以看到 shared_expert 和 shared_expert_gate 两个组件——后者是一个标量门，控制共享专家的输出强度。

Expert Parallelism

当专家数量很多（如 DeepSeek-V2 的 160 个专家），所有专家的参数无法放在一张 GPU 上。这时需要 Expert Parallelism：不同专家分布在不同的 GPU 上，token 根据路由结果被发送到对应 GPU 做 FFN 计算，然后返回原来的 GPU。

这个过程需要 all-to-all 通信——每张 GPU 可能要把 token 发送给其他所有 GPU，同时接收其他 GPU 发来的 token。All-to-all 通信是 MoE 训练和推理的主要开销来源，也是大规模 MoE 系统设计的核心难点。

细粒度专家分片

DeepSeekMoE（2024）提出了一个关键发现：在总参数量和计算量不变的前提下，把每个专家拆成更多、更小的专家，效果反而更好。

举例来说，一个 MoE 层有 8 个专家、top-2，和另一个有 64 个更小专家、top-8 的设计，总参数和计算量可以完全相同。区别在于组合选择空间：C(8,2) = 28，而 C(64,8) ≈ 4.4 × 10⁹。后者能组合出的专家搭配远多于前者，每种组合可以学到更专一的知识。

DeepSeekMoE 论文在其实验设置中报告了细粒度专家的收益；具体切分粒度和节省比例不能脱离模型规模、数据和训练配置直接泛化。参考：DeepSeekMoE。

专家数量和 Top-K 的选择

N（专家数）和 K（激活数）不是随意设定的，背后有几个共同考量：

计算预算。K × 单个专家参数量应该匹配目标 FLOPs。Mixtral 选 top-2 是让每个 token 的计算量接近 2 个 dense 专家（约 13B 激活参数）。DeepSeek-V3 选 top-8，但每个专家更小（FFN hidden dim 1408），总活跃参数控制在 37B。
细粒度分片。DeepSeekMoE 的实验显示，在其设置下细粒度专家能带来更好的参数/计算效率；是否总是更优仍要看模型、数据、路由和系统实现。DeepSeekMoE / DeepSeek-V2 / DeepSeek-V3 的公开设计体现了更细粒度专家和共享专家的方向；未正式公开的后续模型不要写成确定配置。
通信开销。多 GPU 训练中，K 越大，每个 token 要被发送到越多 GPU 上做计算，all-to-all 通信量线性增长。DeepSeek-V3 引入了 node-limited routing：限制每个 token 最多跨 4 个节点，即使 top-K=8 也只和 4 个节点通信。
组合多样性。C(N,K) 衡量模型能组合出多少种不同的专家搭配。N 越大、K/N 越小，组合空间越大。下面用代码算一下各模型的组合数。

import math

# 各模型的 MoE 配置对比
configs = [
    ("Mixtral 8×7B", 8, 2),
    ("DeepSeek-V2", 160, 6),
    ("Qwen2.5-MoE", 64, 8),
    ("DeepSeek-V3", 256, 8),
]

print("=== MoE 专家组合选择空间 C(N, K) ===")
print()
print(f"{'模型':20s} {'N':>4s} {'K':>3s} {'C(N,K)':>14s} {'激活比':>7s}")
for name, n, k in configs:
    c = math.comb(n, k)
    ratio = k / n * 100
    print(f"{name:20s} {n:4d} {k:3d} {c:14.2e} {ratio:6.1f}%")

print()
print("关键观察：")
print("1. Mixtral 的 C(8,2)=28，V4-Pro 的 C(384,6) ≈ 6.6×10¹¹")
print("2. 组合空间增长了 20 多个数量级——模型有更多种专家搭配可用")
print("3. 激活比例从 25%（Mixtral）降到 1.6%（V4-Pro）——更少激活，更多选择")
print()

# V3 vs V4 对比
print("=== 公开 MoE 配置如何解读 ===")
print()
rows = [
    ("", "DeepSeek-V3", "Mixtral 8×7B"),
    ("Routed 专家数", "256", "8"),
    ("Top-K", "8", "2"),
    ("Shared 专家", "有", "无"),
    ("公开重点", "负载均衡/细粒度专家", "简单清晰的 top-2 MoE"),
    ("解读", "大规模系统工程", "教学入门代表"),
]
for label, v3, v4 in rows:
    print(f"{label:14s} {v3:>10s} {v4:>10s}")
print()
print("V4 增加了专家数量但降低了 top-K——更少的激活比例，更大的组合空间。")
print("此外，V4 还引入了 CSA+HCA 混合注意力（支持 1M 上下文）、")
print("Muon 优化器（替代 AdamW）和 FP4 量化感知训练。")

小结

✅ Dense FFN 所有 token 共享一套参数，参数量 = 计算量
✅ MoE 通常替换 FFN，而不是替换 Attention：Attention 负责信息交换，FFN 负责逐 token 加工
✅ MoE 把一个大 FFN 拆成 N 个专家 FFN + 一个路由器
✅ 路由器对每个 token 打分，选 top-k 个专家，加权求和作为输出
✅ 总参数量 = N × 单个专家参数量（可以很大），每次计算量 = k × 单个专家参数量
✅ 参数和计算量「解耦」→ 总参数可以翻 8 倍，计算只翻 2 倍
✅ 训练难题：负载均衡 → 辅助 loss（传统）或动态偏置（DeepSeek-V3，不干扰主 loss）
✅ 推理难题：所有专家都要在显存里 → 量化 / 多 GPU / expert parallelism
✅ 著名 MoE：Mixtral 8×7B（约 47B 总参数、约 13B 激活）、DeepSeek-V2/V3 等公开报告中的大规模 MoE
✅ 激活比例取决于 top-k、shared experts 和非专家参数；不能用固定 10-20% 概括所有 MoE
✅ 细粒度专家分片：相同 FLOPs 下，更多更小的专家优于更少更大的专家
✅ 专家数量和 Top-K 受计算预算、通信开销、组合多样性共同约束 一句话总结：MoE 把 Dense 模型中绑定在一起的参数和计算量拆开了——参数负责存储知识（越多越好），计算负责推理速度（越少越好）。负载均衡是这个架构的核心工程挑战，DeepSeek-V3 的动态偏置方案让辅助 loss 不再是必需品。MoE 的公开演进方向包括更细粒度专家、共享专家、负载均衡和更好的系统实现；不要把未公开模型的传闻当成事实。

作业> 可以让 AI 帮忙解释思路，但不建议直接让 AI "做完这道题"。

作业 1：MoE 参数量与激活量计算Mixtral 8×7B 有 8 个专家 FFN，每个专家的参数量约为 7B。但 MoE 层之外还有共享的 Attention 参数约 1B。计算以下两个值：1. MoE 层的总参数量（8 个专家的总和）2. 每次前向传播实际激活的参数量（top-2 路由，只激活 2 个专家的 FFN + 共享 Attention）小提示：总参数 = 共享参数 + 8 × 单专家参数。激活参数 = 共享参数 + 2 × 单专家参数。

# 作业 1：MoE 参数量与激活量计算shared_params = 1e9   # 共享 Attention 参数（1B）expert_params = 7e9   # 单个专家 FFN 参数（7B）num_experts = 8top_k = 2# TODO: 计算 MoE 层总参数量total_params = None  # 在这里计算# TODO: 计算每次前向传播激活的参数量active_params = None  # 在这里计算assert total_params is not None, "请先计算总参数量"assert active_params is not None, "请先计算激活参数量"expected_total = shared_params + num_experts * expert_paramsexpected_active = shared_params + top_k * expert_paramsassert total_params == expected_total, f"总参数量应为 {expected_total/1e9:.0f}B"assert active_params == expected_active, f"激活参数量应为 {expected_active/1e9:.0f}B"print(f"✅ 作业 1 通过：")print(f"   总参数量: {total_params/1e9:.0f}B")print(f"   激活参数量: {active_params/1e9:.0f}B")print(f"   激活比例: {active_params/total_params:.1%}")print("   MoE 的核心优势：总参数可以很大，但每次只激活一小部分。")

作业 2：负载均衡分析假设有 4 个专家，处理 8 个 token。路由器给每个 token 选 top-1 专家。理想的负载均衡是每个专家处理 2 个 token。实际路由结果如下：专家 0 处理 5 个 token，专家 1 处理 2 个，专家 2 处理 1 个，专家 3 处理 0 个。计算路由熵 $H = - \sum_{i} p_{i} lo g_{2} (p_{i})$ ，其中 $p_{i}$ 是专家 $i$ 处理的 token 比例。并与理想均匀分布的熵 $H_{ma x} = lo g_{2} (4)$ 比较。小提示：均匀分布时 $H = lo g_{2} (N)$ ，熵越低说明负载越不均衡。

# 作业 2：负载均衡分析import math# 实际路由结果expert_counts = [5, 2, 1, 0]  # 每个专家处理的 token 数total_tokens = sum(expert_counts)# TODO: 计算每个专家的 token 比例 p_iproportions = None  # 在这里计算，列表 [p0, p1, p2, p3]# TODO: 计算路由熵 H = -sum(p_i * log2(p_i))# 注意：p_i = 0 时，p_i * log2(p_i) = 0entropy = None  # 在这里计算# 理想最大熵max_entropy = math.log2(4)assert proportions is not None, "请先计算比例"assert entropy is not None, "请先计算熵"expected_props = [5/8, 2/8, 1/8, 0/8]for i, (p, ep) in enumerate(zip(proportions, expected_props)):    assert abs(p - ep) < 0.01, f"专家 {i} 的比例应为 {ep:.3f}，你得到 {p:.3f}"expected_H = -(5/8*math.log2(5/8) + 2/8*math.log2(2/8) + 1/8*math.log2(1/8) + 0)assert abs(entropy - expected_H) < 0.01, f"熵应为 {expected_H:.4f}，你得到 {entropy:.4f}"print(f"✅ 作业 2 通过：")print(f"   实际熵: {entropy:.3f} bits")print(f"   最大熵: {max_entropy:.3f} bits（均匀分布）")print(f"   均衡度: {entropy/max_entropy:.1%}")print("   熵越低，说明路由器越偷懒——把大部分 token 都发给了少数专家。")

作业 3：MoE 推理显存估算MoE 模型推理时，所有专家的参数都需要加载到显存中（即使每次只用 top-k 个）。假设一个 MoE 模型：共享 Attention 参数 1B，8 个专家 FFN 每个 7B，FP16 推理。计算推理时模型权重占用的显存（单位 GB，FP16 每个参数 2 bytes）。小提示：推理时全部参数都要加载，总参数 × 2 bytes。

# 作业 3：MoE 推理显存估算shared_params = 1e9       # 1B 共享参数expert_params = 7e9       # 7B 每个专家num_experts = 8bytes_per_param = 2       # FP16# TODO: 计算总参数量total_params = None  # 在这里计算# TODO: 计算推理显存（GB）memory_gb = None  # 在这里计算assert total_params is not None, "请先计算总参数量"assert memory_gb is not None, "请先计算显存"expected_total = shared_params + num_experts * expert_paramsexpected_mem = expected_total * bytes_per_param / 1e9assert total_params == expected_total, f"总参数量应为 {expected_total/1e9:.0f}B"assert abs(memory_gb - expected_mem) < 0.1, f"显存应为 {expected_mem:.1f} GB"print(f"✅ 作业 3 通过：")print(f"   总参数量: {total_params/1e9:.0f}B")print(f"   推理显存: {memory_gb:.1f} GB（仅模型权重，不含 KV Cache）")print("   MoE 推理的显存瓶颈：即使只激活 2 个专家，全部 8 个都要加载。")

1. 普通 Transformer 的 FFN 层​

1.1 为什么 MoE 通常替换 FFN，而不是替换 Attention​

2. MoE 的核心思想​

2.1 对比：普通 Transformer Block vs MoE Block​

2.2 打印 HuggingFace 里的真实 MoE 模型​

3. MoE 的参数 vs 计算量​

4. MoE 的训练难题：负载均衡​

4.1 MoE 训练最佳实践​

5. MoE 的推理难题：所有专家都要加载​

6. 著名的 MoE 模型​

7. 为什么 MoE 有效：参数和计算的分离​

8. MoE 的进阶话题​

小结​

作业> 可以让 AI 帮忙解释思路，但不建议直接让 AI "做完这道题"。​