biscuitの博客

1. 一次训练前向传播的调用链#

预训练循环中真正启动模型计算的是下面两行：

1
res = model(input_ids, labels=labels)
2
loss = res.loss + res.aux_loss

这里的 model 是 MiniMindForCausalLM，不是仅包含 Transformer 主干的 MiniMindModel。一次调用会依次经过：

1
MiniMindForCausalLM.forward
2
└── MiniMindModel.forward
3
    ├── Token Embedding
4
    ├── MiniMindBlock.forward × L
5
    └── Final RMSNorm
6
└── LM Head
7
└── Shifted Cross Entropy

用一个函数概括，就是：

$I$ 是输入 Token ID，$Y$ 是训练标签，$Z$ 是词表 logits，$\mathcal{L}$ 是训练损失，$\theta$ 表示 Embedding、Attention、FFN、Norm 和 LM Head 中的全部可训练参数。

本文统一使用以下符号：

2. MiniMindConfig 将架构选择变成具体维度#

模型配置中的核心代码可以压缩成：

1
class MiniMindConfig(PretrainedConfig):
2
    def __init__(self, hidden_size=768, num_hidden_layers=8, **kwargs):
3
        self.hidden_size = hidden_size
4
        self.num_hidden_layers = num_hidden_layers
5
        self.vocab_size = kwargs.get("vocab_size", 6400)
6

7
        self.num_attention_heads = kwargs.get("num_attention_heads", 8)
8
        self.num_key_value_heads = kwargs.get("num_key_value_heads", 4)
9
        self.head_dim = kwargs.get(
10
            "head_dim", self.hidden_size // self.num_attention_heads
11
        )
12

13
        self.intermediate_size = kwargs.get(
14
            "intermediate_size",
15
            math.ceil(hidden_size * math.pi / 64) * 64
16
        )
17
        self.max_position_embeddings = kwargs.get(
18
            "max_position_embeddings", 32768
19
        )

Config 本身不处理输入张量。它的作用是确定模型中各参数矩阵的尺寸。默认配置对应：

这里有两个容易忽略的比例。首先，Query 总维度为：

其次，GQA 中每组 K/V 需要服务的 Query 头数为：

因此 MiniMind 默认有 8 个 Query 头、4 个 K/V 头，每个 K/V 头会对应 2 个 Query 头。max_position_embeddings 只决定 RoPE 缓冲区的可用位置范围，并不表示每条训练样本一定长达 32768 个 Token；训练脚本仍可以使用更短的 max_seq_len。

3. 模型初始化对应参数化函数的组成#

MiniMindModel.__init__ 创建 Transformer 主干：

1
class MiniMindModel(nn.Module):
2
    def __init__(self, config):
3
        self.embed_tokens = nn.Embedding(
4
            config.vocab_size,
5
            config.hidden_size
6
        )
7
        self.dropout = nn.Dropout(config.dropout)
8
        self.layers = nn.ModuleList([
9
            MiniMindBlock(l, config)
10
            for l in range(config.num_hidden_layers)
11
        ])
12
        self.norm = RMSNorm(
13
            config.hidden_size,
14
            eps=config.rms_norm_eps
15
        )

它对应四组数学对象：

其中 $E$ 是 Embedding 矩阵，每个 MiniMindBlock 是一个保持隐藏状态形状不变的函数，$\gamma_{\mathrm{final}}$ 是最终 RMSNorm 的可训练缩放参数。

ModuleList 只是把 $L$ 个 Block 注册为模型子模块，并不会在初始化阶段执行它们。真正的复合运算发生在 forward 的循环中。

MiniMindForCausalLM 在主干之后增加 LM Head：

1
class MiniMindForCausalLM(PreTrainedModel, GenerationMixin):
2
    def __init__(self, config=None):
3
        self.model = MiniMindModel(self.config)
4
        self.lm_head = nn.Linear(
5
            self.config.hidden_size,
6
            self.config.vocab_size,
7
            bias=False
8
        )
9

10
        if self.config.tie_word_embeddings:
11
            self.model.embed_tokens.weight = self.lm_head.weight

LM Head 的权重形状为：

当 tie_word_embeddings=True 时：

这不是复制一份数值相同的矩阵，而是让输入 Embedding 和输出 LM Head 引用同一个可训练参数。输入端用它按 Token ID 查行，输出端则执行矩阵乘法：

权重共享将两块原本各自需要 $Vd$ 个参数的矩阵合并为一块，但“查表”和“输出分类”仍然是两种不同运算。

4. Token ID 查表得到初始隐藏状态#

MiniMindModel.forward 首先读取输入形状并调用 Embedding：

1
def forward(self, input_ids, attention_mask=None,
2
            past_key_values=None, use_cache=False, **kwargs):
3
    batch_size, seq_length = input_ids.shape
4

5
    hidden_states = self.dropout(
6
        self.embed_tokens(input_ids)
7
    )

