拼出一个真正的 GPT 结构

从一个零件到一台完整的机器

上一篇我们做出了注意力头，让模型能看见上下文。但一个注意力头还不是 GPT，就像一个气缸还不是发动机。

这一篇把剩下的零件补齐——多头注意力、前馈网络、残差连接、层归一化、位置编码——再把它们拼装成一个完整的 GPT 模型类。这一篇代码偏多，但每个零件都很短小，我们一个一个来。先把这一篇会用到的几个超参数列出来，它们是这台机器的尺寸规格：

n_embd = 64        # 每个字向量的长度
n_head = 4         # 注意力头的数量
n_layer = 4        # Transformer 层数
block_size = 32    # 上下文长度
dropout = 0.1      # 一个防止死记硬背的小零件，第七篇细讲

多头注意力：几个头各看各的

上一篇是一个注意力头。但语言里的关系是多种多样的——有的关系是语法上的（主语和谓语），有的是语义上的（这个代词指代前面哪个名词）。指望一个头同时盯住所有这些关系，太勉强。

办法很简单：放几个头并行跑，让它们各看各的。每个头独立地做一遍上一篇那套查询、键、值的运算，关注前文的不同侧面;跑完，把几个头的输出拼接起来，再用一个线性层混合一下。这就是多头注意力（multi-head attention）。

先把上一篇的注意力头封装成一个标准模块：

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


class Head(nn.Module):
    """单个注意力头，就是上一篇那套运算"""
    def __init__(self, head_size: int):
        super().__init__()
        self.key = nn.Linear(n_embd, head_size, bias=False)
        self.query = nn.Linear(n_embd, head_size, bias=False)
        self.value = nn.Linear(n_embd, head_size, bias=False)
        # 下三角矩阵，用于因果掩码；register_buffer 表示它不是旋钮、不参与训练
        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
        k = self.key(x)
        q = self.query(x)
        wei = q @ k.transpose(-2, -1) * k.shape[-1] ** -0.5     # 算关注度并缩放
        wei = wei.masked_fill(self.tril[:T, :T] == 0, float("-inf"))  # 屏蔽未来
        wei = F.softmax(wei, dim=-1)
        wei = self.dropout(wei)
        v = self.value(x)
        return wei @ v


class MultiHeadAttention(nn.Module):
    """多个注意力头并行，再拼接、混合"""
    def __init__(self, num_heads: int, head_size: int):
        super().__init__()
        self.heads = nn.ModuleList([Head(head_size) for _ in range(num_heads)])
        self.proj = nn.Linear(n_embd, n_embd)        # 拼接后的混合层
        self.dropout = nn.Dropout(dropout)

    def forward(self, x):
        out = torch.cat([h(x) for h in self.heads], dim=-1)   # 拼接各头输出
        return self.dropout(self.proj(out))

Head 就是上一篇代码的模块化版本。MultiHeadAttention 把若干个 Head 装进 nn.ModuleList，前向时让它们各跑一遍、torch.cat 拼起来，再过一个线性层 proj 把各头的信息融合。

前馈网络：让模型消化一下

注意力做的事，是让每个字"收集"前文的信息。但收集回来之后，还得有个环节让模型对这些信息做一番加工、思考。这个环节就是前馈网络（feed-forward network）。

它很朴素：对每个位置的向量，独立地做两次线性变换，中间夹一个非线性激活函数。结构上先把向量放大到 4 倍宽，过一道激活，再压回原来的宽度。

class FeedForward(nn.Module):
    """逐位置的小型神经网络，负责加工注意力收集来的信息"""
    def __init__(self, n_embd: int):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(n_embd, 4 * n_embd),    # 放大到 4 倍宽
            nn.ReLU(),                         # 非线性激活
            nn.Linear(4 * n_embd, n_embd),    # 压回原宽度
            nn.Dropout(dropout),
        )

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

中间那个 nn.ReLU() 是激活函数，作用是给模型引入"非线性"。如果全是线性变换，再多层叠起来效果也等同于一层，模型表达不了复杂的规律;夹一道非线性，模型才能拟合真实语言里那些拐弯抹角的关系。可以这样分工记忆：注意力负责"字与字之间互相看",前馈网络负责"每个字自己想一想"。

残差连接：给梯度留一条高速路

现在有了两个主要部件——多头注意力和前馈网络。我们想把它们叠很多层，让模型更强。但深度网络有个老问题：层一多，第四篇讲的反向传播在往回传梯度时会层层衰减，传到最前面几层时已经微弱得几乎不起作用，那些层就训不动了。

残差连接（residual connection）是解决这个问题的经典办法，而且简单到出人意料：把一个部件的输入，直接加到它的输出上。

写出来就是 x = x + 部件(x)，而不是 x = 部件(x)。

它为什么管用？因为这个"加回去"的动作，相当于在网络里开了一条直通的高速路。反向传播的梯度，除了走部件内部那条层层衰减的小路，还能顺着这条加法高速路畅通无阻地传到前面去。有了它，几十层上百层的网络才训得动。这是让深度网络成立的关键一招。

层归一化：让数值不要乱跑

还有一个稳定训练的零件：层归一化（layer normalization）。

数据穿过一层层运算时，数值的大小会飘——这一层输出的数普遍很大，下一层又普遍很小，忽大忽小会让训练变得不稳定、难收敛。层归一化做的事，是在每个部件处理之前，把每个位置的向量重新缩放到一个标准的范围（均值 0、方差 1 附近），相当于每过一道工序就把数值"扶正"一次。

