从零实现 nanoGPT（七）：融合 QKV，让模型实现向原版靠拢

从零实现 nanoGPT（七）：融合 QKV，让模型实现向原版靠拢

上一篇，我们完善了 GPT 的训练流程，加入验证集评估、学习率调度、梯度累积、梯度裁剪和断点续训。

到这里，我们已经实现了一个能够训练、保存和生成文本的小型 GPT。

不过，如果对比我们自己编写的模型和原版 nanoGPT，就会发现一个问题：

明明使用的都是 Transformer 架构，为什么代码的实现方式却存在不少差异？

这是因为前几篇的目标是理解原理。

例如，在第四篇中，我们把多个注意力头分别定义，再将它们的输出拼接起来。这种写法非常直观，适合学习多头注意力的工作机制。

但实际项目中，我们可以把这些操作重新组织，减少分散的矩阵运算，使计算更集中。

这一篇，我们主要对 model.py 进行四项改造：

1. 融合 QKV 投影：将多个注意力头的 Q、K、V 投影合并计算。
2. 使用 PyTorch SDPA：替换手动实现的注意力计算。
3. 共享输入输出权重：让 Token Embedding 与输出预测层使用同一组参数。
4. 改进参数初始化：统一初始化，并控制残差分支的初始权重尺度。

最后，我们还会优化生成时的输出投影，编写测试代码，验证新模型是否仍然保持正确的因果注意力行为。

本文接续之前的 GPTLanguageModel，不改变训练数据格式，也不改变基本的训练目标。

---

一、回顾之前的多头注意力实现

1.1 为什么最开始使用独立的注意力头？

第四篇中，我们首先实现了单头因果自注意力：

class Head(nn.Module):
    ...

随后通过 ModuleList 组合多个注意力头：

self.heads = nn.ModuleList([
    Head(...)
    for _ in range(n_head)
])

前向传播时，每个注意力头分别计算：
$$Qi=XWiQ\; Q\_i=XW\_i^Q$$ Qi=XWiQ
$$Ki=XWiK\; K\_i=XW\_i^K$$ Ki=XWiK
$$Vi=XWiV\; V\_i=XW\_i^V$$ Vi=XWiV

再计算：
$$headi=softmax\; (\; QiKiT\; dk\; +M)\; Vi\; head\_i=\; \backslash operatorname\{softmax\}\; \backslash left(\; \backslash frac\{Q\_iK\_i^T\}\{\backslash sqrt\{d\_k\}\}+M\; \backslash right)V\_i$$ headi=softmax(dk QiKiT+M)Vi

其中， $MM$ M 是因果遮罩。

最后，将所有注意力头的输出拼接：
$$H=Concat(head1,...,headh)H=\backslash operatorname\{Concat\}(head\_1,\backslash ldots,head\_h)$$ H=Concat(head1,...,headh)

这种方式的优点是，每个注意力头都是独立的模块，非常容易理解。

但是，它也意味着我们需要创建多个单独的线性层。

1.2 多个注意力头真的需要分别计算 QKV 吗？

假设模型参数为：

n_embd = 128
n_head = 4

head_size = n_embd // n_head

得到：

head_size = 32

每个注意力头都需要三个线性映射：

Query：128 → 32
Key：  128 → 32
Value：128 → 32

四个头合起来，一共有：

4 × 3 = 12 个线性映射

实际上，我们不需要分别调用这 12 个线性层。

以 Query 为例，可以把四个头的投影矩阵拼接为一个更大的矩阵：
$$WQ=$$W1Q,W2Q,W3Q,W4Q W^Q=W_1\^Q,W_2\^Q,W_3\^Q,W_4\^Q WQ=W1Q,W2Q,W3Q,W4Q

原来每个矩阵的形状是：
$$(128,32)(128,32)$$ (128,32)

拼接后，相应的数学矩阵形状变为：
$$(128,128)(128,128)$$ (128,128)

因此，我们只需要执行一次更大的线性映射，就能得到所有注意力头的 Query。

Key 和 Value 同样如此。

进一步，我们还可以把 Q、K、V 的投影合并为一个线性层。

---

二、融合 QKV：一次计算得到三个投影

2.1 定义融合投影层

原来的实现中，Q、K、V 分别由不同的线性层得到。

现在改为：

self.qkv = nn.Linear(
    n_embd,
    3 * n_embd,
    bias=False,
)

假设输入形状为：

(B, T, C)

其中：

