从零实现 nanoGPT（四）：从单头注意力搭建完整 GPT

从零实现 nanoGPT（四）：从单头注意力搭建完整 GPT

上一篇，我们实现了单头因果自注意力，让模型能够根据上下文预测下一个字符。

相比第二篇的 Bigram 模型，单头注意力已经能够利用当前位置之前的信息，而不再只依赖当前字符。

但我们的模型仍然比较简单：

字符 ID
   ↓
字符嵌入 + 位置嵌入
   ↓
单头因果自注意力
   ↓
线性预测层
   ↓
下一个字符的预测分数

这里有几个值得思考的问题：

1. 一个注意力头是否足够？能否同时学习多种不同的信息关联方式？
2. 注意力完成信息汇总后，能否进一步加工得到的特征？
3. 如果不断增加网络层数，怎样让模型更容易训练？

为了解决这些问题，我们将在单头注意力的基础上，加入多头注意力、前馈网络、残差连接和 LayerNorm，最终组装出完整的 Transformer Block。

接着，通过堆叠多个 Block，搭建一个能够训练、预测和生成文本的小型 GPT。

本篇不再修改第一篇的数据预处理，也不改变第二篇的训练目标：

给定一段字符序列，让模型预测每个位置的下一个字符。

---

一、从单头注意力到多头注意力

1.1 为什么需要多个注意力头？

上一篇，我们实现了单头因果自注意力。

它的核心是通过输入特征生成三个矩阵：

• Query（Q）：当前位置想获取什么信息。
• Key（K）：每个位置提供什么信息用于匹配。
• Value（V）：每个位置实际提供什么特征。

注意力计算公式为：
$$softmax\; (\; QKT\; dk\; +M)\; V\; \backslash operatorname\{softmax\}\; \backslash left(\; \backslash frac\{QK^T\}\{\backslash sqrt\{d\_k\}\}+M\; \backslash right)V$$ softmax(dk QKT+M)V

其中， $MM$ M 是因果掩码，用于阻止模型关注未来位置。

不过，一个注意力头只有一套独立的 Q、K、V 投影参数。

例如，输入一句话：

The cat sat on the mat.

模型在处理某个位置时，可能需要关注不同的信息，例如相邻字符关系、单词内部结构以及更远位置的上下文。

单个注意力头也可能学习到多种模式，但我们希望模型能够并行学习多种信息汇总方式。

因此，可以同时使用多个注意力头。

这就是 Multi-Head Attention（多头注意力）。

需要注意，我们并没有规定某个头必须负责语法，另一个头必须负责其他语言现象。每个头学到什么，是训练过程中由数据和优化共同决定的。

1.2 多头注意力的计算过程

假设设置：

n_embd = 128
n_head = 4
head_size = n_embd // n_head

这里：

• n_embd：每个位置的特征维度。
• n_head：注意力头的数量。
• head_size：每个头输出的特征维度。

因为：
$$128/4=32128\; /\; 4\; =\; 32$$ 128/4=32

所以，每个注意力头输出 32 维特征。

假设输入张量形状为：

(B, T, 128)

其中：

• B：batch_size，一次输入多少条序列。
• T：序列长度。
• 128：每个位置的特征维度。

多头注意力的处理过程如下：

                输入 x
             (B, T, 128)
                  │
       ┌──────────┼──────────┬──────────┐
       │          │          │          │
     Head 1     Head 2     Head 3     Head 4
       │          │          │          │
   (B,T,32)  (B,T,32)  (B,T,32)  (B,T,32)
       │          │          │          │
       └──────────┴──────────┴──────────┘
                  │
              torch.cat
                  │
             (B, T, 128)
                  │
             Linear 投影
                  │
             (B, T, 128)

四个头分别学习独立的 Q、K、V 投影。

得到结果后，沿着最后一个特征维度拼接：

out = torch.cat(
    [head(x) for head in self.heads],
    dim=-1
)

dim=-1 表示沿最后一个维度拼接。

因此：

(B, T, 32) × 4

得到：

(B, T, 128)

再经过线性投影：

out = self.proj(out)

这一步会在拼接后的特征空间中进行线性变换，使不同头的输出能够进一步混合。

1.3 更新上一篇的 Head

我们不需要重新设计单头注意力，直接复用上一篇的 Head 即可。

不过，这次增加一个 dropout 参数。

在 model.py 中更新：