输入是整数矩阵(input_ids [B, T] input_ids [B, T] 就是一个 B 行 T 列的整数矩阵,每一行是一条样本的 token id 序列)：

例子

例如用预训练的配置：预训练: batch_size=128, max_seq_len=512 → input_ids [128, 512]

1
            ┌─── T = 512 个 token ───┐
2
       样本0 │ 1  882  53  ...  0  0 │
3
       样本1 │ 1  204  91  ...  0  0 │
4
       样本2 │ 1   77  16  ...  0  0 │
5
        ... │                       │
6
    样本127  │ 1  431  62  ...  0  0 │
7
            └───────────────────────┘
8
            ↑ B = 128 行

Embedding 按 ID 从 $E$ 中取行：

所以形状发生变化，查表后得到的是三维张量，多出来的那一维就是 hidden_size=768 —— 原来每个格子是一个整数(token id),现在每个格子被替换成一个 768 维向量：

这里不是把 $[B,T]$ 乘成 $[B,T,d]$，也不是对整数 ID 做连续数值计算，而是执行 $B\times T$ 次查表。

相同 Token ID 在不同位置会先取得 同一个 基础向量，Embedding 层完全不知道位置。区分它们的工作全部由后面的 Attention + RoPE 完成 —— 上下文融合之后,两个同一词的表示才会分化。

而且这个 [B, T, 768] 的形状在整个 8 层主干里保持不变 —— 每个 Block 进去是 [B,T,768],出来还是 [B,T,768]。只有最后过 LM Head 时最后一维才从 768 变成 6400。

Dropout 不改变张量形状。训练阶段它对元素施加随机掩码并进行尺度补偿，推理阶段则相当于恒等映射：

值得注意的是当前实践中 config.dropout 默认是 0.0,而三个训练脚本都没有覆盖它。所以 nn.Dropout(0.0) 在这个项目里全程等于恒等映射 —— 包括 Attention 里的 attn_dropout 和 resid_dropout 也一样。

1
self.dropout = kwargs.get("dropout", 0.0)   # config 默认 0.0

TIP

现代 LLM 预训练通常不用 dropout:数据量足够大、每个样本基本只见一次(epochs=1),过拟合风险很低,而 dropout 会拖慢收敛。Dropout 是”数据少、多轮重复训练”时代的正则手段。不过 LoRA 那次训了 5 个 epoch、数据只有 1.6 万条 —— 那个场景其实是有过拟合风险的,dropout 保持 0 可能是退化原因之一。

MiniMind 没有把位置向量直接加到 Embedding 上。它从预计算的 RoPE 表中切出当前位置：

1
start_pos = (
2
    past_key_values[0][0].shape[1]
3
    if past_key_values[0] is not None else 0
4
)
5

6
position_embeddings = (
7
    self.freqs_cos[start_pos:start_pos + seq_length],
8
    self.freqs_sin[start_pos:start_pos + seq_length]
9
)

若当前输入长度为 $T$，则：

训练时没有历史 KV Cache，通常有 start_pos=0；逐 Token 推理时，start_pos 等于已经缓存的历史长度。position_embeddings 会继续传入 Attention，在 Q/K 投影后参与旋转，而不会直接修改此处的 $H^{(0)}$。

5. MiniMindBlock 对应两次 Pre-Norm 残差更新#

单个 Decoder Block 的主干代码是：

1
def forward(self, hidden_states, position_embeddings,
2
            past_key_value=None, use_cache=False,
3
            attention_mask=None):
4
    residual = hidden_states
5
    hidden_states, present_key_value = self.self_attn(
6
        self.input_layernorm(hidden_states),
7
        position_embeddings,
8
        past_key_value,
9
        use_cache,
10
        attention_mask
11
    )
12
    hidden_states = hidden_states + residual
13

14
    hidden_states = hidden_states + self.mlp(
15
        self.post_attention_layernorm(hidden_states)
16
    )
17
    return hidden_states, present_key_value

设进入第 $\ell$ 个 Block 的隐藏状态为 $H^{(\ell)}$，代码可以逐行翻译为：

代码反复覆盖 hidden_states，数学表达则给每个中间结果单独命名。第一处 residual 显式保存 $H^{(\ell)}$；第二次残差没有再创建变量，因为此时 hidden_states 本身已经是 $U^{(\ell)}$。

整条主干上的形状始终保持：

只有形状一致，Attention 输出和 FFN 输出才能分别与残差逐元素相加。present_key_value 是供生成阶段复用的旁路输出，不会代替当前层的隐藏状态。

6. Block 堆叠形成完整 Transformer 主干#

MiniMindModel.forward 用循环依次调用全部 Block：

1
presents = []
2
for layer, past_key_value in zip(self.layers, past_key_values):
3
    hidden_states, present = layer(
4
        hidden_states,
5
        position_embeddings,
6
        past_key_value=past_key_value,
7
        use_cache=use_cache,
8
        attention_mask=attention_mask
9
    )