• B：Batch Size。
• T：序列长度。
• C：Embedding 维度。

经过投影之后：

输入 x          (B, T, C)
                  ↓
         nn.Linear(C, 3C)
                  ↓
融合后的 QKV     (B, T, 3C)

例如：

B = 2
T = 8
C = 128

输入：

x.shape = (2, 8, 128)

经过融合投影：

qkv.shape = (2, 8, 384)

随后沿着最后一个维度拆分：

q, k, v = self.qkv(x).chunk(3, dim=-1)

得到：

q.shape = (2, 8, 128)
k.shape = (2, 8, 128)
v.shape = (2, 8, 128)

2.2 chunk(3, dim=-1) 是什么意思？

这里可以先看一个简单例子：

import torch

x = torch.arange(12).reshape(2, 6)

a, b, c = x.chunk(3, dim=-1)

print(x)
print(a)
print(b)
print(c)

原始张量：

[[ 0,  1,  2,  3,  4,  5],
 [ 6,  7,  8,  9, 10, 11]]

沿最后一个维度分成三份：

a = [[0, 1],
     [6, 7]]

b = [[2, 3],
     [8, 9]]

c = [[4, 5],
     [10, 11]]

因此：

q, k, v = self.qkv(x).chunk(3, dim=-1)

就是将最后一个维度上的 3 * n_embd 个特征，分成 Q、K、V 三组。

需要注意，融合 QKV 并不是让 Query、Key 和 Value 共享权重。

它们在融合矩阵中仍然对应不同的参数区域。

我们只是把原来分散的矩阵计算组织为一次更大的线性运算。

原版 nanoGPT 也采用这种融合投影方式。

---

三、如何把融合后的 QKV 拆成多个注意力头？

现在，Q、K、V 都已经计算完成。

但还有一个问题：

每个张量的形状都是 (B, T, C) ，怎么区分不同的注意力头？

这一步需要使用 reshape() 和 transpose()。

3.1 使用 reshape 拆分特征维度

继续使用：

B = 2
T = 8
C = 128

n_head = 4
head_size = 32

此时：

q.shape = (2, 8, 128)

我们希望将最后的 128 维拆成：

4 个注意力头 × 每头 32 维

因此执行：

q = q.reshape(B, T, n_head, head_size)

得到：

(2, 8, 4, 32)

各维度分别表示：

Batch Size
    ↓
   (2, 8, 4, 32)
        ↑  ↑  ↑
        │  │  └── 每个头的特征维度
        │  └───── 注意力头数量
        └──────── 序列长度

这里没有引入新的注意力参数，只是重新组织张量的维度。

3.2 为什么还需要 transpose？

注意力计算时，我们希望每个注意力头能够独立处理整段序列。

因此，通常把张量组织成：

(B, n_head, T, head_size)

而不是：

(B, T, n_head, head_size)

执行：

q = q.transpose(1, 2)

得到：

(2, 4, 8, 32)

这一步交换了第 1 维和第 2 维。

在这个维度布局下：

• 第 0 维对应不同样本。
• 第 1 维对应不同注意力头。
• 第 2 维对应序列中的不同位置。
• 第 3 维对应每个头内部的特征。

K 和 V 也进行相同操作：

def split_heads(tensor):
    return tensor.reshape(
        B, T, self.n_head, self.head_size
    ).transpose(1, 2)

q = split_heads(q)
k = split_heads(k)
v = split_heads(v)

最终：

Q：(B, H, T, D)
K：(B, H, T, D)
V：(B, H, T, D)

其中：
$$D=CH\; D=\backslash frac\{C\}\{H\}$$ D=HC

这种组织方式允许我们通过批量矩阵运算，同时计算所有注意力头，而不需要手动循环每一个头。

---

四、使用 scaled_dot_product_attention

4.1 回顾手动实现的注意力

在第三篇中，我们手动实现过缩放点积注意力：

scores = q @ k.transpose(-2, -1)

scores = scores / math.sqrt(head_size)

scores = scores.masked_fill(
    mask == 0,
    float("-inf"),
)

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

weights = dropout(weights)

out = weights @ 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

这段实现非常适合理解注意力机制。

但当我们真正训练模型时，可以交给 PyTorch 提供的专用接口。

4.2 SDPA 是什么？

PyTorch 提供了：

F.scaled_dot_product_attention(...)

其中，SDPA 是 Scaled Dot-Product Attention 的缩写。

它将注意力计算封装为一个函数，并可以根据设备、数据类型、张量形状等条件选择合适的计算后端。