class Head(nn.Module):
    """单头因果自注意力"""

    def __init__(self, n_embd, head_size, block_size, dropout=0.0):
        super().__init__()

        self.query = nn.Linear(n_embd, head_size, bias=False)
        self.key = nn.Linear(n_embd, head_size, bias=False)
        self.value = nn.Linear(n_embd, head_size, bias=False)

        self.register_buffer(
            "tril",
            torch.tril(torch.ones(block_size, block_size))
        )

        self.dropout = nn.Dropout(dropout)

    def forward(self, x):
        B, T, C = x.shape

        q = self.query(x)
        k = self.key(x)
        v = self.value(x)

        scores = q @ k.transpose(-2, -1)
        scores = scores * (k.shape[-1] ** -0.5)

        scores = scores.masked_fill(
            self.tril[:T, :T] == 0,
            float("-inf")
        )

        weights = F.softmax(scores, dim=-1)
        weights = self.dropout(weights)

        return weights @ v

整体逻辑和上一篇一致：

1. 将输入映射为 Q、K、V。
2. 计算注意力分数。
3. 根据特征维度进行缩放。
4. 使用下三角掩码阻止访问未来位置。
5. 使用 Softmax 得到注意力权重。
6. 将注意力权重与 V 相乘。

唯一新增的是：

self.dropout = nn.Dropout(dropout)

以及：

weights = self.dropout(weights)

Dropout 在训练过程中随机将部分注意力权重置零，用于正则化，降低模型对特定连接的过度依赖。

这里的 Dropout 并不会重新归一化剩余的注意力权重，而是对保留下来的元素进行相应缩放。

调用：

model.train()

Dropout 生效。

调用：

model.eval()

Dropout 关闭。

另外，由于第三篇的 AttentionLanguageModel 仍然可以使用默认 dropout=0.0，所以这里增加默认参数后，不会破坏上一篇原有的模型调用方式。

1.4 实现 MultiHeadAttention

现在，我们可以把多个 Head 组合起来。

class MultiHeadAttention(nn.Module):
    """多头因果自注意力"""

    def __init__(self, n_embd, n_head, block_size, dropout):
        super().__init__()

        if n_embd % n_head != 0:
            raise ValueError("n_embd 必须能被 n_head 整除")

        head_size = n_embd // n_head

        self.heads = nn.ModuleList([
            Head(n_embd, head_size, block_size, dropout)
            for _ in range(n_head)
        ])

        self.proj = nn.Linear(n_embd, n_embd)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x):
        out = torch.cat(
            [head(x) for head in self.heads],
            dim=-1
        )

        out = self.proj(out)
        return self.dropout(out)

这里有两个值得注意的地方。

第一，为什么使用 nn.ModuleList ？

self.heads = nn.ModuleList([...])

nn.ModuleList 可以让 PyTorch 正确登记其中的子模块。

这些子模块的参数会被包含在：

model.parameters()

中，因此能够参与优化。

如果只是将多个注意力头存放在普通 Python 列表中，PyTorch 不会通过这个列表自动登记子模块。

第二，为什么还需要 self.proj ？

self.proj = nn.Linear(n_embd, n_embd)

因为 torch.cat 只是将多个头输出的特征拼接起来。

线性投影能够进一步学习如何组合各个头的信息。

最后再增加一次 Dropout，完成多头注意力模块。

---

二、前馈网络：加工注意力汇总后的信息

有了多头注意力之后，模型可以从上下文中收集信息。

但这是否意味着每个位置的特征已经得到了充分处理？

还不够。

注意力主要负责不同位置之间的信息交互。

在信息汇总完成后，我们还希望模型对当前位置的特征进行进一步变换。

这就是 Feed Forward Network，简称 FFN（前馈网络）。

2.1 FFN 的结构

本篇采用以下结构：

输入特征
   ↓
Linear(n_embd, 4 * n_embd)
   ↓
GELU
   ↓
Linear(4 * n_embd, n_embd)
   ↓
Dropout
   ↓
输出特征

例如：

n_embd = 128

那么特征维度依次变化：

128 → 512 → 128

首先扩展特征维度，再经过非线性激活，最后映射回原来的维度。

数学表达式为：
$$W2GELU(W1x+b1)+b2\; W\_2\backslash operatorname\{GELU\}(W\_1x+b\_1)+b\_2$$ W2GELU(W1x+b1)+b2

在这个公式中，暂时省略了 Dropout。

需要说明的是，FFN 并不会直接在不同 token 位置之间交换信息。