10
    presents.append(present)
11

12
hidden_states = self.norm(hidden_states)
13
return hidden_states, presents, aux_loss

这对应函数复合：

虽然每层输入输出形状都为 $[B,T,d]$，不同层拥有各自独立的 Attention、FFN 和 Norm 参数，因此每一层学习的变换并不相同。

循环结束后执行最终 RMSNorm：

presents 保存每一层产生的 KV Cache；训练时默认 use_cache=False，其中的元素为 None。aux_loss 用于 MoE 路由负载均衡；普通 FFN 配置下没有 MoE 层，因此它为 0。两者都是主隐藏状态之外的附加输出。

7. LM Head 将隐藏状态变成词表 Logits#

MiniMindForCausalLM.forward 先调用主干，再执行输出投影：

1
hidden_states, past_key_values, aux_loss = self.model(
2
    input_ids,
3
    attention_mask,
4
    past_key_values,
5
    use_cache,
6
    **kwargs
7
)
8

9
slice_indices = (
10
    slice(-logits_to_keep, None)
11
    if isinstance(logits_to_keep, int)
12
    else logits_to_keep
13
)
14
logits = self.lm_head(
15
    hidden_states[:, slice_indices, :]
16
)

对每个位置的最终隐藏向量 $h_{b,t}\in\mathbb{R}^{d}$，LM Head 计算：

按照 PyTorch 的批量张量写法，等价于：

因此最后一个维度从隐藏特征 $d$ 变成词表类别 $V$：

$Z$ 是 logits，不是概率。对某个位置应用 Softmax 才得到条件概率：

logits_to_keep 是推理优化参数，可以只计算最后若干位置的 logits。默认值为 0，而 Python 中 -0 等于 0，所以 slice(-0, None) 实际是 slice(0, None)，训练默认仍会保留全部 $T$ 个位置。

代码没有在 LM Head 后显式调用 Softmax，因为训练使用的 F.cross_entropy 会在数值更稳定的实现中合并 Log-Softmax 与负对数似然。

8. 标签错位把输入序列变成下一个 Token 目标#

预训练数据集先把输入复制为标签，再把 PAD 位置设为 -100：

1
tokens = [bos_token_id] + tokens + [eos_token_id]
2
input_ids = tokens + [pad_token_id] * (
3
    max_length - len(tokens)
4
)
5

6
input_ids = torch.tensor(input_ids, dtype=torch.long)
7
labels = input_ids.clone()
8
labels[input_ids == pad_token_id] = -100
9
return input_ids, labels

看到 labels = input_ids.clone() 时，容易误以为模型在学习原样复制输入。真正的时间错位发生在模型内部：

1
if labels is not None:
2
    x = logits[..., :-1, :].contiguous()
3
    y = labels[..., 1:].contiguous()
4

5
    loss = F.cross_entropy(
6
        x.view(-1, x.size(-1)),
7
        y.view(-1),
8
        ignore_index=-100
9
    )

假设一条补齐后的序列是：

1
input_ids = [BOS, Thinking, Machines, EOS, PAD]
2
labels    = [BOS, Thinking, Machines, EOS, -100]

切片后形成：

也就是：

若完整 logits 形状为 $[B,T,V]$，切片和展平对应：

F.cross_entropy 要求每一行包含 $V$ 个类别分数，并为这一行提供一个整数类别 ID，所以代码把 Batch 维和时间维合并为一个样本维。contiguous() 先确保切片结果在内存中连续，使后续 view 能安全重塑。

设有效目标位置集合为：

默认平均交叉熵为：

这正是因果语言模型的 Token 级负对数似然。ignore_index=-100 会从求和与平均分母中排除 PAD 目标。交叉熵与 NLL 的完整关系可以回看KL 散度与交叉熵笔记。

因果掩码和标签错位承担不同职责：Causal Mask 保证位置 $t$ 的隐藏状态不能读取未来 Token；标签错位则规定位置 $t$ 应该预测 Token $t+1$。缺少前者会泄露答案，缺少后者则不会形成标准的 Next-Token Prediction。

例子#

为了把维度看清，假设使用以下配置：

一次训练前向传播的形状为：

这个表揭示了完整前向传播中两次关键的最后一维变化：Embedding 把整数 ID 变成 $d$ 维隐藏向量，LM Head 再把 $d$ 维隐藏向量变成 $V$ 维词表分数。中间的 $L$ 个 Decoder Block 都保持 $[B,T,d]$ 不变。

10. 从代码阅读过渡到组件实现#

现在可以把 MiniMind 主干压缩成一组代码与数学对象的对照：

到这里，已经能从训练脚本一路追踪到标量 Loss，但仍把 self.self_attn(...) 当作一个整体。下一篇将进入 Attention.forward，逐段对应 Q/K/V 投影、QK Norm、RoPE、GQA、Causal Mask、注意力矩阵和 KV Cache。