因此，我们可以写成：

out = F.scaled_dot_product_attention(
    q,
    k,
    v,
    dropout_p=0.0,
    is_causal=True,
)

这里需要注意三个参数。

（1）Q、K、V 的形状

本篇使用：

(B, H, T, D)

这是 SDPA 支持的常见输入布局。

（2）is_causal=True

它表示使用因果注意力。

模型在位置 $tt$ t 预测时，只能利用当前位置及其之前的信息，不能读取未来位置。

因此，本篇不再需要像第三篇那样手动创建下三角遮罩。

（3）dropout_p

它控制注意力概率上的 Dropout。

这里有一个容易忽略的细节：

scaled_dot_product_attention() 不会自动根据外层模型的 train() 或 eval() 状态关闭 Dropout。

因此，我们应该写成：

dropout_p=(
    self.dropout if self.training else 0.0
)

这样在调用：

model.eval()

之后，注意力中的 Dropout 才会关闭。

4.3 SDPA 就是 FlashAttention 吗？

不是。

这里需要区分两个概念：

• SDPA：PyTorch 提供的缩放点积注意力计算接口。
• FlashAttention：一种针对注意力计算的高效实现方法。

在支持的 CUDA 环境下，SDPA 可能使用 FlashAttention 等优化后端。

但是否使用某种优化实现，取决于当前设备和运行条件。

例如，我们在 Mac 上使用：

device = "mps"

这并不意味着调用 SDPA 后就一定使用了 FlashAttention。

因此，本篇的改造目标是让代码组织更加合理，并使用 PyTorch 提供的注意力接口。

至于实际运行速度是否提升，还需要针对具体设备进行测试。

---

五、实现新的 CausalSelfAttention

现在，我们已经完成了理论准备。

接下来，将融合 QKV、注意力头拆分和 SDPA 组合起来。

在 model.py 中新增：

class CausalSelfAttention(nn.Module):
    def __init__(self, n_embd, n_head, dropout):
        super().__init__()

        if n_head <= 0 or n_embd % n_head != 0:
            raise ValueError(
                "n_head 必须为正，且能整除 n_embd"
            )

        self.n_head = n_head
        self.head_size = n_embd // n_head
        self.dropout = dropout

        # 一次计算所有头的 Q、K、V
        self.qkv = nn.Linear(
            n_embd,
            3 * n_embd,
            bias=False,
        )

        # 多头输出投影
        self.c_proj = nn.Linear(
            n_embd,
            n_embd,
        )

        self.output_dropout = nn.Dropout(dropout)

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

        # 1. 融合计算 Q、K、V
        q, k, v = self.qkv(x).chunk(3, dim=-1)

        # 2. 拆分注意力头
        def split_heads(tensor):
            return tensor.reshape(
                B, T, self.n_head, self.head_size
            ).transpose(1, 2)

        q = split_heads(q)
        k = split_heads(k)
        v = split_heads(v)

        # 3. 因果自注意力
        out = F.scaled_dot_product_attention(
            q,
            k,
            v,
            dropout_p=(
                self.dropout if self.training else 0.0
            ),
            is_causal=True,
        )

        # 4. 合并注意力头
        out = out.transpose(1, 2).contiguous()
        out = out.view(B, T, C)

        # 5. 输出投影
        out = self.c_proj(out)

        return self.output_dropout(out)

5.1 注意力计算前后的形状变化

整个过程可以总结为：

输入 x
(B, T, C)
    │
    ▼
融合 QKV 投影
(B, T, 3C)
    │
    ▼
拆分 Q、K、V
各自 (B, T, C)
    │
    ▼
拆分多个注意力头
各自 (B, H, T, D)
    │
    ▼
Scaled Dot-Product Attention
(B, H, T, D)
    │
    ▼
合并注意力头
(B, T, C)
    │
    ▼
输出线性投影
(B, T, C)

5.2 为什么需要 contiguous()？

注意这一段：

out = out.transpose(1, 2).contiguous()
out = out.view(B, T, C)

SDPA 返回的形状是：

(B, H, T, D)

但后续 Transformer Block 需要：

(B, T, C)

因此，我们首先交换维度：

out = out.transpose(1, 2)

得到：

(B, T, H, D)

然后合并最后两个维度：

out = out.view(B, T, C)

这里调用 contiguous()，是因为 transpose() 产生的张量通常不是连续内存布局。