对于输入：

(B, T, n_embd)

它在每个位置上应用相同的特征变换，最后输出：

(B, T, n_embd)

因此：

• 注意力：从上下文中获取信息。
• FFN：对当前位置的特征进行加工。

两者配合，才能形成更完整的特征更新过程。

2.2 为什么需要 GELU？

如果只是连续使用两个线性层：

nn.Linear(128, 512)
nn.Linear(512, 128)

那么它们本质上仍然可以合并为一个线性变换。

因为：
$$y=W2(W1x+b1)+b2\; y=W\_2(W\_1x+b\_1)+b\_2$$ y=W2(W1x+b1)+b2

展开得到：
$$y=(W2W1)x+(W2b1+b2)\; y=(W\_2W\_1)x+(W\_2b\_1+b\_2)$$ y=(W2W1)x+(W2b1+b2)

仍然是仿射变换。

因此，我们需要在中间加入非线性激活函数。

本篇使用：

nn.GELU()

GELU 是 GPT 系列模型中常见的激活函数之一。

它能够为网络引入非线性，使模型可以学习更复杂的特征关系。

2.3 实现 FeedForward

代码非常简单：

class FeedForward(nn.Module):
    """逐位置前馈网络"""

    def __init__(self, n_embd, dropout):
        super().__init__()

        self.net = nn.Sequential(
            nn.Linear(n_embd, 4 * n_embd),
            nn.GELU(),
            nn.Linear(4 * n_embd, n_embd),
            nn.Dropout(dropout)
        )

    def forward(self, x):
        return self.net(x)

这里使用：

nn.Sequential(...)

将多个操作按顺序连接起来。

对于输入：

x.shape == (16, 64, 128)

前馈网络输出仍然是：

(16, 64, 128)

保持输入输出维度相同非常重要，因为下一步我们要使用残差连接。

---

三、残差连接：为什么不直接替换特征？

现在，我们已经有两个模块：

MultiHeadAttention
FeedForward

最直接的想法是把它们连续调用：

x = self.attention(x)
x = self.feed_forward(x)

但这样存在一个问题：

每经过一个模块，原来的特征都会被新的输出替换。

当模型层数越来越多时，这种结构可能增加深层网络的优化难度。

因此，Transformer 使用了残差连接。

3.1 残差连接的数学形式

普通的特征变换：
$$y=F(x)y=F(x)$$ y=F(x)

残差连接改为：
$$y=x+F(x)y=x+F(x)$$ y=x+F(x)

也就是：

保留原来的特征，再加上新模块计算出的变化。

例如：

x = x + self.attention(x)

可以理解为：

原始特征 x ─────────────┐
                       +
Attention(x) ──────────┘
                       │
                    新特征

模型不一定需要每一层重新构建整个特征表示，而可以在原有表示的基础上进行更新。

3.2 为什么残差连接有利于梯度传播？

假设：
$$y=x+F(x)y=x+F(x)$$ y=x+F(x)

那么：
$$I+\; \partial F(x)\partial x\; I+\backslash frac\{\backslash partial\; F(x)\}\{\backslash partial\; x\}$$ I+∂x∂F(x)

其中， $II$ I 是单位映射对应的雅可比矩阵。

与单纯的 $F(x)F(x)$ F(x) 相比，残差结构额外提供了一条直接经过加法的梯度传播路径。

这有助于缓解深层网络训练中的优化困难。

不过，残差连接并不意味着梯度一定不会消失或爆炸，它只是改善深层网络训练的重要手段之一。

3.3 残差连接为什么要求形状一致？

观察：

x = x + self.attention(x)

这是张量之间的逐元素相加。

因此，两个张量需要具有相同的形状（或者满足 PyTorch 的广播规则）。

在这里，我们刻意让它们形状完全相同：

x.shape == (B, T, n_embd)

attention(x).shape == (B, T, n_embd)

前馈网络也是如此：

feed_forward(x).shape == (B, T, n_embd)

这正是前面设计多头注意力输出投影和 FFN 最后一层时，都使用 n_embd 作为输出维度的原因。

---

四、LayerNorm：让特征尺度更加稳定

解决了残差连接后，还有一个问题：

随着多个模块不断叠加，特征的数值分布可能发生变化，进而影响训练稳定性。

Transformer 中通常使用 LayerNorm 进行归一化。

4.1 LayerNorm 的计算原理

