24 KiB

Raw Blame History

第 13 课：FFN 与 MoE

这一课只解决一个问题：Attention 汇总上下文之后，MiniMind 如何用 FFN 或 MoE 对每个 token 的 hidden 向量做非线性变换。

0. 本节主线

Transformer block 里，Attention 和 FFN/MoE 分工不同：

Attention:
让每个 token 从上下文 token 里取信息。

FFN / MoE:
对每个 token 自己的 hidden 向量做非线性变换。

MiniMind block 的后半段是：

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，最后加权求和。

一句话：

FFN 是所有 token 共享同一个专家；MoE 是每个 token 先经过 router 选择少数专家，再加权汇总专家输出。

1. 原理讲解

1.1 FFN 到底干什么

第 11 课讲 Attention 时，我们一直在看 token 之间的信息流：

位置 i 的 token
-> 用 Q 去看上下文 K
-> 用注意力权重汇总上下文 V

这解决的是“token 之间怎么交流”。

但拿到上下文信息后，每个 token 自己的向量还要继续加工。这个加工主要由 FFN 完成。

FFN 有一个重要特点：

它不混合不同位置的 token。

如果输入是：

hidden_states: [batch, seq, hidden_size]

那么 FFN 对每个位置独立做同一套变换：

hidden_states[:, 0, :] -> FFN -> output[:, 0, :]
hidden_states[:, 1, :] -> FFN -> output[:, 1, :]
hidden_states[:, 2, :] -> FFN -> output[:, 2, :]

所以可以这样理解：

Attention 负责横向：token 与 token 之间的信息交换。
FFN 负责纵向：单个 token 的 hidden 维度内部加工。

1.2 MiniMind 的 FFN 为什么有三个线性层

普通 MLP 可能写成：

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 表示逐元素相乘。对应到源码就是：

W_up    -> up_proj
W_gate  -> gate_proj
W_down  -> down_proj

这里有三条线性投影：

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)

这一步的意思是：

up_proj 产生候选特征；
gate_proj 决定哪些特征更应该通过。

最后 down_proj 把维度从 intermediate_size 降回 hidden_size，这样才能和 residual 主干相加。

1.3 FFN 为什么必须回到 hidden_size

Transformer block 的主干 shape 一直是：

[batch, seq, hidden_size]

FFN 中间可以扩维：

[batch, seq, hidden_size]
-> gate/up
-> [batch, seq, intermediate_size]

但输出必须回到：

[batch, seq, hidden_size]

因为 block 里还有 residual：

hidden_states = hidden_states + mlp(norm(hidden_states))

如果 MLP 输出不是 hidden_size，就无法和原来的 hidden_states 相加。

1.4 MoE 是什么

MoE 是 Mixture of Experts，中文可以理解成“专家混合”。

Dense FFN 是：

每个 token 都走同一个 FFN。

MoE 是：

准备多个 FFN expert；
每个 token 先通过 router 打分；
只选择 top-k 个 expert 计算；
把 expert 输出按 router 权重加权求和。

这里的 expert 不是一个抽象名字，而是一个真正的子网络。

在 MiniMind 里，一个 expert 本质上就是一个独立的 FeedForward：

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 可能走：

token A -> expert 2
token B -> expert 0
token C -> expert 2
token D -> expert 3

MoE 的核心取舍是：

总参数量变大，因为有多个 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，可能得到：

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。

比如：

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 分数是：

scores[3] = [0.05, 0.10, 0.80, 0.05]

那么：

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. 源码阅读顺序图

这节源码按这个顺序读：

MiniMindConfig
-> MiniMindBlock
-> FeedForward
-> MOEFeedForward
-> MiniMindModel.forward 里的 aux_loss 汇总

对应文件都在：

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:

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:

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:

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:

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:

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:

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:

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:

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. 本节必须会写 / 暂时不要求

必须会写：

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}

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}

MoE token dispatch：

找到每个 expert 对应的 token
只对这些 token 调用 expert
用 index_add_ 写回原位置

MoE aux loss：


L_{\text{aux}}
= E \cdot \lambda_{\text{router}} \sum_{e=1}^{E}\mathrm{load}_e \cdot \mathrm{prob}_e

暂时不要求：

1. fused MoE kernel
2. expert parallel / 多机专家并行
3. capacity factor / token dropping
4. DeepSpeed-MoE / Megatron-LM 工程实现
5. 更复杂的 router 正则化

本节重点是对齐 MiniMind 的 PyTorch 原生实现。

5. 手写模块

本节你要补的是：

course/impl/core/model_parts.py