view() 要求目标形状与张量当前的内存布局兼容。

使用：

out.contiguous()

可以得到连续布局的张量，使后续 view() 正常执行。

当然，也可以直接使用：

out = out.transpose(1, 2).reshape(B, T, C)

因为 reshape() 会在可能时返回视图，在必要时重新组织数据。

本篇保留 contiguous() + view()，以便清楚地展示维度交换和合并两个步骤。

---

六、更新 Transformer Block

第四篇中，我们实现了：

LayerNorm
    ↓
Multi-Head Attention
    ↓
Residual Connection
    ↓
LayerNorm
    ↓
FeedForward
    ↓
Residual Connection

现在，我们只需要将原来的多头注意力替换为新的 CausalSelfAttention。

class Block(nn.Module):
    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 = CausalSelfAttention(
            n_embd,
            n_head,
            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

这里仍然保留：

block_size

虽然新的注意力模块不再依赖预先生成的固定下三角遮罩，但保留这个构造参数，可以让前面创建 Block 的代码保持一致。

同时，我们依然使用第四篇的 Pre-LN 结构：
$$x\prime =x+Attention(LN(x))x\text{'}=x+\backslash operatorname\{Attention\}(\backslash operatorname\{LN\}(x))$$ x′=x+Attention(LN(x))
$$y=x\prime +FFN(LN(x\prime ))y=x\text{'}+\backslash operatorname\{FFN\}(\backslash operatorname\{LN\}(x\text{'}))$$ y=x′+FFN(LN(x′))

也就是说，本篇并没有改变 Transformer Block 的数学结构。

我们优化的是注意力内部的计算组织方式。

---

七、共享输入 Embedding 与输出层权重

除了融合 QKV，原版 nanoGPT 还使用了一项常见优化：Weight Tying（权重共享） 。

7.1 为什么输入和输出可以共享权重？

首先回顾输入嵌入层：

self.token_embedding_table = nn.Embedding(
    vocab_size,
    n_embd,
)

其权重形状为：

(vocab_size, n_embd)

假设：

vocab_size = 65
n_embd = 128

那么：

Embedding Weight = (65, 128)

输入一个字符 ID 后，Embedding 会取出对应的 128 维向量。

接着观察输出预测层：

self.lm_head = nn.Linear(
    n_embd,
    vocab_size,
    bias=False,
)

它负责把隐藏特征转换成词表中每个字符的预测分数。

对应的权重形状也是：

(vocab_size, n_embd)

可以看到，两者具有完全相同的参数形状。

因此，我们可以让它们共享同一组参数。

7.2 如何实现权重共享？

只需要执行：

self.lm_head.weight = (
    self.token_embedding_table.weight
)

这里不是复制参数的数值，而是让两个模块引用同一个 Parameter。

可以使用：

assert (
    model.lm_head.weight
    is model.token_embedding_table.weight
)

检查它们是否为同一个参数对象。

需要注意，虽然共享的是同一个权重矩阵，但输入和输出的使用方式不同。

输入端使用 Embedding 查表：
$$ei=E$$i e_i=Ei ei=Ei

输出端进行线性映射：
$$logits=hETlogits=hE^T$$ logits=hET

因此，共享的是同一份可学习参数，而不是让两个模块执行相同的运算。

7.3 权重共享节省了多少参数？

如果不共享，输入 Embedding 和输出 Linear 各有：
$$V\times CV\backslash times\; C$$ V×C

个参数。

共享后，可以减少：
$$V\times CV\backslash times\; C$$ V×C

个独立参数。

例如：

vocab_size = 65
n_embd = 128

减少：
$$65\times 128=832065\backslash times128=8320$$ 65×128=8320

个参数。

对于字符级语言模型，这个数量不算大。

但如果以后将词表扩大到：

vocab_size = 8192
n_embd = 256

那么可以减少：
$$8192\times 256=20971528192\backslash times256=2097152$$ 8192×256=2097152

个独立参数。

这已经超过 200 万。

共享权重的另一层意义是：输入表示和输出预测共同训练同一个参数矩阵。

在反向传播过程中，两条计算路径产生的梯度会累积到同一个 Parameter 上。

---

八、统一参数初始化

前几篇中，我们直接创建：

nn.Linear(...)
nn.Embedding(...)

它们会采用 PyTorch 各自的默认初始化方法。

现在，我们希望更明确地控制模型刚开始训练时的参数分布。

8.1 为什么要初始化参数？

