minimind/course/09_embedding_norm_residual_...

659 lines
18 KiB
Markdown
Raw Blame History

This file contains ambiguous Unicode characters!

This file contains ambiguous Unicode characters that may be confused with others in your current locale. If your use case is intentional and legitimate, you can safely ignore this warning. Use the Escape button to highlight these characters.

# 第 9 课:Embedding、RMSNorm 与残差 Block
这一课只解决一个问题:`input_ids` 进入 MiniMind 后,如何变成 `hidden_states`,并在每个 block 里保持 `[batch, seq, hidden_size]` 这条主干 shape。
## 目录
- [0. 本节主线](#l09-mainline)
- [1. 本节要懂的 5 个原理](#l09-principles)
- [2. 变量流转](#l09-flow)
- [3. 原理一:Embedding 把 token id 变成向量](#l09-embedding)
- [4. 原理二:主干 shape 必须保持不变](#l09-main-shape)
- [5. 原理三:RMSNorm 只归一化最后一维](#l09-rmsnorm)
- [6. 原理四:残差连接要求 shape 对齐](#l09-residual)
- [7. 原理五:final norm 和 lm_head 完成输出转换](#l09-final-norm-lm-head)
- [8. 实验验证](#l09-experiment)
- [9. 本节检查](#l09-check)
- [10. 下一课](#l09-next)
<a id="l09-mainline"></a>
## 0. 本节主线
MiniMind 主干数据流是:
```text
input_ids [batch, seq]
-> embed_tokens
-> hidden_states [batch, seq, hidden_size]
-> input_layernorm
-> self_attn
-> residual add
-> post_attention_layernorm
-> mlp
-> residual add
-> 下一个 block
-> final norm
-> lm_head
-> logits [batch, seq, vocab_size]
```
所以本节的核心不是 Attention 公式,而是先看懂 Transformer block 的外壳:
```text
Embedding 负责升维;
RMSNorm 负责稳定每个 token 的 hidden 向量;
Attention/MLP 负责变换 hidden_states;
Residual 负责把变换结果加回主干;
整个 block 输入输出 shape 都保持 [batch, seq, hidden_size]。
```
<a id="l09-principles"></a>
## 1. 本节要懂的 5 个原理
| 原理 | 要理解什么 | 源码证据 |
|---|---|---|
| Embedding 把 token id 变成向量 | `input_ids` 是整数,进入模型后变成 `hidden_states` | `model/model_minimind.py:196-204`, `model/model_minimind.py:209-215` |
| 主干 shape 必须保持不变 | block 输入输出都要是 `[batch, seq, hidden_size]`,后面才能层层相接 | `model/model_minimind.py:186-194`, `model/model_minimind.py:220-230` |
| RMSNorm 只归一化最后一维 | 它不改变 batch/seq/hidden shape,只稳定 hidden 向量尺度 | `model/model_minimind.py:50-60` |
| 残差连接要求 shape 对齐 | Attention 和 MLP 输出必须能加回原来的 `hidden_states` | `model/model_minimind.py:186-194` |
| final norm 和 lm_head 完成输出转换 | 主干最后仍是 hidden,再映射成词表 logits | `model/model_minimind.py:230-232`, `model/model_minimind.py:245-249` |
学完本节,你应该能读懂 `MiniMindModel.forward` 的主干流程,并能说清楚为什么每个 block 都要保持同一个 hidden shape。
<a id="l09-flow"></a>
## 2. 变量流转
最小 shape 流:
```text
input_ids
shape = [batch_size, seq_len]
dtype = torch.long
含义 = token id
embed_tokens(input_ids)
shape = [batch_size, seq_len, hidden_size]
含义 = 每个 token 的向量表示
MiniMindBlock(hidden_states)
shape = [batch_size, seq_len, hidden_size]
含义 = 被 Attention 和 MLP 改写过的 token 表示
lm_head(hidden_states)
shape = [batch_size, seq_len, vocab_size]
含义 = 每个位置对下一个 token 的预测分数
```
这一节先记住一条规则:
```text
只要还在 Transformer 主干里,shape 基本保持 [batch, seq, hidden_size]。
```
<a id="l09-embedding"></a>
## 3. 原理一:Embedding 把 token id 变成向量
### 原理讲解
tokenizer 输出的 `input_ids` 只是整数序列,例如:
```text
[1, 345, 1024, 2]
```
这些数字本身没有连续空间里的几何意义。模型第一步要做的是查表:
```text
token id -> hidden_size 维向量
```
这个查表矩阵就是 embedding:
```text
embed_tokens.weight: [vocab_size, hidden_size]
```
如果 `input_ids.shape = [2, 6]`,`hidden_size=64`,那么 embedding 后:
```text
hidden_states.shape = [2, 6, 64]
```
从这一刻开始,模型不再处理整数 token id,而是在处理每个 token 的向量表示。
### 源码证据 A:embedding 矩阵在哪里创建
文件:`model/model_minimind.py:196-204`
看它是为了理解:模型如何根据 config 创建输入 embedding。
源码:
```python
class MiniMindModel(nn.Module):
def __init__(self, config: MiniMindConfig):
super().__init__()
self.config = config
self.vocab_size, self.num_hidden_layers = config.vocab_size, config.num_hidden_layers
self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size)
self.dropout = nn.Dropout(config.dropout)
self.layers = nn.ModuleList([MiniMindBlock(l, config) for l in range(self.num_hidden_layers)])
self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
```
这段代码说明:
- `nn.Embedding(config.vocab_size, config.hidden_size)` 创建 token 查表矩阵。
- embedding 的行数等于词表大小。
- embedding 的列数等于模型 hidden 维度。
- `layers` 和 `norm` 都是在 embedding 后继续处理 hidden_states 的模块。
### 源码证据 B:input_ids 如何变成 hidden_states
文件:`model/model_minimind.py:209-215`
看它是为了理解:forward 开始时变量含义如何变化。
源码:
```python
def forward(self, input_ids, attention_mask=None, past_key_values=None, use_cache=False, **kwargs):
batch_size, seq_length = input_ids.shape
if hasattr(past_key_values, 'layers'): past_key_values = None
past_key_values = past_key_values or [None] * len(self.layers)
start_pos = past_key_values[0][0].shape[1] if past_key_values[0] is not None else 0
hidden_states = self.dropout(self.embed_tokens(input_ids))
```
这段代码说明:
- 进入 forward 时,`input_ids` 只有 `[batch, seq]` 两维。
- `self.embed_tokens(input_ids)` 后,多出最后一维 `hidden_size`。
- 变量名从 `input_ids` 变成 `hidden_states`,说明后面处理的是向量表示。
### 理解到这一步就够
你应该能说清楚:
```text
input_ids 是 token 的编号;
hidden_states 是 token 的向量;
embedding 是编号到向量的查表。
```
暂时不用看:
- embedding 初始化分布。
- embedding 语义空间如何形成。
- 词向量可视化。
<a id="l09-main-shape"></a>
## 4. 原理二:主干 shape 必须保持不变
### 原理讲解
Transformer block 是可以一层接一层堆起来的,前提是:
```text
第 N 层输出 shape = 第 N+1 层输入 shape
```
MiniMind 主干里这个 shape 是:
```text
[batch_size, seq_len, hidden_size]
```
Attention 会改写每个 token 对上下文的吸收方式,MLP 会改写每个 token 自身的非线性表示,但它们最终都要回到 `hidden_size` 维。否则下一层 block 就接不上。
这就是为什么你读源码时要一直盯住这条主干:
```text
hidden_states
```
它像一条流水线:每个 block 接过来,处理完,再交给下一个 block。
### 源码证据 A:block 输入输出都是 hidden_states
文件:`model/model_minimind.py:186-194`
看它是为了理解:一个 block 如何接收并返回同一个主干变量。
源码:
```python
def forward(self, hidden_states, position_embeddings, past_key_value=None, use_cache=False, attention_mask=None):
residual = hidden_states
hidden_states, present_key_value = self.self_attn(
self.input_layernorm(hidden_states), position_embeddings,
past_key_value, use_cache, attention_mask
)
hidden_states += residual
hidden_states = hidden_states + self.mlp(self.post_attention_layernorm(hidden_states))
return hidden_states, present_key_value
```
这段代码说明:
- block 输入变量叫 `hidden_states`。
- Attention 的输出仍然赋值给 `hidden_states`。
- MLP 处理后仍然回到 `hidden_states`。
- block 最后返回新的 `hidden_states` 给下一层。
### 源码证据 B:多层 block 串起来
文件:`model/model_minimind.py:220-230`
看它是为了理解:同一个主干变量如何穿过所有层。
源码:
```python
presents = []
for layer, past_key_value in zip(self.layers, past_key_values):
hidden_states, present = layer(
hidden_states,
position_embeddings,
past_key_value=past_key_value,
use_cache=use_cache,
attention_mask=attention_mask
)
presents.append(present)
hidden_states = self.norm(hidden_states)
```
这段代码说明:
- 每一层都接收上一层的 `hidden_states`。
- 每一层都返回更新后的 `hidden_states`。
- 循环结束后再做 final RMSNorm。
### 理解到这一步就够
你应该能画出:
```text
hidden_states -> block0 -> hidden_states -> block1 -> hidden_states -> final norm
```
暂时不用看:
- `present_key_value` 的缓存内容。
- RoPE 的位置编码细节。
- Attention 内部矩阵乘法。
<a id="l09-rmsnorm"></a>
## 5. 原理三:RMSNorm 只归一化最后一维
### 原理讲解
RMSNorm 是一种归一化层。它处理的对象是每个 token 的 hidden 向量。
如果:
```text
hidden_states.shape = [batch, seq, hidden_size]
```
RMSNorm 做的是沿最后一维计算均方根:
```text
sqrt(mean(x^2 over hidden_size))
```
然后把这个 token 的 hidden 向量缩放到更稳定的尺度。
重要的是:
```text
RMSNorm 不改变 shape
```
输入是 `[batch, seq, hidden_size]`,输出仍然是 `[batch, seq, hidden_size]`。
MiniMind 使用的是 Pre-Norm 结构,也就是在两个子步骤里都先 norm,再送进对应模块。注意这不是两条并列路径,而是同一个 block 里的顺序执行:
```text
hidden_states -> RMSNorm -> Attention -> residual add
hidden_states -> RMSNorm -> MLP -> residual add
```
### 源码证据 A:RMSNorm 的实现
文件:`model/model_minimind.py:50-60`
看它是为了理解:RMSNorm 为什么只影响数值尺度,不改变 shape。
源码:
```python
class RMSNorm(torch.nn.Module):
def __init__(self, dim: int, eps: float = 1e-5):
super().__init__()
self.eps = eps
self.weight = nn.Parameter(torch.ones(dim))
def norm(self, x):
return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
def forward(self, x):
return (self.weight * self.norm(x.float())).type_as(x)
```
这段代码说明:
- `mean(-1, keepdim=True)` 表示沿最后一维计算。
- `keepdim=True` 保留维度,所以可以广播回原 shape。
- `self.weight` 的长度是 `hidden_size`,用于每个 hidden 维度的缩放。
- `type_as(x)` 会把输出转回输入 dtype。
### 源码证据 B:RMSNorm 放在 block 的两个位置
文件:`model/model_minimind.py:178-184`
看它是为了理解:一个 block 里为什么有两个 norm。
源码:
```python
class MiniMindBlock(nn.Module):
def __init__(self, layer_id: int, config: MiniMindConfig):
super().__init__()
self.self_attn = Attention(config)
self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
self.post_attention_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
self.mlp = FeedForward(config) if not config.use_moe else MOEFeedForward(config)
```
这段代码说明:
- `input_layernorm` 在 Attention 前。
- `post_attention_layernorm` 在 MLP 前。
- 两个 norm 的维度都是 `hidden_size`。
### 理解到这一步就够
你应该能说清楚:
```text
RMSNorm 是对每个 token 的 hidden 向量做尺度归一化;
它不会改变 batch、seq、hidden 三个维度。
```
暂时不用看:
- RMSNorm 和 LayerNorm 的严格差异。
- 归一化为什么有助于深层训练稳定。
<a id="l09-residual"></a>
## 6. 原理四:残差连接要求 shape 对齐
### 原理讲解
残差连接就是把模块输出加回原来的主干:
```text
new_hidden = module(norm(hidden)) + old_hidden
```
它的好处是:每个模块不需要从零重写全部表示,而是在原表示上增加一个变化量。
但残差相加有一个硬要求:
```text
两个张量 shape 必须一致
```
所以 Attention 输出必须是 `[batch, seq, hidden_size]`,MLP 输出也必须是 `[batch, seq, hidden_size]`。这也是为什么 Attention 末尾有 `o_proj` 映射回 `hidden_size`,FFN 末尾有 `down_proj` 映射回 `hidden_size`。这些细节后面会继续展开。
### 源码证据 A:Attention 残差
文件:`model/model_minimind.py:186-193`
看它是为了理解:Attention 输出如何加回 block 输入。
源码:
```python
def forward(self, hidden_states, position_embeddings, past_key_value=None, use_cache=False, attention_mask=None):
residual = hidden_states
hidden_states, present_key_value = self.self_attn(
self.input_layernorm(hidden_states), position_embeddings,
past_key_value, use_cache, attention_mask
)
hidden_states += residual
```
这段代码说明:
- `residual` 保存的是进入 Attention 前的主干。
- `self_attn(...)` 输出新的 `hidden_states`。
- `hidden_states += residual` 把 Attention 输出加回原主干。
### 源码证据 B:MLP 残差
文件:`model/model_minimind.py:192-194`
看它是为了理解:MLP 输出如何加回 Attention 后的主干。
源码:
```python
hidden_states += residual
hidden_states = hidden_states + self.mlp(self.post_attention_layernorm(hidden_states))
return hidden_states, present_key_value
```
这段代码说明:
- Attention 残差后,`hidden_states` 继续作为 MLP 的输入主干。
- MLP 前先做 `post_attention_layernorm`。
- MLP 输出再次加回 `hidden_states`。
### 理解到这一步就够
你应该能说清楚:
```text
残差连接不是拼接,也不是替换;
它是相同 shape 的张量相加。
```
暂时不用看:
- 残差连接的梯度流数学证明。
- DeepNorm、ReZero 等残差变体。
<a id="l09-final-norm-lm-head"></a>
## 7. 原理五:final norm 和 lm_head 完成输出转换
### 原理讲解
所有 block 处理完后,模型还没有直接输出 token 概率。
这时的主干仍然是:
```text
hidden_states: [batch, seq, hidden_size]
```
MiniMind 先做一次 final RMSNorm,然后 `MiniMindForCausalLM` 用 `lm_head` 把最后一维从 `hidden_size` 映射成 `vocab_size`:
```text
[batch, seq, hidden_size]
-> lm_head
-> [batch, seq, vocab_size]
```
这个 `[batch, seq, vocab_size]` 就是 logits。后面 causal LM loss 会把 logits 和 labels shift 后做 cross entropy。
### 源码证据 A:主干最后的 final norm
文件:`model/model_minimind.py:220-232`
看它是为了理解:所有 block 之后,模型返回的还是 hidden_states。
源码:
```python
for layer, past_key_value in zip(self.layers, past_key_values):
hidden_states, present = layer(
hidden_states,
position_embeddings,
past_key_value=past_key_value,
use_cache=use_cache,
attention_mask=attention_mask
)
presents.append(present)
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
```
这段代码说明:
- block 循环结束后做 final norm。
- `MiniMindModel` 返回的是 hidden_states,不是 logits。
- logits 是外层 `MiniMindForCausalLM` 继续生成的。
### 源码证据 B:lm_head 把 hidden 映射成 logits
文件:`model/model_minimind.py:245-249`
看它是为了理解:什么时候从 hidden_size 变成 vocab_size。
源码:
```python
def forward(self, input_ids, attention_mask=None, past_key_values=None, use_cache=False, logits_to_keep=0, labels=None, **kwargs):
hidden_states, past_key_values, aux_loss = self.model(input_ids, attention_mask, past_key_values, use_cache, **kwargs)
slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
logits = self.lm_head(hidden_states[:, slice_indices, :])
loss = None
```
这段代码说明:
- `self.model(...)` 只负责主干 hidden 表示。
- `self.lm_head(...)` 负责变成词表 logits。
- logits 的最后一维是 `vocab_size`。
### 理解到这一步就够
你应该能说清楚:
```text
MiniMindModel 输出 hidden_states;
MiniMindForCausalLM 再用 lm_head 输出 logits。
```
暂时不用看:
- `logits_to_keep` 的推理优化用途。
- loss shift 的细节,第 10 课会集中讲。
<a id="l09-experiment"></a>
## 8. 实验验证
### 实验 A:跟踪一个 tiny 模型的 block shape
这个实验验证:
```text
input_ids
-> embedding
-> layernorm
-> attention
-> residual add
-> layernorm
-> mlp
-> residual add
-> final norm
-> logits
```
运行:
```bash
cd /home/sun/minimind
PYTHONDONTWRITEBYTECODE=1 python course/labs/trace_block_shapes.py
```
记录:
```text
input_ids.shape =
embed_tokens.weight.shape =
after embed_tokens + dropout shape =
block input shape =
after input_layernorm shape =
self_attn output shape =
after attention residual add shape =
mlp output shape =
block output shape =
logits.shape =
```
你应该看到:
```text
input_ids.shape = (2, 6)
after embed_tokens + dropout shape = (2, 6, 64)
block input shape = (2, 6, 64)
block output shape = (2, 6, 64)
logits.shape = (2, 6, 6400)
```
这说明:
- embedding 把 `[batch, seq]` 升成 `[batch, seq, hidden_size]`。
- block 内部多次变换,但主干 shape 不变。
- `lm_head` 最后把 hidden 映射到 `vocab_size`。
### 实验 B:改变 hidden_size 看主干 shape
运行:
```bash
PYTHONDONTWRITEBYTECODE=1 python course/labs/trace_block_shapes.py --hidden_size 96 --num_hidden_layers 3
```
记录:
```text
after embed_tokens + dropout shape =
每个 block output shape =
logits.shape =
```
你应该看到:
```text
hidden_states 的最后一维从 64 变成 96;
logits 的最后一维仍然是 vocab_size=6400。
```
这说明:
- `hidden_size` 控制主干向量宽度。
- `vocab_size` 控制最终输出类别数。
- block 数量变多,只是主干多走几层,不改变每层输入输出 shape。
<a id="l09-check"></a>
## 9. 本节检查
如果你真懂了本节,应该能不看答案说清楚:
1. `input_ids` 和 `hidden_states` 的 shape 和含义分别是什么。
2. embedding 为什么会让张量从二维变成三维。
3. RMSNorm 沿哪一维做归一化,为什么不改变 shape。
4. 残差连接为什么要求 Attention/MLP 输出都是 `[batch, seq, hidden_size]`。
5. `MiniMindModel` 输出的是 logits 还是 hidden_states。
6. logits 的最后一维为什么是 `vocab_size`。
<a id="l09-next"></a>
## 10. 下一课
第 10 课进入 `Causal LM forward 与 loss shift`:我们会集中看 `MiniMindForCausalLM.forward`,把 `hidden_states -> logits -> shifted logits/labels -> cross entropy` 这条链路讲清楚。