对于某个位置上的特征向量：
$$x=(x1,x2,...,xd)\; x=(x\_1,x\_2,\backslash ldots,x\_d)$$ x=(x1,x2,...,xd)

首先计算均值：
$$\mu =1d\; \sum i=1d\; xi\; \backslash mu=\backslash frac\{1\}\{d\}\backslash sum\_\{i=1\}^\{d\}x\_i$$ μ=d1i=1∑dxi

然后计算方差：

\frac{1}{d}\sum_{i=1}^{d}(x_i-\mu)^2
$$归一化：\; xi-\mu \; \sigma 2+ϵ\; 归一化：\; \backslash frac\{x\_i-\backslash mu\}\{\backslash sqrt\{\backslash sigma^2+\backslash epsilon\}\}$$ 归一化：σ2+ϵ xi−μ

最后使用两个可训练参数：
$$yi=\gamma i\; x^i\; +\beta i\; y\_i=\backslash gamma\_i\backslash hat\{x\}\_i+\backslash beta\_i$$ yi=γix^i+βi

其中：

• $\gamma \backslash gamma$ γ：可训练缩放参数。
• $\beta \backslash beta$ β：可训练偏移参数。
• $ϵ\backslash epsilon$ ϵ：避免分母为零的小常数。

PyTorch 已经提供了实现：

nn.LayerNorm(n_embd)

我们不需要手动定义 $\gamma \backslash gamma$ γ 和 $\beta \backslash beta$ β。

4.2 LayerNorm 对哪个维度进行归一化？

假设输入为：

x.shape == (16, 64, 128)

那么：

nn.LayerNorm(128)

会针对每个样本、每个位置上的 128 维特征单独进行归一化。

也就是说：

Batch 1
  Token 1 → 对 128 维特征归一化
  Token 2 → 对 128 维特征归一化
  ...

Batch 2
  Token 1 → 对 128 维特征归一化
  ...

它不会把不同 batch 或不同 token 位置混在一起计算均值和方差。

4.3 Pre-LN 与 Post-LN

LayerNorm 可以放在不同的位置。

经典 Transformer 中的一种安排是 Post-LN：

x = layer_norm(x + attention(x))

而我们这次采用 Pre-LN：

x = x + attention(layer_norm(x))

区别在于：

• Post-LN：先进行残差相加，再归一化。
• Pre-LN：先归一化，再进入子模块，最后进行残差相加。

Pre-LN 在现代 GPT 实现中非常常见，也通常有利于较深网络的优化稳定性。

对于两个子模块，我们分别使用独立的 LayerNorm：

x = x + attention(layer_norm_1(x))
x = x + feed_forward(layer_norm_2(x))

这里不共享 LayerNorm 参数，因为注意力和 FFN 接收到的是不同阶段的特征。

---

五、把所有组件组装成 Transformer Block

现在，我们终于集齐了一个 Transformer Block 所需要的核心组件：

1. Multi-Head Attention
2. Feed Forward Network
3. Residual Connection
4. LayerNorm

接下来把它们组合起来。

5.1 Transformer Block 的完整结构

本篇采用 Pre-LN 架构：

                  输入 x
                     │
          ┌──────────┤
          │          ↓
          │       LayerNorm
          │          ↓
          │    MultiHeadAttention
          │          │
          └───────── Add
                     │
                     x
                     │
          ┌──────────┤
          │          ↓
          │       LayerNorm
          │          ↓
          │      FeedForward
          │          │
          └───────── Add
                     │
                    输出

它对应的数学表达式为：
$$h=x+MHA(\; LN1\; (x))h=x+\backslash operatorname\{MHA\}(\backslash operatorname\{LN\}\_1(x))$$ h=x+MHA(LN1(x))
$$y=h+FFN(\; LN2\; (h))y=h+\backslash operatorname\{FFN\}(\backslash operatorname\{LN\}\_2(h))$$ y=h+FFN(LN2(h))

注意，第二个 LayerNorm 处理的是经过注意力更新后的 $hh$ h，不是最初的 $xx$ x。

5.2 实现 Block

在 model.py 中新增：

class Block(nn.Module):
    """Transformer Block"""

    def __init__(self, n_embd, n_head, block_size, dropout):
        super().__init__()

        self.ln1 = nn.LayerNorm(n_embd)
        self.ln2 = nn.LayerNorm(n_embd)

        self.attention = MultiHeadAttention(
            n_embd,
            n_head,
            block_size,
            dropout
        )

        self.feed_forward = FeedForward(n_embd, dropout)

    def forward(self, x):
        x = x + self.attention(self.ln1(x))
        x = x + self.feed_forward(self.ln2(x))
        return x