我们已经知道，神经网络训练时会通过梯度不断调整参数。

但在训练开始之前，模型并不知道应该如何预测文本。

参数通常需要从某种随机分布中进行初始化。

初始化尺度如果不合理，可能导致前向传播的激活值或者反向传播的梯度出现不合适的大小。

因此，初始化并不是简单地给参数赋予随机值，而是需要选择合理的分布和尺度。

8.2 使用正态分布初始化

本篇使用：
$$W\sim N(0,0.022)W\backslash sim\backslash mathcal\{N\}(0,0.02^2)$$ W∼N(0,0.022)

即：

mean = 0.0
std = 0.02

在 GPTLanguageModel 中增加：

def _init_weights(self, module):
    if isinstance(module, nn.Linear):
        nn.init.normal_(
            module.weight,
            mean=0.0,
            std=0.02,
        )

        if module.bias is not None:
            nn.init.zeros_(module.bias)

    elif isinstance(module, nn.Embedding):
        nn.init.normal_(
            module.weight,
            mean=0.0,
            std=0.02,
        )

其中：

nn.init.normal_(...)

表示使用正态分布初始化参数。

而：

nn.init.zeros_(...)

表示将偏置初始化为 0。

LayerNorm 在本篇中继续使用 PyTorch 默认的初始化方式。

8.3 为什么残差分支还需要特殊初始化？

第四篇中，我们使用了残差连接：
$$xl+1\; =xl+Fl(xl)\; x\_\{l+1\}=x\_l+F\_l(x\_l)$$ xl+1=xl+Fl(xl)

如果堆叠多个 Transformer Block，每一层都会在原有特征的基础上增加一个残差分支的输出。

当模型层数较深时，我们希望控制这些残差分支在训练初期的输出尺度。

因此，可以针对每条残差分支最后一个线性投影层，使用更小的初始化标准差：
$$0.02\; 2N\; \backslash frac\{0.02\}\{\backslash sqrt\{2N\}\}$$ 2N 0.02

其中：

• $NN$ N 是 Transformer Block 数量。
• 系数 2 对应每个 Block 中的注意力和 FFN 两条残差分支。

例如：

n_layer = 2

对应：
$$0.024\; =0.01\; \backslash frac\{0.02\}\{\backslash sqrt4\}\; =0.01$$ 4 0.02=0.01

因此，我们可以写：

residual_std = (
    0.02 / math.sqrt(2 * n_layer)
)

然后对注意力输出投影和 FFN 最后一个线性层分别初始化。

for block in self.blocks:
    nn.init.normal_(
        block.attention.c_proj.weight,
        mean=0.0,
        std=residual_std,
    )

    nn.init.normal_(
        block.feed_forward.net[2].weight,
        mean=0.0,
        std=residual_std,
    )

这里沿用第四篇的 FFN 结构：

nn.Sequential(
    nn.Linear(n_embd, 4 * n_embd),  # 0
    nn.GELU(),                    # 1
    nn.Linear(4 * n_embd, n_embd), # 2
    nn.Dropout(dropout),           # 3
)

所以：

net[2]

对应第二个线性层，也就是 FFN 的输出投影。

这类初始化方法可以帮助控制初始残差输出的尺度，但不代表能够完全消除梯度爆炸或梯度消失问题。

---

九、生成时只计算最后一个位置的 Logits

之前，我们的 forward() 无论训练还是生成，都会对所有位置计算词表预测分数。

例如：

隐藏特征：
(B, T, C)

经过 lm_head：

(B, T, vocab_size)

训练时这样做是必要的。

因为我们需要预测每个位置对应的下一个字符。

但生成时并不一样。

9.1 为什么生成只需要最后一个位置？

假设当前输入：

ROMEO:

我们希望模型预测下一个字符。

虽然 Transformer 需要利用整段上下文进行计算，但真正用于下一步采样的是最后一个位置的预测结果。

因此，生成时没有必要对所有位置再执行一次词表投影。

可以修改为：

if targets is not None:
    logits = self.lm_head(x)
else:
    logits = self.lm_head(x[:, -1:, :])

这里：

x[:, -1:, :]

表示取出最后一个序列位置，并保留序列长度这一维。

它的形状为：

(B, 1, C)

经过输出层：

(B, 1, vocab_size)

相比于对全部 $TT$ T 个位置做词表投影，这样减少了生成阶段的输出投影计算量。

不过需要注意：

这并不代表 Transformer 只计算最后一个位置。