PyTorch 直接提供了 nn.LayerNorm，拿来用即可。把多头注意力、前馈网络、残差连接、层归一化组合到一起，就是一个 Transformer 块（Block）——GPT 就是由一摞这样的块堆成的：

class Block(nn.Module):
    """一个 Transformer 块：注意力 + 前馈，各自带残差和层归一化"""
    def __init__(self, n_embd: int, n_head: int):
        super().__init__()
        head_size = n_embd // n_head
        self.sa = MultiHeadAttention(n_head, head_size)
        self.ffwd = FeedForward(n_embd)
        self.ln1 = nn.LayerNorm(n_embd)
        self.ln2 = nn.LayerNorm(n_embd)

    def forward(self, x):
        x = x + self.sa(self.ln1(x))      # 先层归一化，再注意力，结果残差加回
        x = x + self.ffwd(self.ln2(x))    # 先层归一化，再前馈，结果残差加回
        return x

forward 里那两行，把这一篇的零件全用上了：ln1、ln2 是层归一化，sa、ffwd 是两个主部件，而每行开头的 x = x + 就是残差连接。

位置编码：告诉模型字的先后

最后一个零件，补一个注意力机制的盲点。

回想上一篇，注意力是拿查询和键做点积来算关注度。这个运算有个特点：它只看两个字向量的内容，不看它们的先后位置。也就是说，"猫追狗"和"狗追猫",在注意力眼里字向量是同一批，它分不出顺序——可这两句话意思完全相反。

得想办法把"位置"这个信息也喂给模型。办法很巧：再建一张表，专门为位置编号。第 0 个位置一个向量，第 1 个位置一个向量……这张位置表叫位置编码（positional embedding）。然后把"字本身的向量"和"它所在位置的向量"相加，作为这个字真正进入模型的表示。这样每个字就同时带上了"我是什么字"和"我在第几个"两份信息。

把零件拼成 GPT

零件齐了，拼最终的模型。它的流程是：查字向量、查位置向量、两者相加、穿过一摞 Transformer 块、最后用一个线性层把每个位置的向量映射回 vocab_size 个打分。

class GPTLanguageModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.token_embedding_table = nn.Embedding(vocab_size, n_embd)      # 字向量表
        self.position_embedding_table = nn.Embedding(block_size, n_embd)   # 位置向量表
        self.blocks = nn.Sequential(*[Block(n_embd, n_head) for _ in range(n_layer)])
        self.ln_f = nn.LayerNorm(n_embd)               # 最后一道层归一化
        self.lm_head = nn.Linear(n_embd, vocab_size)   # 映射回 vocab_size 个打分

    def forward(self, idx, targets=None):
        B, T = idx.shape
        tok_emb = self.token_embedding_table(idx)                      # (B,T,n_embd)
        pos_emb = self.position_embedding_table(torch.arange(T))       # (T,n_embd)
        x = tok_emb + pos_emb                                          # 字 + 位置
        x = self.blocks(x)                                             # 穿过所有块
        x = self.ln_f(x)
        logits = self.lm_head(x)                                       # (B,T,vocab_size)
        if targets is None:
            loss = None
        else:
            B, T, C = logits.shape
            loss = F.cross_entropy(logits.view(B * T, C), targets.view(B * T))
        return logits, loss

    def generate(self, idx, max_new_tokens: int):
        for _ in range(max_new_tokens):
            idx_cond = idx[:, -block_size:]      # 只保留最近 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

对比一下第三篇的 Bigram，骨架惊人地像：一样有 forward 返回 logits 和 loss，一样有 generate 做接龙，一样用交叉熵算损失。变的只是中间——从"查一张表直接出打分",变成了"查字向量和位置向量、穿过一摞 Transformer 块、再出打分"。

generate 里多了一行 idx[:, -block_size:]，因为位置表只有 block_size 个位置，输入不能超过这个长度，所以生成时只把最近的 block_size 个字喂进去。

这就是一个完整的 GPT，GPT 三个字母全齐了：Generative（能生成）、Pre-trained（待会儿要训练）、Transformer（由 Transformer 块堆成）。它和 OpenAI 的 GPT 是同一套结构，差别只在尺寸——我们的 n_layer 是 4，大模型可能是几十上百;我们的 n_embd 是 64，大模型是好几千。原理一字不差。

本篇要点

• 多头注意力让几个注意力头并行，各自关注前文的不同侧面，再拼接融合。
• 前馈网络对每个位置独立做加工;注意力负责"字与字互相看",前馈负责"每个字自己想"。
• 残差连接 x = x + 部件(x) 给反向传播的梯度开了一条高速路，让深层网络训得动。
• 层归一化在每个部件前把数值缩放到标准范围，稳定训练。
• 注意力本身分不清字的先后，位置编码用一张位置向量表补上顺序信息，与字向量相加。
• 完整 GPT = 字向量加位置向量，穿过一摞 Transformer 块，再用线性层映射回打分;与真 GPT 同构，只是尺寸小。

下一篇

模型搭好了，训练循环第三篇也有了。下一篇把两者合在一起，完整训练这个迷你 GPT——设好超参数，跑起来，盯着损失下降，看它生成的文字一轮比一轮像样。

参考资料

• Attention Is All You Need（Transformer 原始论文）
• nanoGPT 项目源码
• 博客内 ml-basics 系列《Transformer》《批归一化》《Dropout 详解》