虽然只有两行核心代码：

x = x + self.attention(self.ln1(x))
x = x + self.feed_forward(self.ln2(x))

但这已经实现了完整的 Transformer Block。

一个 Block 的输入和输出形状保持一致：

(B, T, n_embd)

因此，我们可以连续堆叠多个 Block。

---

六、堆叠 Transformer Block，构建 GPT

前面我们已经完成了最关键的 Transformer Block。

但为什么还要堆叠多个 Block？

因为经过一次注意力计算，模型虽然可以获取上下文信息，但这只是一次特征更新。

如果使用多层结构：

Embedding
   ↓
Block 1
   ↓
Block 2
   ↓
Block 3
   ↓
Block 4

后面的 Block 就可以在前面已经加工过的特征上，继续进行注意力计算和非线性变换。

通过反复整合上下文，模型有机会学习更加复杂的特征表示。

6.1 GPT 的完整结构

最终的小型 GPT 结构如下：

                     字符 ID
                        │
              ┌─────────┴─────────┐
              ↓                   ↓
         Token Embedding     Position Embedding
              │                   │
              └─────────┬─────────┘
                        │
                       Add
                        │
                     Dropout
                        │
                 Transformer Block
                        │
                 Transformer Block
                        │
                        ...
                        │
                 Transformer Block
                        │
                    LayerNorm
                        │
                     LM Head
                        │
               (B, T, vocab_size)
                        │
                      logits

这里有几个新名称：

• token_embedding_table：字符嵌入表。
• position_embedding_table：位置嵌入表。
• blocks：多个 Transformer Block。
• ln_f：最终的 LayerNorm。
• lm_head：将隐藏特征映射到整个字符词表的预测层。

其中，Token Embedding 和 Position Embedding 已经在上一篇实现过了。

因此，这次我们直接复用，而不是重新讲解。

6.2 完整实现 GPTLanguageModel

为了方便对照，我们保留此前的：

BigramLanguageModel
AttentionLanguageModel

并在 model.py 中新增：

GPTLanguageModel

文件顶部继续使用：

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

完整模型如下：

class GPTLanguageModel(nn.Module):
    """小型 GPT 语言模型"""

    def __init__(
        self,
        vocab_size,
        n_embd=128,
        n_head=4,
        n_layer=4,
        block_size=128,
        dropout=0.1
    ):
        super().__init__()

        self.block_size = block_size

        # 字符嵌入
        self.token_embedding_table = nn.Embedding(
            vocab_size,
            n_embd
        )

        # 位置嵌入
        self.position_embedding_table = nn.Embedding(
            block_size,
            n_embd
        )

        self.embedding_dropout = nn.Dropout(dropout)

        # 堆叠多个 Transformer Block
        self.blocks = nn.Sequential(*[
            Block(n_embd, n_head, block_size, dropout)
            for _ in range(n_layer)
        ])

        # 最终归一化
        self.ln_f = nn.LayerNorm(n_embd)

        # 语言模型预测层
        self.lm_head = nn.Linear(n_embd, vocab_size)

    def forward(self, idx, targets=None):
        B, T = idx.shape

        if T > self.block_size:
            raise ValueError("输入长度不能超过 block_size")

        # 生成位置索引
        positions = torch.arange(T, device=idx.device)

        # Token Embedding
        token_emb = self.token_embedding_table(idx)

        # Position Embedding
        position_emb = self.position_embedding_table(positions)

        # 合并字符信息与位置信息
        x = token_emb + position_emb
        x = self.embedding_dropout(x)

        # 经过多个 Transformer Block
        x = self.blocks(x)

        # 归一化与词表预测
        x = self.ln_f(x)
        logits = self.lm_head(x)

        # 计算交叉熵损失
        loss = None

        if targets is not None:
            loss = F.cross_entropy(
                logits.reshape(B * T, -1),
                targets.reshape(B * T)
            )

        return logits, loss

    @torch.no_grad()
    def generate(self, idx, max_new_tokens):
        for _ in range(max_new_tokens):

            # 只保留最近 block_size 个字符
            idx_cond = idx[:, -self.block_size:]

            # 获取模型预测
            logits, _ = self(idx_cond)

            # 只取最后一个位置
            logits = logits[:, -1, :]

            # 转换成概率
            probs = F.softmax(logits, dim=-1)

            # 采样下一个字符
            idx_next = torch.multinomial(
                probs,
                num_samples=1
            )

            # 拼接新字符
            idx = torch.cat((idx, idx_next), dim=1)

        return idx