我们仍然需要对当前上下文执行 Transformer 计算。

本篇还没有实现 KV Cache，因此随着生成文本变长，模型仍然需要重复处理滑动窗口中的上下文。

---

十、完整的 model.py

前面我们分别修改了注意力模块、Block、权重共享和参数初始化。

为了方便实际操作，下面给出一份整合后的 model.py。

它保留之前 GPTLanguageModel 的构造方式，并继续提供：

model(idx, targets)

用于训练，以及：

model.generate(...)

用于文本生成。

import math

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


# =====================================
# 1. 融合 QKV 的因果自注意力
# =====================================

class CausalSelfAttention(nn.Module):
    def __init__(self, n_embd, n_head, dropout):
        super().__init__()

        if n_head <= 0 or n_embd % n_head != 0:
            raise ValueError(
                "n_head 必须为正，且能整除 n_embd"
            )

        self.n_head = n_head
        self.head_size = n_embd // n_head
        self.dropout = dropout

        self.qkv = nn.Linear(
            n_embd,
            3 * n_embd,
            bias=False,
        )

        self.c_proj = nn.Linear(n_embd, n_embd)
        self.output_dropout = nn.Dropout(dropout)

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

        q, k, v = self.qkv(x).chunk(3, dim=-1)

        def split_heads(tensor):
            return tensor.reshape(
                B, T, self.n_head, self.head_size
            ).transpose(1, 2)

        q = split_heads(q)
        k = split_heads(k)
        v = split_heads(v)

        out = F.scaled_dot_product_attention(
            q,
            k,
            v,
            dropout_p=(
                self.dropout if self.training else 0.0
            ),
            is_causal=True,
        )

        out = out.transpose(1, 2).contiguous()
        out = out.view(B, T, C)

        out = self.c_proj(out)
        return self.output_dropout(out)


# =====================================
# 2. 前馈网络
# =====================================

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)


# =====================================
# 3. Transformer Block
# =====================================

class Block(nn.Module):
    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 = CausalSelfAttention(
            n_embd,
            n_head,
            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


# =====================================
# 4. GPT 模型
# =====================================

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

        if min(vocab_size, n_embd, n_layer, block_size) <= 0:
            raise ValueError(
                "模型维度和词表大小必须大于 0"
            )

        self.block_size = block_size

        # Token Embedding
        self.token_embedding_table = nn.Embedding(
            vocab_size,
            n_embd,
        )

        # Position Embedding
        self.position_embedding_table = nn.Embedding(
            block_size,
            n_embd,
        )

        self.dropout = nn.Dropout(dropout)

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

        # Final LayerNorm
        self.ln_f = nn.LayerNorm(n_embd)

        # 输出词表投影
        self.lm_head = nn.Linear(
            n_embd,
            vocab_size,
            bias=False,
        )

        # 统一初始化
        self.apply(self._init_weights)

        # 输入输出权重共享
        self.lm_head.weight = (
            self.token_embedding_table.weight
        )

        # 残差投影缩放初始化
        residual_std = (
            0.02 / math.sqrt(2 * n_layer)
        )

        for block in self.blocks:
            nn.init.normal_(
                block.attention.c_proj.weight,
                mean=0.0,
                std=residual_std,
            )

            nn.init.normal_(
                block.feed_forward.net[2].weight,
                mean=0.0,
                std=residual_std,
            )

    def _init_weights(self, module):
        if isinstance(module, nn.Linear):
            nn.init.normal_(
                module.weight,
                mean=0.0,
                std=0.02,
            )

            if module.bias is not None:
                nn.init.zeros_(module.bias)

        elif isinstance(module, nn.Embedding):
            nn.init.normal_(
                module.weight,
                mean=0.0,
                std=0.02,
            )

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

        if T == 0 or T > self.block_size:
            raise ValueError(
                f"输入序列长度必须位于 1 到 "
                f"{self.block_size} 之间"
            )

        # Token Embedding
        tok_emb = self.token_embedding_table(idx)

        # Position Embedding
        positions = torch.arange(
            T,
            device=idx.device,
        )

        pos_emb = self.position_embedding_table(
            positions
        )

        # 合并输入特征
        x = self.dropout(tok_emb + pos_emb)

        # Transformer
        x = self.blocks(x)
        x = self.ln_f(x)

        loss = None

        if targets is not None:
            # 训练：所有位置都参与预测
            logits = self.lm_head(x)

            loss = F.cross_entropy(
                logits.reshape(B * T, -1),
                targets.reshape(B * T),
            )
        else:
            # 推理：只计算最后一个位置
            logits = self.lm_head(
                x[:, -1:, :]
            )

        return logits, loss

    @torch.no_grad()
    def generate(
        self,
        idx,
        max_new_tokens,
        temperature=1.0,
        top_k=None,
    ):
        if temperature <= 0:
            raise ValueError(
                "temperature 必须大于 0"
            )

        if top_k is not None and top_k <= 0:
            raise ValueError(
                "top_k 必须大于 0"
            )

        was_training = self.training
        self.eval()

        try:
            for _ in range(max_new_tokens):
                # 保留最近 block_size 个 Token
                idx_cond = idx[:, -self.block_size:]

                logits, _ = self(idx_cond)

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

                # Temperature
                logits = logits / temperature

                # Top-K
                if top_k is not None:
                    k = min(top_k, logits.size(-1))

                    values, _ = torch.topk(
                        logits,
                        k,
                    )

                    logits = logits.masked_fill(
                        logits < values[:, [-1]],
                        float("-inf"),
                    )

                # 计算概率
                probs = F.softmax(logits, dim=-1)

                # 随机采样
                idx_next = torch.multinomial(
                    probs,
                    num_samples=1,
                )

                # 拼接生成结果
                idx = torch.cat(
                    (idx, idx_next),
                    dim=1,
                )

        finally:
            self.train(was_training)

        return idx

这一版继续沿用之前的字符级词表，使用全部有效目标位置计算交叉熵。

需要说明的是，这里选择保持教学项目的模块命名，而不是直接复制 nanoGPT 源码中的 transformer.wte、attn.c_attn 等命名。

这样既方便理解每一步改造，也减少了对我们之前训练脚本的修改。

---

十一、验证改造后的模型是否正确

代码能够运行，并不意味着模型实现一定正确。

尤其是注意力中的维度变换和因果遮罩，如果出现错误，模型仍然可能输出形状正确的张量，但实际上读取了不应该看到的未来信息。

所以，我们还需要编写测试代码。

新建：

test_model_v2.py

写入：

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,
)

