# 第 13 课：FFN 与 MoE
这一课只解决一个问题：Attention 汇总上下文之后，MiniMind 如何用 FFN 或 MoE 对每个 token 的 hidden 向量做非线性变换。
## 目录
- [0. 本节主线](#l13-mainline)
- [1. 原理讲解](#l13-principle)
- [2. 源码阅读顺序图](#l13-reading-order)
- [3. MiniMind 源码走读](#l13-source-walkthrough)
- [4. 本节必须会写 / 暂时不要求](#l13-must-write)
- [5. 手写模块](#l13-handwrite)
- [6. 对齐测试](#l13-alignment-test)
- [7. 阶段组装](#l13-stage-assembly)
- [8. 本节检查](#l13-check)
- [9. 下一课](#l13-next)
## 0. 本节主线
Transformer block 里，Attention 和 FFN/MoE 分工不同：
```text
Attention:
让每个 token 从上下文 token 里取信息。
FFN / MoE:
对每个 token 自己的 hidden 向量做非线性变换。
```
MiniMind block 的后半段是：
```text
attention 输出后的 hidden_states
-> post_attention_layernorm
-> FeedForward 或 MOEFeedForward
-> residual add
-> 下一个 block
```
Dense FFN 的核心公式可以写成：
$$
\begin{aligned}
u &= W_{\text{up}}x \\
g &= W_{\text{gate}}x \\
h &= \mathrm{SiLU}(g) \odot u \\
y &= W_{\text{down}}h
\end{aligned}
$$
也就是：先扩维得到候选特征 `u` 和门控特征 `g`，门控后再降回 `hidden_size`。
MoE 的核心公式可以写成：
$$
\begin{aligned}
r &= W_{\text{router}}x \\
p &= \mathrm{softmax}(r) \\
S &= \mathrm{TopK}(p, k) \\
\alpha_e &= \frac{p_e}{\sum_{j \in S} p_j}, \quad e \in S \\
y &= \sum_{e \in S}\alpha_e \cdot \mathrm{Expert}_e(x)
\end{aligned}
$$
也就是：router 先给每个 expert 一个概率，只选 top-k 个 expert 计算，再把被选中的概率归一化成权重 $\alpha_e$，最后加权求和。
一句话：
```text
FFN 是所有 token 共享同一个专家；MoE 是每个 token 先经过 router 选择少数专家，再加权汇总专家输出。
```
## 1. 原理讲解
### 1.1 FFN 到底干什么
第 11 课讲 Attention 时，我们一直在看 token 之间的信息流：
```text
位置 i 的 token
-> 用 Q 去看上下文 K
-> 用注意力权重汇总上下文 V
```
这解决的是“token 之间怎么交流”。
但拿到上下文信息后，每个 token 自己的向量还要继续加工。这个加工主要由 FFN 完成。
FFN 有一个重要特点：
```text
它不混合不同位置的 token。
```
如果输入是：
```text
hidden_states: [batch, seq, hidden_size]
```
那么 FFN 对每个位置独立做同一套变换：
```text
hidden_states[:, 0, :] -> FFN -> output[:, 0, :]
hidden_states[:, 1, :] -> FFN -> output[:, 1, :]
hidden_states[:, 2, :] -> FFN -> output[:, 2, :]
```
所以可以这样理解：
```text
Attention 负责横向：token 与 token 之间的信息交换。
FFN 负责纵向：单个 token 的 hidden 维度内部加工。
```
### 1.2 MiniMind 的 FFN 为什么有三个线性层
普通 MLP 可能写成：
```text
down( activation(up(x)) )
```
MiniMind 用的是 SwiGLU 风格。对单个 token 的 hidden 向量 `x`，可以写成：
$$
\begin{aligned}
u &= W_{\text{up}}x \\
g &= W_{\text{gate}}x \\
h &= \mathrm{SiLU}(g) \odot u \\
y &= W_{\text{down}}h
\end{aligned}
$$
如果放回一个 batch，形状是：
| 变量 | 形状 | 含义 |
|---|---|---|
| $B$ | 标量 | batch size |
| $S$ | 标量 | sequence length |
| $H$ | 标量 | `hidden_size` |
| $I$ | 标量 | `intermediate_size` |
| $x$ | `[B, S, H]` | FFN 输入，每个 token 一个 hidden 向量 |
| $u$ | `[B, S, I]` | $W_{\text{up}}x$，候选特征 |
| $g$ | `[B, S, I]` | $W_{\text{gate}}x$，门控特征 |
| $h$ | `[B, S, I]` | $\mathrm{SiLU}(g) \odot u$，门控后的中间特征 |
| $y$ | `[B, S, H]` | $W_{\text{down}}h$，FFN 输出 |
这里 $\odot$ 表示逐元素相乘。对应到源码就是：
```text
W_up -> up_proj
W_gate -> gate_proj
W_down -> down_proj
```
这里有三条线性投影：
```text
gate_proj: hidden_size -> intermediate_size
up_proj: hidden_size -> intermediate_size
down_proj: intermediate_size -> hidden_size
```
`gate_proj` 经过 `silu` 后像一个门：
$$
\mathrm{SiLU}(W_{\text{gate}}x)
$$
它会和 `up_proj(x)` 逐元素相乘：
$$
\mathrm{SiLU}(W_{\text{gate}}x) \odot (W_{\text{up}}x)
$$
这一步的意思是：
```text
up_proj 产生候选特征；
gate_proj 决定哪些特征更应该通过。
```
最后 `down_proj` 把维度从 `intermediate_size` 降回 `hidden_size`，这样才能和 residual 主干相加。
### 1.3 FFN 为什么必须回到 hidden_size
Transformer block 的主干 shape 一直是：
```text
[batch, seq, hidden_size]
```
FFN 中间可以扩维：
```text
[batch, seq, hidden_size]
-> gate/up
-> [batch, seq, intermediate_size]
```
但输出必须回到：
```text
[batch, seq, hidden_size]
```
因为 block 里还有 residual：
```text
hidden_states = hidden_states + mlp(norm(hidden_states))
```
如果 MLP 输出不是 `hidden_size`，就无法和原来的 `hidden_states` 相加。
### 1.4 MoE 是什么
MoE 是 Mixture of Experts，中文可以理解成“专家混合”。
Dense FFN 是：
```text
每个 token 都走同一个 FFN。
```
MoE 是：
```text
准备多个 FFN expert；
每个 token 先通过 router 打分；
只选择 top-k 个 expert 计算；
把 expert 输出按 router 权重加权求和。
```
这里的 `expert` 不是一个抽象名字，而是一个真正的子网络。
在 MiniMind 里，一个 expert 本质上就是一个独立的 `FeedForward`：
```text
expert 0 = FeedForward(...)
expert 1 = FeedForward(...)
expert 2 = FeedForward(...)
expert 3 = FeedForward(...)
```
它们结构一样，但参数各自独立。也就是说，`Expert_0(x)` 和 `Expert_2(x)` 都是 FFN 计算，但用的是两套不同的权重。
如果有 4 个 experts，top-k=1，那么某个 token 可能走：
```text
token A -> expert 2
token B -> expert 0
token C -> expert 2
token D -> expert 3
```
MoE 的核心取舍是：
```text
总参数量变大，因为有多个 experts；
每个 token 激活的参数量不一定同比变大，因为只走少数 experts。
```
这就是 README 里 `198M-A64M` 这种写法的含义：总参数更多，但单 token 激活参数接近较小模型。
### 1.5 Router 干什么
Router 是一个线性层。对单个 token 的 hidden 向量 `x`，它先产生每个 expert 的 logit：
$$
r = W_{\text{router}}x
$$
其中：
$$
r \in \mathbb{R}^{E}
$$
这里 $E$ 表示 expert 数量，也就是 `num_experts`。
然后用 softmax 变成概率：
$$
p_e = \frac{\exp(r_e)}{\sum_{j=1}^{E}\exp(r_j)}
$$
如果 `num_experts=4`，可能得到：
```text
p = [0.05, 0.10, 0.80, 0.05]
```
top-1 routing 会选择 expert 2。
如果 `num_experts_per_tok=2`，可能选择 expert 2 和 expert 1，再按权重加权：
$$
\begin{aligned}
S &= \mathrm{TopK}(p, 2) = \{2, 1\} \\
\alpha_2 &= \frac{0.80}{0.80 + 0.10} \\
\alpha_1 &= \frac{0.10}{0.80 + 0.10} \\
y &= \alpha_2 \cdot \mathrm{Expert}_2(x) + \alpha_1 \cdot \mathrm{Expert}_1(x)
\end{aligned}
$$
MiniMind 默认主线是 `4 experts / top-1`，所以先把 top-1 跑通最重要。
把它放回源码里的 batch 计算，形状是：
| 变量 | 形状 | 含义 |
|---|---|---|
| $B$ | 标量 | batch size |
| $S$ | 标量 | sequence length |
| $H$ | 标量 | `hidden_size` |
| $N$ | 标量 | 展平后的 token 数，$N = B \times S$ |
| $E$ | 标量 | expert 总数，`num_experts` |
| $K$ | 标量 | 每个 token 选择几个 expert，`num_experts_per_tok` |
| $x$ | `[B, S, H]` | MoE 输入 |
| $x_{\text{flat}}$ | `[N, H]` | 把 batch 和 seq 展平成 token 维 |
| $r$ | `[N, E]` | router logits，softmax 前的分数 |
| $p$ / `scores` | `[N, E]` | router softmax 概率 |
| $p_{t,e}$ | 标量 | 第 `t` 个 token 分给第 `e` 个 expert 的概率，即 `scores[t, e]` |
| `topk_idx` | `[N, K]` | 每个 token 选中的 expert id |
| `topk_weight` | `[N, K]` | 每个 token 选中 expert 的权重 |
| $y_{\text{flat}}$ | `[N, H]` | 展平后的 MoE 输出 |
| $y$ | `[B, S, H]` | reshape 回原始 batch/seq 的输出 |
### 1.6 MoE 为什么需要 aux loss
Router 有一个常见问题：它可能总是把大部分 token 分给少数 expert。
比如：
```text
expert 0: 95% token
expert 1: 2% token
expert 2: 2% token
expert 3: 1% token
```
这样其他 expert 学不到东西，MoE 退化成“一个很忙的 expert + 几个闲置 expert”。
所以 MoE 常用一个负载均衡辅助损失。先把符号说清楚：
| 符号 | 形状 | 含义 | 对应源码 |
|---|---|---|---|
| $t$ | 标量下标 | 第 `t` 个 token。源码里先把 `[B, S]` 展平成 `N = B * S` 个 token。 | `x_flat = x.view(-1, hidden_dim)` |
| $e$ | 标量下标 | 第 `e` 个 expert。比如 `e=2` 就是 `self.experts[2]`。 | `for i, expert in enumerate(self.experts)` |
| $E$ | 标量 | expert 总数。 | `config.num_experts` |
| $K$ | 标量 | 每个 token 选几个 expert。MiniMind 默认是 `1`。 | `config.num_experts_per_tok` |
| $p$ / `scores` | `[N, E]` | 所有 token 对所有 expert 的 router 概率。 | `scores = F.softmax(...)` |
| $p_{t,e}$ | 标量 | 第 `t` 个 token 被 router 分给第 `e` 个 expert 的 softmax 概率。 | `scores[t, e]` |
| $\mathrm{load}_e$ | 标量 | top-k 选择后，实际有多少比例的 token 选中了 expert `e`。 | `load = F.one_hot(topk_idx, ...).float().mean(0)` |
| $\mathrm{prob}_e$ | 标量 | router 对 expert `e` 的平均偏好，即所有 token 的 $p_{t,e}$ 平均值。 | `scores.mean(0)` |
| $L_{\text{aux}}$ | 标量 | 辅助损失，用来惩罚 router 总偏向少数 expert。 | `self.aux_loss = ...` |
其中最容易混的是 $p_{t,e}$。它不是最终输出，也不是标签，而是 router 的“分配概率”：
$$
p_{t,e} = \mathrm{softmax}(W_{\text{router}}x_t)_e
$$
如果第 3 个 token 的 router 分数是：
```text
scores[3] = [0.05, 0.10, 0.80, 0.05]
```
那么：
```text
p_{3,0} = 0.05
p_{3,1} = 0.10
p_{3,2} = 0.80
p_{3,3} = 0.05
```
如果是 top-1 routing，这个 token 会选择 expert 2。
但是 aux loss 不只看“最终选了谁”，也看 router 对所有 expert 的平均概率分布。
设一个 batch 展平后有 `N` 个 token，`E` 个 experts。MiniMind 里的负载均衡量可以这样理解：
$$
\begin{aligned}
\mathrm{load}_e
&= \frac{1}{N}\sum_{t=1}^{N}\mathbf{1}[\text{token } t \text{ 选择 expert } e] \\
\mathrm{prob}_e
&= \frac{1}{N}\sum_{t=1}^{N}p_{t,e} \\
L_{\text{aux}}
&= E \cdot \lambda_{\text{router}} \sum_{e=1}^{E}\mathrm{load}_e \cdot \mathrm{prob}_e
\end{aligned}
$$
这里 $\lambda_{\text{router}}$ 对应源码里的 `router_aux_loss_coef`。
这里的 $L_{\text{aux}}$ 不是语言模型的 next-token loss。语言模型主 loss 关心“下一个 token 预测对不对”；$L_{\text{aux}}$ 只关心“MoE 的专家使用是否太偏”。
训练时真正优化的是：
$$
L_{\text{train}} = L_{\text{lm}} + L_{\text{aux}}
$$
所以它的作用是给 router 一个轻微约束：在不破坏主任务的前提下，尽量不要让所有 token 都挤到同一个 expert。
注意一个源码细节：`F.one_hot(topk_idx, E)` 的形状是 `[N, K, E]`，MiniMind 接着 `.mean(0)` 后得到 `[K, E]`。默认 `K=1` 时，它等价于上面公式里的每个 expert 一个 $\mathrm{load}_e$；如果 `K>1`，源码会保留 top-k 每个位置的负载，再和 `scores.mean(0)` 广播相乘，最后 `sum()` 成一个标量 `aux_loss`。
## 2. 源码阅读顺序图
这节源码按这个顺序读：
```text
MiniMindConfig
-> MiniMindBlock
-> FeedForward
-> MOEFeedForward
-> MiniMindModel.forward 里的 aux_loss 汇总
```
对应文件都在：
```text
model/model_minimind.py
```
先看 config，是为了知道 dense FFN 和 MoE 由哪些参数控制。
再看 block，是为了知道 FFN/MoE 插在 Transformer block 的哪个位置。
然后看 dense FFN 和 MoE 的具体计算。
最后看 aux loss 如何被模型 forward 返回。
## 3. MiniMind 源码走读
### 第 1 步：Config 决定 FFN 和 MoE 尺寸
File: `model/model_minimind.py:25-45`
Read this to understand: FFN/MoE 的维度、专家数量、top-k 和 aux loss 系数都来自 config。
Code/config/template excerpt:
```python
self.hidden_act = kwargs.get("hidden_act", 'silu')
self.intermediate_size = kwargs.get("intermediate_size", math.ceil(hidden_size * math.pi / 64) * 64)
...
self.num_experts = kwargs.get("num_experts", 4)
self.num_experts_per_tok = kwargs.get("num_experts_per_tok", 1)
self.moe_intermediate_size = kwargs.get("moe_intermediate_size", self.intermediate_size)
self.norm_topk_prob = kwargs.get("norm_topk_prob", True)
self.router_aux_loss_coef = kwargs.get("router_aux_loss_coef", 5e-4)
```
This code shows:
- `hidden_act='silu'` 决定 FFN 激活函数。
- `intermediate_size` 是 dense FFN 的中间维度。
- `num_experts` 是 MoE expert 数量。
- `num_experts_per_tok` 是每个 token 选择几个 expert。
- `router_aux_loss_coef` 控制 MoE 辅助损失权重。
### 第 2 步：Block 决定用 dense FFN 还是 MoE
File: `model/model_minimind.py:178-194`
Read this to understand: MoE 不是额外插入一层，而是替代普通 FFN 的位置。
Code/config/template excerpt:
```python
self.mlp = FeedForward(config) if not config.use_moe else MOEFeedForward(config)
...
hidden_states = hidden_states + self.mlp(self.post_attention_layernorm(hidden_states))
```
This code shows:
- `config.use_moe=False` 时，block 使用普通 `FeedForward`。
- `config.use_moe=True` 时，block 使用 `MOEFeedForward`。
- 两者都必须输入 `[batch, seq, hidden_size]`，输出也必须是 `[batch, seq, hidden_size]`。
- 因为输出要和 `hidden_states` 做 residual add。
### 第 3 步：Dense FeedForward 是 SwiGLU
File: `model/model_minimind.py:136-146`
Read this to understand: MiniMind 的 FFN 不是简单两层 MLP，而是 gate/up/down 三投影。
Code/config/template excerpt:
```python
class FeedForward(nn.Module):
def __init__(self, config: MiniMindConfig, intermediate_size: int = None):
super().__init__()
intermediate_size = intermediate_size or config.intermediate_size
self.gate_proj = nn.Linear(config.hidden_size, intermediate_size, bias=False)
self.down_proj = nn.Linear(intermediate_size, config.hidden_size, bias=False)
self.up_proj = nn.Linear(config.hidden_size, intermediate_size, bias=False)
self.act_fn = ACT2FN[config.hidden_act]
def forward(self, x):
return self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
```
这段源码对应的数学公式是：
$$
y = W_{\text{down}}\left(\mathrm{SiLU}(W_{\text{gate}}x) \odot W_{\text{up}}x\right)
$$
This code shows:
- `gate_proj` 和 `up_proj` 都把 `hidden_size` 扩到 `intermediate_size`。
- `return self.down_proj(...)` 是源码写法，不是课程里的数学公式写法。
- `down_proj` 把维度降回 `hidden_size`。
- 输入输出 shape 不变，中间 shape 变大。
### 第 4 步：MoE 先把 token 展平
File: `model/model_minimind.py:156-160`
Read this to understand: Router 是按 token 路由的，所以先把 `[batch, seq]` 合并成 token 维。
Code/config/template excerpt:
```python
batch_size, seq_len, hidden_dim = x.shape
x_flat = x.view(-1, hidden_dim)
scores = F.softmax(self.gate(x_flat), dim=-1)
topk_weight, topk_idx = torch.topk(scores, k=self.config.num_experts_per_tok, dim=-1, sorted=False)
```
This code shows:
- `x_flat.shape = [batch * seq, hidden_dim]`。
- `self.gate(x_flat).shape = [batch * seq, num_experts]`。
- `scores.shape = [batch * seq, num_experts]`，是每个 token 对每个 expert 的概率。
- `topk_idx.shape = [batch * seq, num_experts_per_tok]`，是每个 token 选择的 expert id。
- `topk_weight.shape = [batch * seq, num_experts_per_tok]`，是对应 expert 的权重。
### 第 5 步：MoE 对每个 expert 收集 token 并写回
File: `model/model_minimind.py:161-170`
Read this to understand: MiniMind 的 MoE 实现是按 expert 循环，找到分给这个 expert 的 token，再把 expert 输出加回对应 token 位置。
Code/config/template excerpt:
```python
if self.config.norm_topk_prob: topk_weight = topk_weight / (topk_weight.sum(dim=-1, keepdim=True) + 1e-20)
y = torch.zeros_like(x_flat)
for i, expert in enumerate(self.experts):
mask = (topk_idx == i)
if mask.any():
token_idx = mask.any(dim=-1).nonzero().flatten()
weight = topk_weight[mask].view(-1, 1)
y.index_add_(0, token_idx, (expert(x_flat[token_idx]) * weight).to(y.dtype))
```
This code shows:
- `norm_topk_prob` 会把被选中的 top-k 权重重新归一化。
- `mask.shape = [batch * seq, num_experts_per_tok]`，找出哪些 token 选择了当前 expert。
- `token_idx.shape = [M]`，其中 `M` 是当前 expert 实际接收到的 token 数。
- `x_flat[token_idx].shape = [M, hidden_dim]`，这是当前 expert 的输入。
- `expert(x_flat[token_idx]).shape = [M, hidden_dim]`，这是当前 expert 的输出。
- `weight.shape = [M, 1]`，用于给当前 expert 的输出加权。
- `index_add_` 把 `[M, hidden_dim]` 的 expert 输出加回 `y.shape = [batch * seq, hidden_dim]` 的原 token 行号。
top-1 时，每个 token 只会写回一次。top-k 大于 1 时，同一个 token 会从多个 expert 收到输出，所以要加权相加。
### 第 6 步：MoE 计算 aux loss
File: `model/model_minimind.py:171-176`
Read this to understand: aux loss 是 router 的负载均衡约束，不是语言模型主 loss。
Code/config/template excerpt:
```python
if self.training and self.config.router_aux_loss_coef > 0:
load = F.one_hot(topk_idx, self.config.num_experts).float().mean(0)
self.aux_loss = (load * scores.mean(0)).sum() * self.config.num_experts * self.config.router_aux_loss_coef
else:
self.aux_loss = scores.new_zeros(1).squeeze()
return y.view(batch_size, seq_len, hidden_dim)
```
This code shows:
- 只有训练模式才计算 aux loss。
- `F.one_hot(topk_idx, num_experts).shape = [batch * seq, num_experts_per_tok, num_experts]`。
- `load = ...mean(0)` 的源码形状是 `[num_experts_per_tok, num_experts]`。
- MiniMind 默认 `num_experts_per_tok=1` 时，可以把它理解成每个 expert 被选中的比例，也就是公式里的 $\mathrm{load}_e$。
- `scores.mean(0).shape = [num_experts]`，表示 router 给每个 expert 的平均概率，也就是公式里的 $\mathrm{prob}_e$。
- `(load * scores.mean(0)).sum()` 最后变成一个标量，也就是 `aux_loss`。
- 输出 reshape 回 `[batch, seq, hidden_dim]`。
### 第 7 步：MiniMindModel 汇总所有 MoE 层的 aux loss
File: `model/model_minimind.py:230-232`
Read this to understand: 每层 MoE 都有自己的 aux loss，最终会加总后返回。
Code/config/template excerpt:
```python
hidden_states = self.norm(hidden_states)
aux_loss = sum([l.mlp.aux_loss for l in self.layers if isinstance(l.mlp, MOEFeedForward)], hidden_states.new_zeros(1).squeeze())
return hidden_states, presents, aux_loss
```
This code shows:
- dense FFN 没有 aux loss。
- 只有 `MOEFeedForward` 层的 `aux_loss` 会被汇总。
- CausalLM forward 会把 `aux_loss` 返回出去，后续训练脚本再把它加到训练 loss 里。
### 第 8 步：训练脚本把 aux loss 加到主 loss
File: `trainer/train_pretrain.py:35-38`
Read this to understand: `L_aux` 最终是怎么参与反向传播的。
Code/config/template excerpt:
```python
with autocast_ctx:
res = model(input_ids, labels=labels)
loss = res.loss + res.aux_loss
loss = loss / args.accumulation_steps
```
This code shows:
- `res.loss` 是 causal LM 的 next-token loss，也就是 $L_{\text{lm}}$。
- `res.aux_loss` 是 MoE router 的负载均衡损失，也就是 $L_{\text{aux}}$。
- `loss = res.loss + res.aux_loss` 才是最终反向传播的训练损失。
## 4. 本节必须会写 / 暂时不要求
必须会写：
1. Dense FeedForward：
$$
\begin{aligned}
u &= W_{\text{up}}x \\
g &= W_{\text{gate}}x \\
h &= \mathrm{SiLU}(g) \odot u \\
y &= W_{\text{down}}h
\end{aligned}
$$
2. Router 与 top-k 加权：
$$
\begin{aligned}
r &= W_{\text{router}}x \\
p &= \mathrm{softmax}(r) \\
S &= \mathrm{TopK}(p, k) \\
y &= \sum_{e \in S}\alpha_e \cdot \mathrm{Expert}_e(x)
\end{aligned}
$$
3. MoE token dispatch：
```text
找到每个 expert 对应的 token
只对这些 token 调用 expert
用 index_add_ 写回原位置
```
4. MoE aux loss：
$$
L_{\text{aux}}
= E \cdot \lambda_{\text{router}} \sum_{e=1}^{E}\mathrm{load}_e \cdot \mathrm{prob}_e
$$
暂时不要求：
```text
1. fused MoE kernel
2. expert parallel / 多机专家并行
3. capacity factor / token dropping
4. DeepSpeed-MoE / Megatron-LM 工程实现
5. 更复杂的 router 正则化
```
本节重点是对齐 MiniMind 的 PyTorch 原生实现。
## 5. 手写模块
本节你要补的是：
```text
course/impl/core/model_parts.py
```
### 5.1 补 `FeedForward`
接口已经在骨架里：
```python
class FeedForward(nn.Module):
def __init__(self, hidden_size: int, intermediate_size: int, hidden_act: str = "silu"):
...
def forward(self, x: torch.Tensor) -> torch.Tensor:
...
```
你要实现的行为：
- 输入：`x: [batch, seq, hidden_size]`
- 输出：`y: [batch, seq, hidden_size]`
- 核心公式：
$$
\begin{aligned}
u &= W_{\text{up}}x \\
g &= W_{\text{gate}}x \\
h &= \mathrm{SiLU}(g) \odot u \\
y &= W_{\text{down}}h
\end{aligned}
$$
对齐源码：
```text
model/model_minimind.py:136-146
```
注意：
- 层名字最好保持 `gate_proj`、`up_proj`、`down_proj`，方便测试脚本拷贝原模型权重。
- `bias=False`，和原源码一致。
- `hidden_act` 先支持 `silu` 即可。
### 5.2 补 `MOEFeedForward`
本节已经在骨架里加入 `MOEFeedForward`。你要实现的行为：
```text
输入:
x: [batch, seq, hidden_size]
中间:
x_flat: [batch * seq, hidden_size]
scores: [batch * seq, num_experts]
topk_idx: [batch * seq, num_experts_per_tok]
topk_weight: [batch * seq, num_experts_per_tok]
输出:
y: [batch, seq, hidden_size]
```
对齐源码：
```text
model/model_minimind.py:148-176
```
实现时先保证 top-1 能跑通，再保证 `num_experts_per_tok > 1` 时能用 `index_add_` 加权累加。
## 6. 对齐测试
本节新增对齐测试：
```text
course/impl/tests/test_ffn_moe_alignment.py
```
运行命令：
```bash
cd /home/sun/minimind
python course/impl/tests/test_ffn_moe_alignment.py
```
现在还没有实现时，这个测试会因为 `NotImplementedError` 失败。等你补完 `FeedForward` 和 `MOEFeedForward` 后，它应该打印：
```text
feedforward_max_abs_diff=...
moe_eval_max_abs_diff=...
moe_train_max_abs_diff=...
moe_aux_loss_abs_diff=...
```
验收标准：
```text
所有 max_abs_diff 都应该接近 0。
一般小于 1e-6 就可以认为对齐。
```
测试做了什么：
```text
1. 实例化 MiniMind 原版 FeedForward / MOEFeedForward。
2. 实例化 course/impl 里的手写版本。
3. 把原版权重复制到手写版本。
4. 用同一个输入 x 前向。
5. 比较输出和 aux_loss。
```
这比单纯看 shape 更严格：它要求你的实现和原源码数值一致。
## 7. 阶段组装
本节完成后，模型结构阶段会多出两个可复用模块：
```text
course/impl/core/model_parts.py::FeedForward
course/impl/core/model_parts.py::MOEFeedForward
```
它们以后会接到教学版 block 中：
```text
hidden_states
-> Attention
-> residual
-> RMSNorm
-> FeedForward 或 MOEFeedForward
-> residual
```
当前模型结构阶段还缺：
```text
RMSNorm
Attention
RoPE
KV cache
MiniMindBlock
CourseMiniMindForCausalLM
```
这些会在后续阶段组装或回补时补齐。第 13 课先把 FFN/MoE 这块手写清楚。
## 8. 本节检查
1. Attention 和 FFN 在 Transformer block 里的分工有什么区别？
2. MiniMind 的 FFN 为什么有 `gate_proj`、`up_proj`、`down_proj` 三个线性层？
3. 为什么 FFN 中间可以是 `intermediate_size`，但最终必须回到 `hidden_size`？
4. MoE 的 router 输出的 `scores` 表示什么？
5. top-1 MoE 和 dense FFN 的计算路径有什么区别？
6. 如果输入 `x.shape = [B, S, H]`，写出 `x_flat`、`scores`、`topk_idx`、`topk_weight`、`y` 的形状。
7. 为什么 MoE 需要 aux loss？它和 causal LM loss 是一回事吗？
8. `index_add_` 在 MoE 里解决了什么问题？
## 9. 下一课
第 14 课进入训练机制：学习率、梯度累积和混合精度。
从下一课开始，我们会从模型结构转向训练脚本。你要看懂的不只是 `loss.backward()`，还要看清：
```text
一个 batch 的 loss 如何累积；
什么时候 optimizer.step；
学习率怎么随 step 变化；
混合精度如何影响 forward/backward。
```