6.3 为什么仍然返回 logits 和 loss？

观察：

return logits, loss

这与第二篇、第三篇保持一致。

当我们只进行预测时：

logits, _ = model(x)

当我们需要训练时：

logits, loss = model(x, y)

这种设计让训练程序不需要因为模型结构变化而重新编写。

无论是 Bigram、单头注意力模型，还是这次实现的 GPT，都可以沿用相同的训练调用方式。

6.4 为什么 generate 不需要重写？

还记得第二篇实现的自回归生成流程吗？

输入已有字符
     ↓
预测下一个字符
     ↓
采样一个字符
     ↓
拼接到输入末尾
     ↓
重复

整个流程没有发生变化。

变化的只是预测下一个字符的模型本身。

因此，generate() 仍然可以按照之前的方式实现。

这里需要特别注意：

idx_cond = idx[:, -self.block_size:]

因为模型的位置嵌入表最多只支持 block_size 个位置。

假设：

block_size = 64

即使我们想生成 300 个字符，每次传入模型的上下文也不能超过 64 个字符。

因此，每次生成时只取最近的 64 个字符。

这意味着生成总长度可以超过 64，但模型单次看到的上下文仍然最多只有 64 个字符。

另外，这种教学版实现每次生成都会重新计算当前上下文，并没有使用 KV Cache 等推理优化技术。

---

七、将 GPT 接入之前的训练程序

模型已经写完了，接下来就可以使用第一篇准备的数据进行训练。

这一部分不需要重新编写训练系统。

我们继续沿用：

data/
├── train.bin
├── val.bin
└── meta.pkl

以及第二篇实现的批次采样函数：

get_batch(
    split,
    datasets,
    batch_size,
    block_size,
    device
)

只需要修改训练脚本中的模型初始化部分。

7.1 导入 GPTLanguageModel

在 train.py 中修改模型导入：

from model import GPTLanguageModel

7.2 设置模型参数

第一次运行完整 GPT 时，可以先采用较小的参数配置：

model_args = {
    "vocab_size": meta["vocab_size"],
    "n_embd": 64,
    "n_head": 4,
    "n_layer": 2,
    "block_size": 64,
    "dropout": 0.1
}

这里：

参数	含义
vocab_size	字符词表大小
n_embd	字符特征维度
n_head	注意力头数量
n_layer	Transformer Block 数量
block_size	单次最大上下文长度
dropout	Dropout 概率

其中：

n_embd = 64
n_head = 4

意味着每个注意力头输出：

head_size = 64 // 4

即 16 维。

模型一共有两个 Transformer Block，每个 Block 都包含四个注意力头和一个 FFN。

这里的参数只是用于跑通实验的较小配置，不是 nanoGPT 官方训练配置。

7.3 初始化模型和优化器

model = GPTLanguageModel(**model_args).to(device)

optimizer = torch.optim.AdamW(
    model.parameters(),
    lr=1e-3
)

这里的：

**model_args

表示将字典中的键值对展开为关键字参数。

例如：

GPTLanguageModel(
    vocab_size=meta["vocab_size"],
    n_embd=64,
    n_head=4,
    n_layer=2,
    block_size=64,
    dropout=0.1
)

两种写法是等价的。

我们继续使用 AdamW 优化器。

7.4 查看模型参数量

在训练之前，可以先统计一下参数数量：

num_params = sum(
    p.numel()
    for p in model.parameters()
)

print(f"模型参数量：{num_params:,}")

p.numel() 返回一个张量所包含的元素数量。

因为模型参数本质上也是张量，所以将所有参数张量的元素数量相加，就能得到模型总参数量。

与之前的 Bigram 模型相比，新增的注意力投影、FFN 和 LayerNorm 都会带来更多参数。

模型结构更加复杂，但训练的计算成本也会增加。

7.5 运行一次训练步骤

取一个训练批次：

xb, yb = get_batch(
    "train",
    datasets,
    batch_size=16,
    block_size=model_args["block_size"],
    device=device
)

其中：

xb.shape == (16, 64)
yb.shape == (16, 64)

对应第一篇构造的错位预测目标。

接下来：

model.train()

logits, loss = model(xb, yb)