model.eval()


# -------------------------
# 1. 测试权重共享
# -------------------------

assert (
    model.lm_head.weight
    is model.token_embedding_table.weight
)


# -------------------------
# 2. 测试因果性
# -------------------------

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

targets = torch.zeros_like(a)

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

assert logits_a.shape == (1, 4, 10)

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


# -------------------------
# 3. 测试推理输出形状
# -------------------------

with torch.no_grad():
    last_logits, _ = model(a)

assert last_logits.shape == (1, 1, 10)

torch.testing.assert_close(
    last_logits,
    logits_a[:, -1:, :],
)


# -------------------------
# 4. 测试 QKV 梯度
# -------------------------

model.train()
model.zero_grad(set_to_none=True)

_, loss = model(a, targets)

loss.backward()

for block in model.blocks:
    grad = block.attention.qkv.weight.grad

    assert grad is not None
    assert torch.isfinite(grad).all().item()


# -------------------------
# 5. 测试文本生成
# -------------------------

model.eval()

generated = model.generate(
    a[:, :1],
    max_new_tokens=12,
)

assert generated.shape == (1, 13)


print("模型改造检查通过")

执行：

python test_model_v2.py

如果全部测试通过，应该看到：

模型改造检查通过

这里最重要的测试是因果性检查。

我们构造了两个输入：

a = [1, 2, 3, 4]
b = [1, 2, 8, 9]

它们前两个 Token 相同，但后两个 Token 不同。

按照因果注意力的要求：

• 第一个位置只能利用第一个 Token。
• 第二个位置只能利用前两个 Token。
• 后面的 Token 发生变化，不应该影响前两个位置的预测结果。

因此：

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

应该成立。

另外，我们使用：

model(a, targets)

取得所有位置的 Logits，而不是直接调用：

model(a)

因为本篇已经修改了推理分支，未提供 targets 时只返回最后一个位置的词表分数。

上述测试可以验证计算行为和基本结构，但不能证明模型在所有配置下都正确，也不能证明它的训练速度一定比原来更快。

如果需要比较性能，还应该在相同设备、模型规模、Batch Size 和序列长度下进行基准测试。

---

十二、使用新模型重新训练

这一篇改变了注意力参数的组织方式，并增加了输入输出权重共享。

因此，虽然模型仍然叫：

GPTLanguageModel