5.1 补 `FeedForward`

接口已经在骨架里：

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}

对齐源码：

model/model_minimind.py:136-146

注意：

层名字最好保持 gate_proj、up_proj、down_proj，方便测试脚本拷贝原模型权重。
bias=False，和原源码一致。
hidden_act 先支持 silu 即可。

5.2 补 `MOEFeedForward`

本节已经在骨架里加入 MOEFeedForward。你要实现的行为：

输入:
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]

对齐源码：

model/model_minimind.py:148-176

实现时先保证 top-1 能跑通，再保证 num_experts_per_tok > 1 时能用 index_add_ 加权累加。

6. 对齐测试

本节新增对齐测试：

course/impl/tests/test_ffn_moe_alignment.py

运行命令：

cd /home/sun/minimind
python course/impl/tests/test_ffn_moe_alignment.py

现在还没有实现时，这个测试会因为 NotImplementedError 失败。等你补完 FeedForward 和 MOEFeedForward 后，它应该打印：

feedforward_max_abs_diff=...
moe_eval_max_abs_diff=...
moe_train_max_abs_diff=...
moe_aux_loss_abs_diff=...

验收标准：

所有 max_abs_diff 都应该接近 0。
一般小于 1e-6 就可以认为对齐。

测试做了什么：

1. 实例化 MiniMind 原版 FeedForward / MOEFeedForward。
2. 实例化 course/impl 里的手写版本。
3. 把原版权重复制到手写版本。
4. 用同一个输入 x 前向。
5. 比较输出和 aux_loss。

这比单纯看 shape 更严格：它要求你的实现和原源码数值一致。

7. 阶段组装

本节完成后，模型结构阶段会多出两个可复用模块：

course/impl/core/model_parts.py::FeedForward
course/impl/core/model_parts.py::MOEFeedForward

它们以后会接到教学版 block 中：

hidden_states
-> Attention
-> residual
-> RMSNorm
-> FeedForward 或 MOEFeedForward
-> residual

当前模型结构阶段还缺：

RMSNorm
Attention
RoPE
KV cache
MiniMindBlock
CourseMiniMindForCausalLM

这些会在后续阶段组装或回补时补齐。第 13 课先把 FFN/MoE 这块手写清楚。

8. 本节检查

Attention 和 FFN 在 Transformer block 里的分工有什么区别？
MiniMind 的 FFN 为什么有 gate_proj、up_proj、down_proj 三个线性层？
为什么 FFN 中间可以是 intermediate_size，但最终必须回到 hidden_size？
MoE 的 router 输出的 scores 表示什么？
top-1 MoE 和 dense FFN 的计算路径有什么区别？
如果输入 x.shape = [B, S, H]，写出 x_flat、scores、topk_idx、topk_weight、y 的形状。
为什么 MoE 需要 aux loss？它和 causal LM loss 是一回事吗？
index_add_ 在 MoE 里解决了什么问题？

9. 下一课

第 14 课进入训练机制：学习率、梯度累积和混合精度。

从下一课开始，我们会从模型结构转向训练脚本。你要看懂的不只是 loss.backward()，还要看清：

一个 batch 的 loss 如何累积；
什么时候 optimizer.step；
学习率怎么随 step 变化；
混合精度如何影响 forward/backward。

24 KiB Raw Blame History Unescape Escape

第 13 课：FFN 与 MoE

目录

0. 本节主线

1. 原理讲解

1.1 FFN 到底干什么

1.2 MiniMind 的 FFN 为什么有三个线性层

1.3 FFN 为什么必须回到 hidden_size

1.4 MoE 是什么

1.5 Router 干什么

1.6 MoE 为什么需要 aux loss

2. 源码阅读顺序图

3. MiniMind 源码走读

第 1 步：Config 决定 FFN 和 MoE 尺寸

第 2 步：Block 决定用 dense FFN 还是 MoE

第 3 步：Dense FeedForward 是 SwiGLU

第 4 步：MoE 先把 token 展平

第 5 步：MoE 对每个 expert 收集 token 并写回

第 6 步：MoE 计算 aux loss

第 7 步：MiniMindModel 汇总所有 MoE 层的 aux loss

第 8 步：训练脚本把 aux loss 加到主 loss

4. 本节必须会写 / 暂时不要求

5. 手写模块

5.1 补 FeedForward

5.2 补 MOEFeedForward

6. 对齐测试

7. 阶段组装

8. 本节检查

9. 下一课

24 KiB

Raw Blame History

5.1 补 `FeedForward`

5.2 补 `MOEFeedForward`