optimizer.zero_grad(set_to_none=True)
loss.backward()
optimizer.step()

print(f"Loss: {loss.item():.4f}")

这里没有新的训练机制。

仍然是：

获取训练数据
     ↓
前向传播
     ↓
计算 CrossEntropy Loss
     ↓
反向传播
     ↓
优化器更新参数

可以看到，我们虽然将模型从 Bigram 升级为多层 GPT，但训练循环的核心代码几乎没有改变。

这正是前几篇统一模型接口的好处。

完整训练时，可以继续沿用第二篇的迭代循环、训练与验证 loss 统计等逻辑。

需要注意，如果此前训练代码使用的是较小的 block_size，这里必须确保取样时使用新模型配置中的 block_size，避免训练样本长度与模型允许的最大上下文长度不一致。

7.6 使用训练后的模型生成文本

训练完成后，可以继续采用之前的方式生成文本。

model.eval()

start = torch.tensor(
    [[meta["stoi"]["\n"]]],
    dtype=torch.long,
    device=device
)

generated = model.generate(
    start,
    max_new_tokens=300
)

text = "".join(
    meta["itos"][i]
    for i in generated[0].tolist()
)

print(text)

这里仍然使用换行符作为生成起点。

model.eval() 与 @torch.no_grad() 具有不同的作用：

• model.eval()：将模型切换到评估模式，关闭 Dropout 等训练专用行为。
• @torch.no_grad()：让生成过程不记录梯度，减少不必要的计算与内存开销。

两者不能互相替代。

同样，在计算验证集 loss 时，也应该切换到评估模式，并且关闭梯度记录。

验证完成后，再调用：

model.train()

切换回训练模式。

---

八、验证我们实现的 GPT 是否正确

模型能够执行前向传播，并不意味着实现一定正确。

例如，如果因果掩码写错了，模型可能在训练时偷偷看到未来字符。

即使 loss 下降，也不代表它正确完成了自回归语言建模。

因此，我们需要编写几个基础测试。

新建文件：

test_gpt.py

8.1 检查前向传播和反向传播

首先创建一个小模型：

import torch

from model import GPTLanguageModel

torch.manual_seed(42)

model = GPTLanguageModel(
    vocab_size=10,
    n_embd=16,
    n_head=4,
    n_layer=2,
    block_size=8,
    dropout=0.1
)

随机构造一个 batch：

x = torch.randint(0, 10, (2, 8))
y = torch.randint(0, 10, (2, 8))

输入形状：

(2, 8)

即：

• 2 条序列。
• 每条序列长度为 8。

然后执行前向传播：

logits, loss = model(x, y)

assert logits.shape == (2, 8, 10)
assert loss.ndim == 0
assert torch.isfinite(loss).item()

这里：

logits.shape == (2, 8, 10)

表示每个样本、每个位置，都会预测 10 个字符对应的 logits。

然后执行反向传播：

loss.backward()

检查所有参数的梯度：

for name, parameter in model.named_parameters():
    assert parameter.grad is not None, name
    assert torch.isfinite(parameter.grad).all().item(), name

如果这些测试通过，说明当前测试输入下：

• 模型前向传播能够执行。
• Loss 能够正常计算。
• 反向传播能够计算出各参数的梯度。
• 梯度没有出现 NaN 或无穷大。

但这还不够，我们还需要检查最关键的因果性。

8.2 验证模型不会看到未来字符

因果语言模型有一个非常重要的性质：

修改未来位置的输入，不应该影响此前位置的预测结果。

例如，构造两个序列：

a = torch.tensor([[1, 2, 3, 4]])
b = torch.tensor([[1, 2, 8, 9]])

它们的前两个字符相同：

a: 1 2 3 4
b: 1 2 8 9
     ↑ ↑
   前两个相同

虽然第三、第四个位置不同，但是前两个位置的预测结果应该相同。

原因很简单：

前两个位置根本不应该看到第三、第四个位置的信息。

我们可以这样测试：

model.eval()

with torch.no_grad():
    logits_a, _ = model(a)
    logits_b, _ = model(b)

torch.testing.assert_close(
    logits_a[:, :2, :],
    logits_b[:, :2, :]
)

这里：

logits_a[:, :2, :]

表示取出所有 batch 中前两个位置对应的完整词表预测分数。

如果因果掩码正确，未来字符的变化就不会影响这些位置。

需要注意，这里必须先调用：

model.eval()

因为 Dropout 在训练模式下具有随机性。