但它和第六篇中的版本已经不再完全一致。

12.1 为什么不能直接续训旧检查点？

之前使用独立注意力头：

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

现在则使用融合投影：

self.qkv = nn.Linear(...)

它们的参数名称和组织方式不同。

另外，我们还增加了权重共享，改变了部分参数的独立性。

因此，如果直接加载旧检查点：

model.load_state_dict(checkpoint["model"])

通常会遇到参数名称不匹配的问题。

从理论上讲，部分独立注意力头的参数可以通过适当的拼接转换为融合 QKV 的形式。

但这种迁移还需要准确处理权重排列、偏置、输出投影以及权重共享等差异。

本篇暂时不实现模型权重迁移。

我们直接开始一个新的实验。

12.2 创建新的训练目录

继续使用第六篇的 train.py：

python train.py \
  --steps 1000 \
  --out-dir out-fused

新模型会使用之前的训练配置，并在新目录中生成：

out-fused/
├── last.pt
└── best.pt

这里使用新的输出目录，而不是覆盖之前的 out-v2。

这样我们就可以保留不同版本的训练结果，方便后续对比。

12.3 继续训练

如果希望继续训练：

python train.py \
  --resume \
  --steps 1000 \
  --out-dir out-fused

程序会从新模型对应的 last.pt 中恢复模型和优化器状态。

12.4 生成文本

继续使用第五篇实现的生成脚本：

python sample.py \
  --checkpoint out-fused/best.pt \
  --prompt "ROMEO:" \
  --tokens 300 \
  --temperature 0.8 \
  --top-k 20

由于我们仍然使用字符级编码，因此不需要修改第一篇的数据预处理方式。

不过，有一点需要注意：

新的模型代码无法保证加载旧版本检查点。

因此，建议将每一阶段的代码提交到 Git，例如：

git add .
git commit -m "feat: fused QKV attention and weight tying"

这样，以后如果需要复现旧实验，就可以切换到对应的 Git 提交版本。

---

十三、本篇总结

从第三篇开始，我们一步一步搭建了 GPT 的注意力机制。

最初，我们手动计算：
$$QKT\to Mask\to Softmax\to VQK^T\backslash rightarrow\; Mask\backslash rightarrow\; Softmax\backslash rightarrow\; V$$ QKT→Mask→Softmax→V

然后通过多个独立的 Head 实现多头注意力。

这种实现帮助我们理解了 Transformer 内部的信息交互过程。

但在实际项目中，理解计算原理只是第一步。

我们还需要思考：能否在不改变核心数学结构的前提下，把计算组织得更加合理？

这一篇，我们主要完成了以下改造：

改造内容	主要作用
融合 QKV	将分散的投影组织为更大的矩阵运算
拆分多头维度	让所有注意力头可以并行计算
使用 SDPA	统一注意力计算接口，并利用适用的优化后端
输入输出权重共享	减少独立参数数量
残差投影特殊初始化	控制残差分支初始权重尺度
推理 Logits 优化	减少生成阶段不必要的词表投影

经过这些修改，我们的模型在计算组织上已经更加接近原版 nanoGPT。

但也需要明确，这些改造并不意味着所有设备上的训练性能都会显著提升。

尤其对于 Mac MPS，具体能够使用哪些底层优化，以及不同形状下的执行效率，都需要通过实际测试来判断。

到这里，我们已经从一个侧重原理的教学模型，逐步发展为结构更加完整、训练流程更加规范的字符级 GPT。

下一篇预告

《从零实现 nanoGPT（八）：从字符级编码到 BPE，让模型学会使用 Token》

到目前为止，我们一直采用最简单的字符级编码：

hello
  ↓
h、e、l、l、o

每个字符对应一个 Token。

这种方式非常适合初学者，因为编码和解码的过程容易理解。

但真实的大语言模型通常不会直接把每一个字符作为独立 Token。

下一篇，我们将引入 BPE（Byte Pair Encoding），学习如何从原始文本中构建一个子词级 Tokenizer。

主要包括：

1. 为什么字符级编码存在局限。
2. BPE 如何通过频繁出现的相邻符号对构建更长的 Token。
3. 如何训练和保存一个 Tokenizer。
4. 如何使用新的 Tokenizer 重新准备训练数据。
5. 词表大小变化后，为什么需要调整模型并重新训练。

通过下一篇，我们将从模型结构进入另一个重要主题：语言模型究竟是怎样把真实文本转换成 Token 的？