即使输入完全相同，两次前向计算也可能得到不同输出。

这会干扰我们对因果性的判断。

8.3 检查文本生成长度

最后检查生成函数：

start = torch.tensor([[1]])

generated = model.generate(
    start,
    max_new_tokens=12
)

assert generated.shape == (1, 13)

assert (
    (generated >= 0) & (generated < 10)
).all().item()

初始输入长度为 1。

生成 12 个新字符后，总长度应该为：
$$1+12=131+12=13$$ 1+12=13

因此：

generated.shape == (1, 13)

虽然我们设置：

block_size = 8

但最终生成的序列长度仍然可以超过 8，因为每次调用模型时都只保留最近 8 个字符作为上下文。

最后添加：

print("GPT 基础检查通过")

运行：

python test_gpt.py

如果没有触发任何断言错误，就会看到：

GPT 基础检查通过

当然，这些只是最基本的功能测试，不意味着模型已经完全经过验证。

后续仍然需要在真实文本上训练，检查训练 loss、验证 loss 以及生成结果。

---

九、我们的模型已经是 GPT 了吗？

经过这一篇，我们已经拥有了一个完整的小型 GPT 主体：

Token Embedding
       +
Position Embedding
       ↓
Transformer Block × N
       ↓
Final LayerNorm
       ↓
LM Head
       ↓
Next Token Prediction

它已经具备 GPT 语言模型的几个关键特征：

• 使用因果注意力，确保每个位置只能访问当前及之前的信息。
• 使用多头注意力，从多个投影子空间汇总上下文。
• 使用前馈网络，对各个位置的特征进行非线性变换。
• 使用残差连接和 LayerNorm，改善深层网络的优化条件。
• 堆叠多个 Transformer Block，逐层更新特征表示。
• 使用交叉熵损失训练下一字符预测任务。
• 使用自回归方式逐个生成新字符。

从整体架构来看，我们已经实现了一个可以训练和生成文本的 Decoder-only Transformer 语言模型。

不过，这还不是与官方 nanoGPT 完全一致的高效实现。

例如，我们目前：

1. 为每个注意力头分别创建 Q、K、V 投影，而不是使用融合的投影计算。
2. 使用手动实现的注意力矩阵运算，还没有引入 PyTorch 的 scaled_dot_product_attention。
3. 尚未实现输入嵌入与输出预测层之间的权重共享。
4. 尚未专门调整 GPT 的参数初始化策略。
5. 生成过程中尚未使用 KV Cache 等优化机制。

这些并不会影响我们对 GPT 基本工作原理的理解。

本系列首先追求结构清晰、代码易懂、功能正确，再逐步讨论性能优化和更接近 nanoGPT 官方代码的工程实现。

---

十、总结：我们完成了什么？

回顾整个系列：

第一篇：让莎士比亚变成模型能读懂的数据

我们完成了：

原始文本
   ↓
字符词表
   ↓
编码成整数
   ↓
训练集、验证集
   ↓
构造 x 和 y

第二篇：用 Bigram 跑通第一个语言模型

我们完成了：

当前字符
   ↓
Embedding
   ↓
下一个字符的预测

第一次打通了模型训练、损失计算、反向传播和字符生成的过程。

第三篇：实现因果自注意力

我们让模型能够利用上下文：

字符嵌入 + 位置嵌入
           ↓
      单头因果注意力
           ↓
        LM Head

第四篇：搭建完整 GPT

我们把单头注意力扩展为多头注意力，并加入 FFN、残差连接、LayerNorm 和多个 Transformer Block：

字符嵌入 + 位置嵌入
           ↓
     Transformer Block
           ↓
     Transformer Block
           ↓
          ...
           ↓
       LayerNorm
           ↓
        LM Head

至此，我们已经从最初的字符编码，一步步搭建出了一个完整的小型 GPT 主体。

但目前还有一个实际问题：

训练完成后，如果关闭程序，模型的参数就会随着进程结束而消失。

我们总不能每次生成文本，都重新训练一遍模型。

因此，下一篇我们将实现模型检查点的保存与加载，并编写独立的 sample.py，让模型能够在训练结束后重新加载，继续生成文本。

同时，我们还会引入两个常见的采样参数：

• temperature：控制生成分布的随机程度。
• top_k：限制每次采样时参与选择的候选字符。

下一篇：

《从零实现 nanoGPT（五）：保存模型与文本生成，让 GPT 真正运行起来》