本章内容

• 开发类 GPT 大语言模型，使其通过训练可以生成类似人类语言的文本
• 对层激活进行归一化以稳定神经网络训练
• 在深度神经网络中添加快捷连接
• 实现 Transformer 块来创建不同规模的 GPT 模型
• 计算 GPT 模型的参数量和存储需求

前面我们已经学习并编写了多头注意力机制，这是大语言模型的核心组件之一。现在，我们将实现大语言模型的其他构建块，并将它们组装成一个类 GPT 模型，然后在第 5 章中训练它来生成类似人类语言的文本。

图 4-1 展示的大语言模型架构包含多个构建块。我们将先自顶向下地了解一下模型架构，再详细探讨各个部分。

4.1 构建一个大语言模型架构

大语言模型，比如 GPT（生成式预训练 Transformer），是旨在一次生成一个词（或词元）的大型深度神经网络架构。然而，尽管这些模型规模很大，但它们的架构并不像你想象的那么复杂，因为许多组件是重复的，稍后我们将详细介绍。图 4-2 提供了一个类 GPT 大语言模型的自顶向下的示意图，并突出显示了其主要组件。

前面几章已经介绍了大语言模型架构的多个方面，比如输入词元化、嵌入以及掩码多头注 力模块。本章将实现 GPT 模型的核心结构，包括其 Transformer 块，之后我们将对其进行训练以生成类似人类语言的文本。

之前，为了简化问题，我们使用了较小的嵌入维度，以便概念和示例可以在单页上展示。现在，我们将扩展到小型 GPT-2 模型的规模，即参数量为 1.24 亿的最小版本，正如 Radford 等人在 论文“Language Models are Unsupervised Multitask Learners”中所描述的。（请注意，尽管原论文中提到的参数量是 1.17 亿，但这一数量后来已被纠正。）在第 6 章中，我们将专注于将预训练权重加载到我们的实现中，并对其进行调整以适应参数量分别为 3.45 亿、7.62 亿和 15.42 亿的更大规模的 GPT-2 模型。

在深度学习和像 GPT 这样的大语言模型中，“参数”指的是模型的可训练权重。这些权重本质上是模型的内部变量，在训练过程中通过调整和优化来最小化特定的损失函数。这种优化使模型能够从训练数据中学习。

例如，在一个由 2048 维×2048 维的权重矩阵（或张量）表示的神经网络层中，矩阵中的每个元素都是一个参数。由于矩阵有 2048 行和 2048 列，因此该层的参数总数为 2048×2048， 即 4 194 304。

GPT-2 与 GPT-3
请注意，我们主要关注 GPT-2，因为 OpenAI 已经公开了该预训练模型的权重，我们将在第 6章中将这些权重加载到实现中。GPT-3在模型架构上与 GPT-2基本相同，只是将参数量从 GPT-2 的 15 亿扩展到了 1750 亿，并且训练的数据量更多了。截至本书撰写时，GPT-3 的权重尚未公开。 GPT-2 是学习实现大语言模型的更好选择，因为它能够在单台笔记本电脑上运行，而 GPT-3 需要依赖 GPU 集群来完成训练和推断。根据 Lambda 实验室的说法，在单个 V100 数据中心 GPU 上训练 GPT-3 需要 355 年，在消费级 RTX 8000 GPU 上则需要 665 年。

可以通过以下 Python 字典指定小型 GPT-2 模型的配置，稍后我们将在代码示例中使用它。

GPT_CONFIG_124M = {
    "vocab_size": 50257,    # Vocabulary size
    "context_length": 1024, # Context length
    "emb_dim": 768,         # Embedding dimension
    "n_heads": 12,          # Number of attention heads
    "n_layers": 12,         # Number of layers
    "drop_rate": 0.1,       # Dropout rate
    "qkv_bias": False       # Query-Key-Value bias
}

在 GPT_CONFIG_124M 字典中，我们使用简洁的变量名来提高代码的可读性并避免代码过长。

• vocab_size 表示会被 BPE 分词器使用的由 50 257 个单词组成的词汇表（参见第 2 章）。
• context_length 指的是模型通过位置嵌入能够处理的最大输入词元数量（参见第 2 章）。 【模型一次能容纳的最大上下文长度，单位是词元（token）】
• emb_dim 表示嵌入维度大小，可以将每个词元转化为 768 维的向量。
• n_heads 表示多头注意力机制中注意力头的数量（参见第 3 章）。
• n_layers 表示模型中的 Transformer 块数量，接下来的讨论中将介绍。
• drop_rate 表示 dropout 机制的强度（0.1 表示有 10%的隐藏单元被随机丢弃），以防止过拟合（参见第 3 章）。
• qkv_bias 指的是是否在多头注意力机制的线性层中添加一个偏置向量，用于查询、键和值的计算。遵循现代大语言模型的做法，我们会在一开始禁用它，但在第 6 章中加载 OpenAI 的预训练 GPT-2 权重时，我们会再回到这个话题（参见第 6 章）。

通过该配置，我们将实现一个 GPT 占位架构（DummyGPTModel），如图 4-3 所示。这将帮助我们从全局上理解各个部分如何协同工作，以及构建完整的 GPT 模型架构还需要编写哪些其他组件。

图 4-3 中的编号框展示了我们处理编写最终 GPT 架构所需的各个概念的顺序。我们将从第(1) 步开始，创建一个名为 DummyGPTModel 的占位符 GPT 主干部分，如代码清单 4-1 所示。

import torch
import torch.nn as nn


class DummyGPTModel(nn.Module):
    def __init__(self, cfg):
        """
        tok_emb（词嵌入）：把每个 token id 映射成一个 emb_dim 维的向量。词表大小是 vocab_size。
        pos_emb（位置嵌入）：给每个位置（最多 context_length 个）一个向量，让模型知道 token 的先后顺序。因为注意力机制本身不区分位置。
        drop_emb（Dropout）：训练时随机丢弃一部分数值，防止过拟合。
        trf_blocks：用 nn.Sequential 把 n_layers 个 Transformer 块串起来。这里用的是占位块。
        final_norm：最终的层归一化，这里也是占位。
        out_head（输出层）：把 emb_dim 维向量投影回 vocab_size 维，得到每个词的分数（logits）。bias=False 是 GPT 的常见做法。
        """
        super().__init__()
        self.tok_emb = nn.Embedding(cfg["vocab_size"], cfg["emb_dim"])
        self.pos_emb = nn.Embedding(cfg["context_length"], cfg["emb_dim"])
        self.drop_emb = nn.Dropout(cfg["drop_rate"])
        
        # Use a placeholder for TransformerBlock
        self.trf_blocks = nn.Sequential(
            *[DummyTransformerBlock(cfg) for _ in range(cfg["n_layers"])])
        
        # Use a placeholder for LayerNorm
        self.final_norm = DummyLayerNorm(cfg["emb_dim"])
        self.out_head = nn.Linear(
            cfg["emb_dim"], cfg["vocab_size"], bias=False
        )

    def forward(self, in_idx):
        """
        in_idx 形状是 (batch_size, seq_len)，是一批 token id。
        tok_embeds：查词嵌入表，形状变成 (batch, seq_len, emb_dim)。
        pos_embeds：torch.arange(seq_len) 生成 [0,1,2,...,seq_len-1]，查位置嵌入表得到 (seq_len, emb_dim)。
        tok_embeds + pos_embeds：两者相加（位置嵌入通过广播加到每个样本上），把"词义"和"位置"信息合并。
        经过 dropout、Transformer 块、归一化后，out_head 输出 logits，形状 (batch, seq_len, vocab_size)。每个位置对整个词表打分，可用来预测下一个词。
        """
        batch_size, seq_len = in_idx.shape
        tok_embeds = self.tok_emb(in_idx)
        pos_embeds = self.pos_emb(torch.arange(seq_len, device=in_idx.device))
        x = tok_embeds + pos_embeds
        x = self.drop_emb(x)
        x = self.trf_blocks(x)
        x = self.final_norm(x)
        logits = self.out_head(x)
        return logits


class DummyTransformerBlock(nn.Module):
    def __init__(self, cfg):
        super().__init__()
        # A simple placeholder

    def forward(self, x):
        # This block does nothing and just returns its input.
        return x


class DummyLayerNorm(nn.Module):
    def __init__(self, normalized_shape, eps=1e-5):
        super().__init__()
        # The parameters here are just to mimic the LayerNorm interface.

    def forward(self, x):
        # This layer does nothing and just returns its input.
        return x

在这段代码中，DummyGPTModel 类基于 PyTorch 的神经网络模块（nn.Module）定义了一个简化版的类 GPT 模型。DummyGPTModel 类中的模型架构包括词元和位置嵌入、dropout、一系列 Transformer 块（DummyTransformerBlock）、最终层归一化（DummyLayerNorm）和线性输出层（out_head）。配置信息通过一个 Python 字典（比如我们之前创建的 GPT_CONFIG_124M） 传入。

forward 方法描述了数据在模型中的处理流程：它首先计算输入索引的词元和位置嵌入，

然后应用 dropout，接着通过 Transformer 块处理数据，再应用归一化，最后使用线性输出层生成 logits。

代码清单 4-1 中的代码已经可以使用。然而，请注意，我们现在使用了占位符（DummyLayerNorm 和 DummyTransformerBlock）来作为 Transformer 块和层归一化（这些将在以后开发）的替代品。

接下来，我们将准备输入数据并初始化一个新的 GPT 模型，以展示其使用方法。基于对分词器的实现（参见第 2 章），现在我们将宏观地概述数据在 GPT 模型中的流入和流出过程，如图 4-4 所示。

为了完成这些步骤，我们将使用第 2 章中介绍的 tiktoken 分词器对包含两个文本输入的 批次进行分词处理，以供 GPT 模型使用：

import tiktoken

tokenizer = tiktoken.get_encoding("gpt2")

batch = []

txt1 = "Every effort moves you"
txt2 = "Every day holds a"

batch.append(torch.tensor(tokenizer.encode(txt1)))
batch.append(torch.tensor(tokenizer.encode(txt2)))
batch = torch.stack(batch, dim=0)
print(batch)

tensor([[6109, 3626, 6100,  345],
        [6109, 1110, 6622,  257]])

接下来，初始化一个参数量为 1.24 亿的 DummyGPTModel 实例，并将分词后的批次数据传递给它：

torch.manual_seed(123)
model = DummyGPTModel(GPT_CONFIG_124M)

logits = model(batch)
print("Output shape:", logits.shape)
print(logits)

模型的输出（通常称为 logits）如下所示：

Output shape: torch.Size([2, 4, 50257])
tensor([[[-0.9289,  0.2748, -0.7557,  ..., -1.6070,  0.2702, -0.5888],
         [-0.4476,  0.1726,  0.5354,  ..., -0.3932,  1.5285,  0.8557],
         [ 0.5680,  1.6053, -0.2155,  ...,  1.1624,  0.1380,  0.7425],
         [ 0.0447,  2.4787, -0.8843,  ...,  1.3219, -0.0864, -0.5856]],

        [[-1.5474, -0.0542, -1.0571,  ..., -1.8061, -0.4494, -0.6747],
         [-0.8422,  0.8243, -0.1098,  ..., -0.1434,  0.2079,  1.2046],
         [ 0.1355,  1.1858, -0.1453,  ...,  0.0869, -0.1590,  0.1552],
         [ 0.1666, -0.8138,  0.2307,  ...,  2.5035, -0.3055, -0.3083]]],
       grad_fn=<UnsafeViewBackward0>)

为什么logits = model(batch)运行这个后的，结果是上面的。上面的看着像是只是经过了embedding层的情况。但是实际的代码确实直接就是能生成最终结果的代码?
in_idx
→ tok_emb + pos_emb (真实层：查表，得到 (2,4,768))
→ drop_emb (真实层：dropout)
→ trf_blocks (占位：DummyTransformerBlock，原样返回 x)
→ final_norm (占位：DummyLayerNorm，原样返回 x)
→ out_head (真实层：Linear(768 → 50257))
→ logits (2,4,50257)

输出张量包含两行，对应于两个文本样本。每个文本样本由 4 个词元组成。每个词元是一个 50 257 维的向量，与分词器的词汇表大小一致。

嵌入的维度为 50 257，因为这些维度中的每一维代表词汇表中的一个唯一词元。在我们实现后处理代码时，将把这些 50 257 维的向量转换回词元 ID，然后将它们解码为单词。

在对 GPT 架构及其输入和输出有了整体了解后，我们将开始编写具体的模块。首先，实现真正的层归一化类，以此来替换之前代码中的 DummyLayerNorm。

4.2 使用层归一化进行归一化激活

由于梯度消失或梯度爆炸等问题，训练深层神经网络有时会变得具有挑战性。这些问题会导致训练过程不稳定，使网络难以有效地调整权重，从而使学习过程难以找到一组最小化损失函数的参数（权重）。换句话说，网络难以学习数据中的潜在模式，从而无法进行准确预测或决策。

---

注意
如果你对神经网络训练和梯度概念不太了解，可以参考 A.4 节，该节对此进行了简要介 绍。然而，深入理解梯度的数学知识并不是理解本书所必需的。

---

现在我们将实现层归一化，以提高神经网络训练的稳定性和效率。层归一化的主要思想是调整神经网络层的激活（输出），使其均值为 0 且方差（单位方差）为 1。这种调整有助于加速权重的有效收敛，并确保训练过程的一致性和可靠性。在 GPT-2 和当前的 Transformer 架构中，层归一化通常在多头注意力模块的前后进行。同时，正如我们在 DummyLayerNorm 占位符中所见，层归一化还应用于最终输出层之前。图 4-5 是一张可视化架构图，它展示了层归一化如何工作。

可以通过以下代码重现图 4-5 中的示例。我们实现了一个具有 5 个输入和 6 个输出的神经网络层，并将其应用于两个输入示例：

torch.manual_seed(123)

# create 2 training examples with 5 dimensions (features) each
batch_example = torch.randn(2, 5) 

layer = nn.Sequential(nn.Linear(5, 6), nn.ReLU())
out = layer(batch_example)
print(out)

nn.Sequential 把多个层按顺序打包成一个模块,前一层的输出自动作为后一层的输入。这里包含两层:
nn.Linear(5, 6):一个全连接(线性)层,输入 5 维,输出 6 维。做的运算是 y = xW^T + b,其中权重 W 形状是 (6, 5),偏置 b 是 6 维。
nn.ReLU():激活层,对上一层的输出逐元素做 max(0, x),引入非线性。

这将输出以下张量，其中第 1 行显示了第一个输入的层输出，第 2 行显示了第二个输入的层输出：

tensor([[0.2260, 0.3470, 0.0000, 0.2216, 0.0000, 0.0000],
        [0.2133, 0.2394, 0.0000, 0.5198, 0.3297, 0.0000]],
       grad_fn=<ReluBackward0>)

我们编写的神经网络层包括一个线性层和一个非线性激活函数 ReLU（修正线性单元），ReLU 是神经网络中的一种标准激活函数。如果你不熟悉 ReLU，可以这样来理解：它只是简单地将负输入值设为 0，从而确保层的输出值都是正值，这也解释了为什么结果层的输出中不包含负值。之后，我们将在 GPT 中使用一种更复杂的激活函数。

在对这些输出应用层归一化之前，先检查一下均值和方差：

mean = out.mean(dim=-1, keepdim=True)
var = out.var(dim=-1, keepdim=True)

print("Mean:\n", mean)
print("Variance:\n", var)

Mean:
 tensor([[0.1324],
        [0.2170]], grad_fn=<MeanBackward1>)
Variance:
 tensor([[0.0231],
        [0.0398]], grad_fn=<VarBackward0>)

在这里的均值张量中，第 1 行是第一个输入行的均值，第 2 行是第二个输入行的均值。

在进行均值或方差等运算时，使用 keepdim=True 可以确保输出张量与输入张量具有相同的维度，尽管这类运算是沿指定的维度 dim 减少张量的。如果没有 keepdim=True，那么返回的均值张量将是一个二维向量 [0.1324, 0.2170] ， 而不是 2×1 维的矩阵 [[0.1324], [0.2170]]。

对张量求均值/方差时，运算会沿指定维度 dim 把张量"压扁"，该维度做完运算后默认会被移除。

例如 out 形状为 [2, 3]，沿 dim=-1（最后一维）求均值：

• 不带 keepdim：out.mean(dim=-1) → 形状变为 [2]，是 1 维张量 [0.1324, 0.2170]（被运算的维度被删掉）。
• 带 keepdim=True：out.mean(dim=-1, keepdim=True) → 形状变为 [2, 1]，维度数量不变，只是该维长度变成 1。

用 keepdim=True 的意义在于保持与原张量维度对齐，方便后续广播运算。比如层归一化中的 out - mean：

• mean 形状 [2, 1]：能和 [2, 3] 正确广播，每行减去自己的均值。✓
• mean 形状 [2]：会被对齐成 [1, 2]，与 [2, 3] 不匹配，报错或算错。✗

一句话总结：keepdim=True 保留被压扁的维度（长度设为 1），让均值/方差能和原张量对齐做广

dim 参数指定了在张量中计算统计量（如均值或方差）时应该沿着哪个维度进行。 正如 图 4-6 所示，对于一个二维张量（如矩阵），使用 dim=-1 进行均值或方差计算与使用 dim=1 是等效的，因为-1 表示张量的最后一个维度，这在二维张量中对应的是列。稍后在将层归一化添加到 GPT 模型时，模型将生成形状为[batch_size, num_tokens, embedding_size]的三维张量。我们仍然可以使用 dim=-1 对最后一个维度进行归一化，以避免需要从 dim=1 转为 dim=2 的情况。

接下来，对之前得到的层输出进行层归一化操作。具体方法是减去均值，并将结果除以方差的平方根（也就是标准差）：

out_norm = (out - mean) / torch.sqrt(var)
print("Normalized layer outputs:\n", out_norm)

mean = out_norm.mean(dim=-1, keepdim=True)
var = out_norm.var(dim=-1, keepdim=True)
print("Mean:\n", mean)
print("Variance:\n", var)

Normalized layer outputs:
 tensor([[ 0.6159,  1.4126, -0.8719,  0.5872, -0.8719, -0.8719],
        [-0.0189,  0.1121, -1.0876,  1.5173,  0.5647, -1.0876]],
       grad_fn=<DivBackward0>)
Mean:
 tensor([[9.9341e-09],
        [1.9868e-08]], grad_fn=<MeanBackward1>)
Variance:
 tensor([[1.0000],
        [1.0000]], grad_fn=<VarBackward0>)

请注意，输出张量中的值-5.9605e-08 是用科学记数法来表示$-5.9605\times 10^{-8}$ ，它的十进制形式是 -0.000 000 059 605。这个值非常接近 0，但由于计算机表示数值的有限精度存在一些小的数值误差，因此不是完全为 0。【我这里的结果和书中有些差异，但核心思想一样大概率是计算机的问题】

为了提高可读性，可以通过将 sci_mode 设置为 False 来关闭科学记数法，从而在打印张量值时避免使用科学记数法：

torch.set_printoptions(sci_mode=False)
print("Mean:\n", mean)
print("Variance:\n", var)

Mean:
 tensor([[    0.0000],
        [    0.0000]], grad_fn=<MeanBackward1>)
Variance:
 tensor([[1.0000],
        [1.0000]], grad_fn=<VarBackward0>)

到目前为止，我们已经一步步地实现了层归一化。

现在，我们将把这一过程封装成一个 PyTorch 模块，以便在后续的 GPT 模型中使用，如代码清单 4-2 所示。

class LayerNorm(nn.Module):
    def __init__(self, emb_dim):
        super().__init__()
        self.eps = 1e-5
        self.scale = nn.Parameter(torch.ones(emb_dim))
        self.shift = nn.Parameter(torch.zeros(emb_dim))

    def forward(self, x):
        mean = x.mean(dim=-1, keepdim=True)
        var = x.var(dim=-1, keepdim=True, unbiased=False)
        norm_x = (x - mean) / torch.sqrt(var + self.eps)
        return self.scale * norm_x + self.shift

这个层归一化的具体实现作用在输入张量 x 的最后一个维度上， 该维度对应于嵌入维度 （emb_dim）。变量 eps 是一个小常数（epsilon），在归一化过程中会被加到方差上以防止除零错误。scale 和 shift 是两个可训练的参数（与输入维度相同），如果在训练过程中发现调整它们可以改善模型的训练任务表现，那么大语言模型会自动进行调整。这使得模型能够学习适合其数据处理的最佳缩放和偏移。

层归一化中 scale 和 shift 的作用

直觉问题：如果只是把数据变成 mean=0、var=1，确实不需要额外参数。scale 和 shift 的意义是让归一化变得可逆、可调，而不是把输出锁死在标准正态分布上。

一个具体例子：某层对一个词的输出向量（emb_dim=4）为 x = [2.0, 4.0, 6.0, 8.0]。

第 1 步 —— 归一化（把数据"拉直"）：

mean = (2+4+6+8)/4 = 5.0
var  = ((2-5)²+(4-5)²+(6-5)²+(8-5)²)/4 = (9+1+1+9)/4 = 5.0
sqrt(5.0) ≈ 2.236
norm_x = [(2-5)/2.236, (4-5)/2.236, (6-5)/2.236, (8-5)/2.236]
       = [-1.342, -0.447, 0.447, 1.342]

此时向量 mean=0、var=1。

归一化前 x:        2── 4 ──── 6 ──── 8        (中心在5, 散得开)
归一化后 norm_x:  -1.34  -0.45  0.45  1.34     (中心在0, 标准幅度)

第 2 步 —— scale 和 shift（把数据"重新摆放"），套公式 out = scale * norm_x + shift：

情况 A：初始 scale=1, shift=0

out = 1 * [-1.342, -0.447, 0.447, 1.342] + 0
    = [-1.342, -0.447, 0.447, 1.342]     ← 原封不动

训练刚开始时就是纯归一化，这两个参数不起作用。

情况 B：模型学到 scale=2, shift=0（想要"更胖"的分布）

out = 2 * [-1.342, -0.447, 0.447, 1.342] + 0
    = [-2.684, -0.894, 0.894, 2.684]     ← 拉伸变宽

scale 控制分布的"胖瘦"（幅度/方差）。

情况 C：模型学到 scale=1, shift=3（想要整体"平移"）

out = 1 * [-1.342, -0.447, 0.447, 1.342] + 3
    = [1.658, 2.553, 3.447, 4.342]        ← 整体右移到 3 附近

shift 控制分布的"位置"（均值）。

关键 —— 模型可以"撤销"归一化：原始 x 的 mean=5、sqrt(var)≈2.236。若模型学到 scale=2.236、shift=5.0：

out = 2.236 * [-1.342, -0.447, 0.447, 1.342] + 5.0
    = [-3.0, -1.0, 1.0, 3.0] + 5.0
    = [2.0, 4.0, 6.0, 8.0]                ← 完全还原成原始 x

说明模型只要愿意就能把归一化完全抵消，所以"加归一化最坏也不会更差"，模型总有退路。

流程总结：

      原始输出 x            归一化后            最终输出
    [2, 4, 6, 8]   ──►  [-1.34 ... 1.34]  ──►  scale·norm_x + shift
   (分布乱、不稳定)      (强制 mean=0,var=1)   (模型决定要什么分布)

• 归一化：负责稳定，把混乱数值拉成标准形态。
• scale/shift：负责灵活，让模型自己决定最终想要多宽、摆在哪。
• 初始化 scale=1、shift=0，训练开始时等价于纯归一化；之后由梯度自动调整。

---

有偏方差
在我们的方差计算方法中有一个实现细节，即设置 unbiased=False。这意味着在方差计算中，我们会使用样本数量 n 作为方差公式的除数。这种方法没有使用贝塞尔修正。贝塞尔修正通常在样本方差的估计中使用 n – 1 作为分母，以调整偏差。因此，这种方法得到的是所谓有偏方差估计。对于嵌入维度 n 非常大的大语言模型（如 GPT-2），使用 n 和 n – 1 的差异在实际中几乎可以忽略。我们选择这种方法是为了确保与 GPT-2 模型的归一化层兼容，并且这种方法反映了 TensorFlow 的默认行为，因为原始 GPT-2 模型是用 TensorFlow 实现的。使用相似的设置可以确保我们的方法与第 6 章中加载的预训练权重兼容。

有偏方差（unbiased=False）

代码里 var = x.var(..., unbiased=False)，指方差计算除以 n 而不是 n-1。

例：x=[2,4,6,8]，mean=5，平方差之和=20
有偏 (÷n):    20/4 = 5.0
无偏 (÷n-1):  20/3 ≈ 6.67

• 除以 n → "有偏"估计
• 除以 n-1 → "无偏"估计（贝塞尔修正）

unbiased 参数取值：

x.var(..., unbiased=False)  # 除以 n     (有偏)
x.var(..., unbiased=True)   # 除以 n-1   (无偏，贝塞尔修正)

• unbiased=True 是 PyTorch 默认值：不写参数（x.var(dim=-1, keepdim=True)）默认就是除以 n-1。所以要用 ÷n 必须显式写 unbiased=False。
• 新版 PyTorch 里 unbiased 正被 correction 替代：unbiased=False ≡ correction=0，unbiased=True ≡ correction=1，含义相同。

哪些函数有这个参数：

需要 unbiased / correction:  var, std, cov   （都和方差有关）
不需要:mean, sum, max, min ...

• 本质是"要不要做贝塞尔修正"，所以只出现在跟方差有血缘关系的函数上：std 是 var 开根号、cov 是方差的推广，都涉及"用样本估计总体离散程度"。
• mean/sum/max/min 不涉及平方距离/离散程度的估计，不存在有偏无偏问题，没有此参数。

为什么 LayerNorm 用有偏（÷n）：

• 差异可忽略：GPT-2 嵌入维度 n 很大（如 768），除以 768 还是 767 几乎没区别。
• 兼容预训练权重（关键）：原始 GPT-2 用 TensorFlow 实现，TF 算方差默认就是除以 n。设 unbiased=False 是为了和它对齐，这样加载官方预训练权重时才不会出偏差。
• 一句话：unbiased=False 即方差除以 n，理论上"有偏"但大维度下可忽略，真正目的是跟原始 GPT-2（TensorFlow）算法保持一致。

——————————

深入理解贝塞尔修正（n-1）

• 背景：常需要用样本(一小部分)去估计总体(全部数据)的真实方差。
• 问题：计算时用的是样本自己的均值 x̄，而 x̄ 天然贴近这批样本，导致算出的平方差之和系统性偏小 → 方差被低估。
• 修正：把除数从 n 改成 n-1（除数变小→结果变大）来补偿。数学上可证明除以 n-1 后估计的期望等于总体真实方差，故称"无偏"。

有偏(偏小):  Σ(x-x̄)² / n
无偏(修正):  Σ(x-x̄)² / (n-1)

• 不是每次都更准，而是长期平均不偏。
• 为什么是 n-1（自由度）：n 个数一旦固定了均值，只剩 n-1 个是自由的，最后一个可由均值反推，不算独立信息。

——————————

为什么"样本到自己均值的距离更小"（偏小的根源）

• 核心性质：样本均值 x̄ 是让 Σ(x-c)² 最小的那个点 c。换任何别的点（包括总体真实均值 μ）平方距离之和都更大。

Σ(x - x̄)²  ≤  Σ(x - μ)²
(用样本均值,最小)  (用真实均值,更大)

• 推导：令 f©=Σ(xᵢ-c)²，求导 f’©=-2Σ(xᵢ-c)=0 → c=Σxᵢ/n=x̄，即均值就是最小点。
• 数字例：x=[2,4,6]，x̄=4，设真值 μ=5

用 x̄=4: (2-4)²+(4-4)²+(6-4)² = 8
用 μ=5: (2-5)²+(4-5)²+(6-5)² = 11→ 8 < 11

• 直觉：样本均值从数据自身算出，自动落在正中间、贴合样本；真实均值是"外部标准"，样本相对它会偏一边。
• 结论：所以用样本均值算方差会系统性偏小 → 需要 ÷(n-1) 补偿。

——————————

三者串起来：LayerNorm 用 ÷n（有偏）→ 有偏是因为"用样本自身均值算方差会偏小"这一数学性质，本该用 ÷(n-1) 修正 → 但 LayerNorm 不是在估计总体，且 n 很大差异可忽略，所以用 ÷n 无妨，真正目的是兼容 GPT-2/TensorFlow。

---

现在，在实际中试用 LayerNorm 模块，并将其应用于批次输入：

ln = LayerNorm(emb_dim=6)
out_ln = ln(out)
mean = out_ln.mean(dim=-1, keepdim=True)
var = out_ln.var(dim=-1, unbiased=False, keepdim=True)

print("Mean:\n", mean)
print("Variance:\n", var)

结果表明，层归一化代码正常工作，成功地将两个输入的值归一化，使其均值为 0，方差为 1。

Mean:
 tensor([[     0.0000],
        [    -0.0000]], grad_fn=<MeanBackward1>)
Variance:
 tensor([[0.9995],
        [0.9997]], grad_fn=<VarBackward0>)

现在我们已经介绍了实现 GPT 架构所需的两个关键构建块，如图 4-7 所示。接下来，我们将研究 GELU 激活函数，它是大语言模型中使用的激活函数之一，而不是前面用过的传统 ReLU 函数。

---

层归一化与批归一化

如果熟悉批归一化（神经网络中一种常用且传统的归一化方法），你可能会想知道它与层归一化的区别。与在批次维度上进行归一化的批归一化不同，层归一化是在特征维度上进行归一化。大语言模型通常需要大量的计算资源，训练或推理时的批次大小可能会受到硬件条 件或具体用例的影响。由于层归一化是对每个输入独立进行归一化，不受批次大小的限制， 因此在这些场景中它提供了更多的灵活性和稳定性。这在分布式训练或在资源受限的环境中部署模型时尤为重要。

---

4.3 实现具有 GELU 激活函数的前馈神经网络

接下来，我们将实现一个作为大语言模型 Transformer 块一部分的小型神经网络子模块。首先要实现的是 GELU 激活函数，该函数在这个神经网络子模块中非常重要。

---

注意
有关在 PyTorch 中实现神经网络的更多信息，请参考 A.5 节。

---

从历史上看，ReLU 激活函数因其在各种神经网络架构中的简单性和有效性而被广泛应用于深度学习。然而，在大语言模型中，除了传统的 ReLU，还有其他几种激活函数，其中两个值得注意的例子是 GELU（Gaussian Error Linear Unit）和 SwiGLU（Swish-gated Linear Unit）。

GELU 和 SwiGLU 是更为复杂且平滑的激活函数，分别结合了高斯分布和 sigmoid 门控线性单元。与较为简单的 ReLU 激活函数相比，它们能够提升深度学习模型的性能。

GELU 激活函数可以通过多种方式实现，其精确的定义为 GELU$(x)=x\cdot \Phi (x)$，其中 $\Phi (x)$ 是标准高斯分布的累积分布函数。然而，在实际操作中，通常我们会使用一种计算量较小的近似实现（原始的 GPT-2 模型也是使用这种通过曲线拟合得到的近似方法进行训练的）：

$$\text{GELU}(x)\approx 0.5\cdot x\cdot \left(1+\tanh \left[\sqrt{\frac{2}{\pi }}\cdot \left(x+0.044715\cdot x^3\right)\right]\right)$$

我们可以将此函数实现为一个 PyTorch 模块，如代码清单 4-3 所示。

class GELU(nn.Module):
    def __init__(self):
        super().__init__()

    def forward(self, x):
        return 0.5 * x * (1 + torch.tanh(
            torch.sqrt(torch.tensor(2.0 / torch.pi)) * 
            (x + 0.044715 * torch.pow(x, 3))
        ))

接下来，为了直观地比较 GELU 函数与 ReLU 函数，可以将它们并排绘制出来。

import matplotlib.pyplot as plt

gelu, relu = GELU(), nn.ReLU()

# Some sample data
x = torch.linspace(-3, 3, 100)
y_gelu, y_relu = gelu(x), relu(x)

plt.figure(figsize=(8, 3))
for i, (y, label) in enumerate(zip([y_gelu, y_relu], ["GELU", "ReLU"]), 1):
    plt.subplot(1, 2, i)
    plt.plot(x, y)
    plt.title(f"{label} activation function")
    plt.xlabel("x")
    plt.ylabel(f"{label}(x)")
    plt.grid(True)

plt.tight_layout()
plt.show()

从图 4-8 的结果可以看到，ReLU（右）是一个分段线性函数，当输入为正数时直接输出输入值，否则输出 0。GELU（左）则是一个平滑的非线性函数，它近似 ReLU，但在几乎所有负值（除了在 x 约等于-0.75 的位置外）上都有非零梯度。

GELU 的平滑特性可以在训练过程中带来更好的优化效果，因为它允许模型参数进行更细微的调整。相比之下，ReLU 在零点处有一个尖锐的拐角（参见图 4-8 的右图），有时会使得优化过程更加困难，特别是在深度或复杂的网络结构中。此外，ReLU 对负输入的输出为 0，而 GELU 对负输入会输出一个小的非零值。这意味着在训练过程中，接收到负输入的神经元仍然可以参与学习，只是贡献程度不如正输入大。

接下来，如代码清单 4-4 所示，我们将使用 GELU 函数来实现小型神经网络模块 FeedForward，该模块将在大语言模型的 Transformer 块中使用。

class FeedForward(nn.Module):
    def __init__(self, cfg):
        super().__init__()
        self.layers = nn.Sequential(
            nn.Linear(cfg["emb_dim"], 4 * cfg["emb_dim"]), # 768  → 3072  放大
            GELU(),
            nn.Linear(4 * cfg["emb_dim"], cfg["emb_dim"]), # 3072 → 768   缩回
        )

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

如你所见，FeedForward 模块是一个小型神经网络，由两个线性层和一个 GELU 激活函数组成。 在参数量为 1.24 亿的 GPT 模型中，该模块通过 GPT_CONFIG_124M 字典接收输入批次，其中每个词元的嵌入维度为 768，即 GPT_CONFIG_124M[“emb_dim”] = 768。

图 4-9 展示了当输入传递给这个小型前馈神经网络时，嵌入维度是如何被操作的。

按照图 4-9 的示例，我们将创建一个词元嵌入维度为 768 的新 FeedForward 模块，然后将一个包含两个样本且每个样本有 3 个词元的批次输入提供给它：

ffn = FeedForward(GPT_CONFIG_124M)

# input shape: [batch_size, num_token, emb_size]
x = torch.rand(2, 3, 768) 
out = ffn(x)
print(out.shape)

可以看到，输出张量的形状与输入张量的形状保持一致。

torch.Size([2, 3, 768])

FeedForward 模块在提升模型学习和泛化能力方面非常关键。虽然该模块的输入和输出维度保持一致，但它通过第一个线性层将嵌入维度扩展到了更高的维度，如图 4-10 所示。扩展之后，应用非线性 GELU 激活函数，然后通过第二个线性变换将维度缩回原始大小。这种设计允许模型探索更丰富的表示空间。

为什么要先升维再降维？为什么是4呢？
下面贴出几个知乎链接来讲解这个：https://www.zhihu.com/people/xu-yu-feng-61/answers
以及一些论文，但我还没看，后续会准备看看
用户想把之前的回答转换成markdown引用格式作为笔记，去掉标题，用 > 符号来格式化。
Geva et al. 2021（EMNLP）《Transformer Feed-Forward Layers Are Key-Value Memories》这篇"FFN 是键值记忆"的论文开创了一个很活跃的研究方向，后续在各大顶会有大量跟进工作。据 Google Scholar，该论文引用量已超 1500 次（Cited by 1561），影响力极大。

直接延续这条线的后续研究：

Geva 团队自己的续作《Transformer Feed-Forward Layers Build Predictions by Promoting Concepts in the Vocabulary Space》(Geva et al., EMNLP 2022)。这是 2021 那篇的直接延伸，进一步论证 FFN 的输出是通过在词表空间"提升某些概念"来构建预测的，把 value 向量投影到词表空间做了可解释性分析。

《Empirical Study on Updating Key-Value Memories in Transformer Feed-forward Layers》(Qiu, Huang, Huang, Fu, 2024)。做消融实验，专门比较"更新 keys（FFN 第一层）"和"更新 values（FFN 第二层）"在知识编辑和微调任务上的差异。核心结论：更新 keys 比更新 values 效果更好，说明调整"模型如何控制知识"可能比直接修改知识本身更有效。代码开源。

由"键值记忆"观点催生的知识编辑方向（顶会成果最多），这是影响最大的一条支线——既然 FFN 存的是事实知识，就可以定位并直接编辑它：

ROME《Locating and Editing Factual Associations in GPT》(Meng et al., NeurIPS 2022)。用因果追踪定位事实知识存在哪些 FFN 层，然后直接改写权重。极高引用量。

MEMIT《Mass-Editing Memory in a Transformer》(Meng et al., ICLR 2023)。ROME 的规模化版本，一次编辑上千条事实。后续还有 PMET、MEND 等一系列 ICLR/ACL/NeurIPS 的知识编辑工作，全部建立在"FFN = 知识记忆"这一前提上。

由此引申的机制可解释性方向：《Analyzing Transformers in Embedding Space》(Dar et al., ACL 2023) 把参数投影到词表空间做解释，直接沿用键值记忆视角；Anthropic 的 Superposition / Toy Models 以及 Sparse Autoencoder（SAE）系列工作，研究 FFN/MLP 层中特征如何以叠加方式存储，是当前机制可解释性最前沿的方向之一。

一个补充视角《Attention is Not All You Need》证明：纯 self-attention 堆叠会导致表征秩坍缩（rank collapse），所有表征趋于同一向量，而加上 FFN（MLP）和残差后能有效缓解——从另一个角度证明了 FFN 的不可或缺。

发展脉络小结：2021 EMNLP（本文）提出 FFN = 键值记忆 → 2022 EMNLP（Geva 续作）FFN 通过词表空间提升概念做预测 → 2022 NeurIPS / 2023 ICLR（ROME、MEMIT）把理论用于定位和编辑事实知识（影响力最大的落地）→ 2023-2025（ACL、SAE 等）向机制可解释性和更细粒度功能分析扩展。

此外，输入维度和输出维度的一致性简化了架构，使我们在后续堆叠多个层时无须调整维度，从而增强了模型的扩展能力。

如图 4-11 所示，现在我们已经实现了大部分大语言模型的构建块。接下来，我们将介绍插入在神经网络不同层之间的快捷连接的概念，这对于提升深度神经网络架构的训练性能非常重要。

4.4 添加快捷连接

让我们讨论一下快捷连接（也称为“跳跃连接”或“残差连接”）的概念。快捷连接最初用于计算机视觉中的深度网络（特别是残差网络），目的是缓解梯度消失问题。梯度消失问题指的是在训练过程中，梯度在反向传播时逐渐变小，导致早期网络层难以有效训练。

如图 4-12 所示，快捷连接通过跳过一个或多个层，为梯度在网络中的流动提供了一条可替代且更短的路径。这是通过将一层的输出添加到后续层的输出中实现的。这也是为什么这种连接被称为跳跃连接。在反向传播训练中，它们在维持梯度流动方面扮演着至关重要的角色。

在代码清单 4-5 中，我们实现了图 4-12 中的神经网络，并展示了如何在前向传播过程中添加快捷连接。

class ExampleDeepNeuralNetwork(nn.Module):
    def __init__(self, layer_sizes, use_shortcut):
        super().__init__()
        self.use_shortcut = use_shortcut
        self.layers = nn.ModuleList([
            nn.Sequential(nn.Linear(layer_sizes[0], layer_sizes[1]), GELU()),
            nn.Sequential(nn.Linear(layer_sizes[1], layer_sizes[2]), GELU()),
            nn.Sequential(nn.Linear(layer_sizes[2], layer_sizes[3]), GELU()),
            nn.Sequential(nn.Linear(layer_sizes[3], layer_sizes[4]), GELU()),
            nn.Sequential(nn.Linear(layer_sizes[4], layer_sizes[5]), GELU())
        ])

    def forward(self, x):
        for layer in self.layers:
            # Compute the output of the current layer
            layer_output = layer(x)
            # Check if shortcut can be applied
            if self.use_shortcut and x.shape == layer_output.shape:
                x = x + layer_output
            else:
                x = layer_output
        return x

上述代码实现了一个具有 5 层的深度神经网络，每层由一个线性层和一个 GELU 激活函数组成。 在前向传播过程中，我们通过各层迭代地传递输入。并且，如果 self.use_shortcut 属性被设置为 True，那么我们就会选择性地添加图 4-12 中所示的快捷连接。

让我们使用这段代码初始化一个没有快捷连接的神经网络。每一层将被初始化为接受包含 3 个输入值的样本，并返回 3 个输出值。最后一层会返回一个输出值：

layer_sizes = [3, 3, 3, 3, 3, 1]  

sample_input = torch.tensor([[1., 0., -1.]])

torch.manual_seed(123)
model_without_shortcut = ExampleDeepNeuralNetwork(
    layer_sizes, use_shortcut=False
)
print_gradients(model_without_shortcut, sample_input)

接下来，实现一个用于在模型的反向传播过程中计算梯度的函数：

def print_gradients(model, x):
    # Forward pass
    output = model(x)
    target = torch.tensor([[0.]])

    # Calculate loss based on how close the target
    # and output are
    loss = nn.MSELoss()
    loss = loss(output, target)
    
    # Backward pass to calculate the gradients
    loss.backward()

    for name, param in model.named_parameters():
        if 'weight' in name:
            # Print the mean absolute gradient of the weights
            print(f"{name} has gradient mean of {param.grad.abs().mean().item()}")

这段代码定义了一个损失函数，用于计算模型输出与用户指定目标（为简化处理，这里设为 0）的接近程度。然后，当调用 loss.backward()时，PyTorch 会计算模型中每一层的损失梯度。我们可以通过 model.named_parameters()迭代权重参数。假设某一层有一个 3×3 的权重参数矩阵，那么该层将有 3×3 的梯度值。我们打印这 3×3 的梯度值的平均绝对值，以得到每一层的单一梯度值，从而更容易比较层与层之间的梯度。

简而言之，.backward()方法是 PyTorch 中的一个便捷方法，它可以在模型训练中计算所需的损失梯度，而无须我们自己手动实现复杂的梯度计算过程，这极大地简化了深度神经网络的使用。

---

注意: 如果你不太了解梯度和神经网络训练的概念，建议阅读 A.4 节和 A.7 节。

---

接下来，使用 print_gradients 函数，并将其应用于没有快捷连接的模型，输出结果如下所示：

layers.0.0.weight has gradient mean of 0.00020173584925942123
layers.1.0.weight has gradient mean of 0.00012011159560643137
layers.2.0.weight has gradient mean of 0.0007152040489017963
layers.3.0.weight has gradient mean of 0.0013988736318424344
layers.4.0.weight has gradient mean of 0.005049645435065031

print_gradients 函数的输出显示，梯度在从最后一层（layers.4）到第 1 层（layers.0） 的过程中逐渐变小，这种现象称为梯度消失问题。

现在，实例化一个包含跳跃连接的模型，并观察它的比较结果：

torch.manual_seed(123)
model_with_shortcut = ExampleDeepNeuralNetwork(
    layer_sizes, use_shortcut=True
)
print_gradients(model_with_shortcut, sample_input)

输出结果如下所示：

layers.0.0.weight has gradient mean of 0.22169792652130127
layers.1.0.weight has gradient mean of 0.20694108307361603
layers.2.0.weight has gradient mean of 0.3289699852466583
layers.3.0.weight has gradient mean of 0.2665732204914093
layers.4.0.weight has gradient mean of 1.3258541822433472

最后一层（layers.4）的梯度仍然大于其他层。然而，梯度值在逐渐接近第 1 层（layers.0）时趋于稳定，并且没有缩小到几乎消失的程度。

总之，快捷连接在解决深度神经网络中梯度消失问题的限制方面非常重要。快捷连接是诸如大语言模型等超大规模模型的核心构建块，它们将通过确保各层之间梯度的稳定流动来帮助实现更有效的训练。在第 5 章中，我们将训练 GPT 模型，其中快捷连接将发挥关键作用。

print_gradients 中的 named_parameters 与 param.grad.abs().mean().item()

先看模型长什么样

用 layer_sizes = [3, 3, 3, 3, 3, 1] 建模型，直接打印 print(model)，真实结构是：

ExampleDeepNeuralNetwork(
  (layers): ModuleList(
    (0-3): 4 x Sequential(
      (0): Linear(in_features=3, out_features=3, bias=True)
      (1): GELU()
    )
    (4): Sequential(
      (0): Linear(in_features=3, out_features=1, bias=True)
      (1): GELU()
    )
  )
)

名字里的数字都来自这棵树：

• layers = 最外层的 ModuleList
• layers 下的 (0)(1)(2)(3)(4) = 5 个 Sequential
• 每个 Sequential 下的 (0) 是 Linear，(1) 是 GELU
• 所以 layers.0.0 = ModuleList 第 0 个 → Sequential 第 0 个 = Linear(3, 3)

named_parameters() 实际输出

for name, param in model.named_parameters():
    print(name, "\t", param.shape)

layers.0.0.weight    torch.Size([3, 3])
layers.0.0.bias      torch.Size([3])
layers.1.0.weight    torch.Size([3, 3])
layers.1.0.bias      torch.Size([3])
layers.2.0.weight    torch.Size([3, 3])
layers.2.0.bias      torch.Size([3])
layers.3.0.weight    torch.Size([3, 3])
layers.3.0.bias      torch.Size([3])
layers.4.0.weight    torch.Size([1, 3])
layers.4.0.bias      torch.Size([1])

要点：

• 每个 Linear 有 weight 和 bias 两个参数，if 'weight' in name 就是只留 weight、过滤掉 bias
• 前 4 层 weight 是 [3, 3]，最后一层是 [1, 3]，跟 layer_sizes 对得上
• GELU（layers.0.1）没有可训练参数，所以不出现在列表里

param 里装的是真实张量，例如 layers.0.0.weight：

Parameter containing:
tensor([[-0.0116,  0.2894,  0.4159],
        [ 0.1153, -0.5157, -0.0993],
        [ 0.4386,  0.1937,  0.4544]], requires_grad=True)

param.grad.abs().mean().item() 逐步拆解（loss.backward() 之后，每个 param 会多出与自身同形状的 .grad）：

第 0 步，原始梯度 param.grad（3×3，有正有负）：

tensor([[ 0.00003, -0.00021,  0.00015],
        [-0.00008,  0.00042, -0.00011],
        [ 0.00019, -0.00005,  0.00007]])

第 1 步 .abs() 逐个取绝对值，形状不变：

tensor([[0.00003, 0.00021, 0.00015],
        [0.00008, 0.00042, 0.00011],
        [0.00019, 0.00005, 0.00007]])

取绝对值是为了避免正负梯度相互抵消，衡量梯度的“大小”而非“方向”。

第 2 步 .mean() 把 9 个数压成一个数（结果仍是张量）：

tensor(0.00014)

第 3 步 .item() 把单元素张量剥成纯 Python 浮点数，才能干净地拼进字符串：

0.00014

为什么要这么压缩：每层 weight 有 9 个梯度值，5 层共 45 个数无法一眼比较；用「绝对值平均」压成每层一个数，就能得到清爽的对比：

layers.0.0.weight has gradient mean of 0.00020    ← 第一层，很小
layers.2.0.weight has gradient mean of 0.00072
layers.4.0.weight has gradient mean of 0.00505    ← 最后一层，明显更大

一列数字扫下来即可看出梯度从后往前越来越小 —— 这就是梯度消失。

补充：.grad 只有在 loss.backward() 之后才有值，之前是 None，所以必须先 backward 再读 param.grad。

PyTorch 自动微分（autograd）
前向传播时，PyTorch 会在后台自动记录一张「计算图」：只要张量 requires_grad=True，用它做的每一步运算都会被记下来，连成图。loss.backward() 就是沿这张图反向走一遍，用链式法则把所有参数的梯度一次性算出来，填进各自的 .grad。不用手写求导公式。

时间线：前向之后、backward 之前，所有 .grad 都是 None（梯度还没算）；但「哪些参数会被算梯度」此时已确定 —— 就是所有 requires_grad=True 的参数。真正的梯度数值要等 backward() 执行后才有。

查看哪些参数将被计算梯度：

for name, param in model.named_parameters():
    print(f"{name:20s} requires_grad={param.requires_grad}  grad={param.grad}")

前向后的 loss 张量带 grad_fn（记录「这个张量由什么运算得来」）。注意 nn.MSELoss() 只是损失函数模块，要先算出 loss 张量才有 grad_fn。建议损失函数用 loss_fn、损失值用 loss，别混名。

loss_fn = nn.MSELoss()
loss = loss_fn(output, target)   # 先算出损失张量
print(loss.grad_fn)              # MseLossBackward0

next_functions 只往回退一步，不是整张图。要看全貌得递归遍历：

def print_graph(fn, depth=0):
    if fn is None:
        return
    print("  " * depth + str(type(fn).__name__))
    for next_fn, _ in fn.next_functions:
        print_graph(next_fn, depth + 1)

print_graph(loss.grad_fn)

图里节点名对照：

• AddmmBackward0 = Linear 层（Wx + b）的反向节点。Addmm = Add + matrix multiply，把「矩阵乘 + 加偏置」融合成一步。一个 Linear 对应一个 AddmmBackward0。（bias=False 时变 MmBackward0；3D 输入变 BmmBackward0）
• MulBackward0 / GeluBackward 等 = GELU 等激活的反向节点
• AccumulateGrad = 图的叶子节点，即参数本身，梯度最终累加到这里的 .grad
• None = 不需要梯度的输入（如 target）
• 末尾的 0 是内部版本编号，无实际含义

易踩的点：

• .grad 是累加的，训练循环每步要先 optimizer.zero_grad() 清零
• backward 后计算图默认释放，想再 backward 一次要 retain_graph=True
• 只有标量才能直接 .backward()（loss 经过求平均是标量，所以可以）
• 反向按「最后一层 → 第一层」顺序算，前层梯度靠后层结果连乘 —— 这也是梯度消失的由来

快捷连接（Shortcut / Residual Connection）的梯度为什么这么大，而且为什么最后一层到倒数第二层又一个剧烈地减小

一、要解决的问题：梯度消失
神经网络训练时，梯度从 loss 出发，反向一层层往输入方向传播。每传过一层，就要乘上"那一跳的局部导数"。在没有快捷连接的普通深层网络里，每一层的局部导数 f' = W · GELU'(...) 通常是一个小于1的小数（因为权重随机初始化、均值接近0）。小于1的数不断连乘，结果指数级衰减，越靠近输入的层梯度越小，最前面的层几乎收不到梯度、学不动。这就是梯度消失。
二、快捷连接是什么
把某一层的输入直接加回到它的输出上：普通层是 x = f(x)，加了快捷连接后变成 x = x + f(x)。相当于在这一层旁边开了一条"绕过去"的捷径，让信号（前向）和梯度（反向）都能直接通过。
三、快捷连接为什么能救梯度（核心机制）
求导时，x + f(x) 对 x 的导数 = 1 + f'。那个 +1 就是捷径的贡献。它把每一跳的因子从"一个小于1的小数 f’“变成了"1 + f’，也就是约等于1的数”。关键对比：

• 没有捷径：每跳乘一个小于1的小数 → 连乘越来越小 → 梯度消失
• 有捷径：每跳乘一个1附近的数（有时略大于1、有时略小于1）→ 连乘保持在原地 → 梯度既不消失也不爆炸，跨层稳定

注意：+1 的真正作用是"把因子拉回1附近让梯度稳定"，不是"无限放大梯度"。真实的 f’ 是很小且有正有负的数，所以真实效果是"稳定波动"而非"单调递增"。
四、结合真实实验数据理解（layer_sizes = [3,3,3,3,3,1]）
有快捷连接时各层梯度均值：

• 第5层 (layers.4): 1.32
• 第4层 (layers.3): 0.26
• 第3层 (layers.2): 0.32
• 第2层 (layers.1): 0.20
• 第1层 (layers.0): 0.22

呈现"第5层突出 → 第4层骤降 → 之后稳定"的三段式。
五、为什么第5层没有快捷连接（关键易错点）
代码里的判断条件：

if self.use_shortcut and x.shape == layer_output.shape:
    x = x + layer_output   # 形状相同才加捷径
else:
    x = layer_output       # 形状不同就不加

捷径要做 输入 + 输出，要求两者形状相同才能相加。各层情况：

• 第1~4层：3→3，输入输出都是3个数，形状相同 → 有捷径 ✅
• 第5层：3→1，要把3个数压成1个数，形状不同、无法相加 → 走 else，没有捷径 ❌

六、反向传播逐跳讲解（梯度从上往下流）
梯度传播顺序：loss → 第5层 → 第4层 → 第3层 → 第2层 → 第1层。

• 起点·第5层 = 1.32：紧挨 loss，梯度是原始值，还没被任何层乘过、没被衰减，所以大。
• 第一跳·第5层→第4层，骤降到 0.26：这一跳跨过第5层，而第5层是3→1没有捷径，因子就是裸的小数 f’≈0.2。1.32 × 0.2 ≈ 0.26。这一跳的行为和"完全没有快捷连接的网络"一样——裸乘小数就衰减。
• 之后每跳·第4→3→2→1，稳定在0.2~0.3：从第4层往下每一跳都有捷径，因子是 1 + f' ≈ 1附近。0.26×1.2≈0.32，0.32×0.63≈0.20，0.20×1.1≈0.22。因子在1上下波动，所以梯度在同一量级小幅起伏，不再往下塌。

接下来，我们会把之前讨论的所有概念（层归一化、GELU 激活函数、前馈神经网络和快捷 连接）连接在一个 Transformer 块中，这是构建 GPT 架构所需的最后一个构建块。

4.5 连接 Transformer 块中的注意力层和线性层

现在，我们将实现 Transformer 块，这是 GPT 和其他大语言模型架构的基本构建块。在参数量为 1.24 亿的 GPT-2 架构中，这个块被重复了许多次，它结合了我们之前提及的多个概念：多头注意力、层归一化、dropout、前馈层和 GELU 激活函数。稍后，我们将把这个 Transformer 块与 GPT 架构的其他部分连接起来。

图 4-13 展示了一个 Transformer 块，它结合了多个组件，包括掩码多头注意力模块（参见第 3 章）和我们之前实现的 FeedForward 模块（参见 4.3 节）。当 Transformer 块处理输入序列时，序列中的每个元素（如单词或子词）都被表示为一个固定大小的向量（此处为 768 维）。 Transformer 块内的操作，包括多头注意力和前馈层，旨在以保持这些向量维度的方式来转换它们。

Transformer 块的核心思想是，自注意力机制在多头注意力块中用于识别和分析输入序列中元素之间的关系。相比之下，前馈神经网络则在每个位置上对数据进行单独的修改。这种组合不仅提供了对输入更细致的理解和处理，而且提升了模型处理复杂数据模式的整体能力。

# If the `previous_chapters.py` file is not available locally,
# you can import it from the `llms-from-scratch` PyPI package.
# For details, see: https://github.com/rasbt/LLMs-from-scratch/tree/main/pkg
# E.g.,
# from llms_from_scratch.ch03 import MultiHeadAttention

from previous_chapters import MultiHeadAttention


class TransformerBlock(nn.Module):
    def __init__(self, cfg):
        super().__init__()
        self.att = MultiHeadAttention(
            d_in=cfg["emb_dim"],
            d_out=cfg["emb_dim"],
            context_length=cfg["context_length"],
            num_heads=cfg["n_heads"], 
            dropout=cfg["drop_rate"],
            qkv_bias=cfg["qkv_bias"])
        self.ff = FeedForward(cfg)
        self.norm1 = LayerNorm(cfg["emb_dim"])
        self.norm2 = LayerNorm(cfg["emb_dim"])
        self.drop_shortcut = nn.Dropout(cfg["drop_rate"])

    def forward(self, x):
        # Shortcut connection for attention block
        shortcut = x
        x = self.norm1(x)
        x = self.att(x)  # Shape [batch_size, num_tokens, emb_size]
        x = self.drop_shortcut(x)
        x = x + shortcut  # Add the original input back

        # Shortcut connection for feed forward block
        shortcut = x
        x = self.norm2(x)
        x = self.ff(x)
        x = self.drop_shortcut(x)
        x = x + shortcut  # Add the original input back

        return x

这段代码定义了一个 TransformerBlock 类，该类在 PyTorch 中实现了一个多头注意力机制（MultiHeadAttention）和一个前馈神经网络（FeedForward），两者都根据提供的配置字典（cfg，比如 GPT_CONFIG_124M）进行配置。

层归一化（LayerNorm）应用于这两个组件之前，而 dropout 应用于这两个组件之后，以便对模型进行正则化并防止过拟合。这种方法也被称为前层归一化（Pre-LayerNorm）。较早的架构 （如最初的 Transformer 模型）在自注意力和前馈神经网络之后才应用层归一化，这种方法被称为后层归一化（Post-LayerNorm），这通常会导致较差的训练效果。

三种 LayerNorm 结构对比

约定:下图中最左边的竖线是残差主干(identity,原样往下),右边分支是子层处理路径,两路在 + 处相加。

一、Post-LN(原始 Transformer,2017)

        x ───────┐  
        │        │
        │        ▼
        │     [子层1]
        │        │
        ▼        ▼
        + ◄──────┘
        │
        ▼
   [层归一化1]   ← LN 骑在相加之后,把残差主干截断
        │
        ▼
      输出（相当于新的 x,继续进入下一层）

数据怎么走:x 进来分两路,左边主干原样往下,右边进[子层1],两路在 + 相加。相加后,整个结果被送进[层归一化1]。

公式:

输出 = 层归一化( x + 子层(x) )
              └────┬──────┘
              整个被包进归一化里

后果:没有一条干净的残差高速路——梯度每往回传一层都要穿过一次 LayerNorm。堆深后靠近输出的梯度量级随深度变大,训练容易不稳,必须靠 warmup(开始用很小的学习率慢慢加)保护。这是它的硬伤。

二、Pre-LN(GPT-2 采用)✅

        x ───────┐  
        │        │
        │        ▼
        │   [层归一化1]
        │        │
        │        ▼
        │     [子层1]
        │        │
        ▼        ▼
        + ◄──────┘
        │
        ▼
      输出

数据怎么走:x 进来分两路,左边主干原样往下,右边先进[层归一化1]、再进[子层1],最后在 + 相加。

关键区别:层归一化被挪进右边支路的开头。滑最左边主干竖线,从 x 到 + 中间不撞任何方框,是一根干净直线。

公式:

输出 = x + 子层( 层归一化(x) )
       ↑       └────┬─────┘
      原样      归一化只作用在支路上

好处:梯度沿干净主干往回传,不被 LayerNorm 缩放,像高速公路一样从最后一层一滑到底(梯度畅通)。因此堆 12 层还是 96 层梯度都稳(可堆很深),对学习率更鲁棒,大幅降低对 warmup 的依赖(实践中通常仍会配合学习率调度使用)。这是它成为现代大模型标准的原因。

小尾巴:主干全程没被归一化过,所以整个模型末尾要再补一个最终层归一化(final LN),把输出归一化后再送去预测。

三、变体(层归一化放在子层之后)

        x ───────┐  主干(残差)
        │        │
        │        ▼
        │     [子层1]
        │        │
        │        ▼
        │   [层归一化1]
        │        │
        ▼        ▼
        + ◄──────┘
        │
        ▼
      输出

数据怎么走:x 进来分两路,左边主干原样往下,右边先进[子层1]、再进[层归一化1],最后在 + 相加。

对了一半:滑最左边主干竖线,从 x 到 + 不撞方框,主干干净。所以"梯度畅通"这点和 Pre-LN 一样好,比 Post-LN 强。

公式:

输出 = x + 层归一化( 子层(x) )
              ↑           ↑
        归一化放在支路末尾   子层拿到没归一化的 x

缺点:

1. (硬伤)子层输入没被归一化。层归一化本该保证送进子层的数据分布稳定,但这里子层先处理原始 x。堆很多层后送进子层的数据尺度越漂越远,该解决的问题没解决。
2. (倾向,非绝对)可能扰乱相加的尺度匹配。残差相加本意是 x + 一个小修正,修正应较小以保留主信号。LN 会把子层输出拉向均值 0、方差 1,尺度可能和原始 x 对不上。不过 LN 带可学习的仿射参数 γ、β,网络有能力把尺度重新学回去,所以这条是倾向性风险,没有第 1 点那么致命。

结果:主干干净这个优点它有,但子层输入没归一化,两头不讨好,不如 Pre-LN。

TransformerBlock 类还实现了前向传播，其中每个组件后面都跟着一个快捷连接，将块 的输入加到其输出上。这个关键特性有助于在训练过程中使梯度在网络中流动，并改善深度模型的学习效果（参见 4.4 节）。

使用我们之前定义的 GPT_CONFIG_124M 配置字典来实例化一个 Transformer 块，并输入一些测试数据：

torch.manual_seed(123)

x = torch.rand(2, 4, 768)  # Shape: [batch_size, num_tokens, emb_dim]
block = TransformerBlock(GPT_CONFIG_124M)
output = block(x)

print("Input shape:", x.shape)
print("Output shape:", output.shape)

Input shape: torch.Size([2, 4, 768])
Output shape: torch.Size([2, 4, 768])

额外提下x.shape，因为在multiheadattention中，b, num_tokens, d_in = x.shape。这里num_tokens（当前序列长度），是这一次真正传进来的序列有多少个 token，随每个 batch 变化。context_length（上下文长度 / 最大长度），模型出厂时就定好最多能处理多少个 token。

如你所见，Transformer 块在输出中维持了输入的维度，这表明 Transformer 架构在处理数据序列时不会改变它们在网络中的形状。

在整个 Transfromer 块架构中保持形状不变并非偶然，而是其设计的一个重要方面。这种设计使其能有效应用于各种序列到序列的任务，其中每个输出向量直接对应一个输入向量，保持一一对应的关系。然而，正如我们在第 3 章中所学到的，输出是一个包含整个输入序列信息的上下文向量。这意味着，虽然序列的物理维度（长度和特征尺寸）在通过 Transformer 块时保持不变，但每个输出向量的内容都要重新编码，以整合来自整个输入序列的上下文信息。

在实现了 Transformer 块之后，我们现在具备了构建 GPT 架构所需的所有构建块。如图 4-14 所示，Transformer 块结合了层归一化、前馈神经网络、GELU 激活函数和快捷连接。正如我们最终会看到的，这个 Transformer 块将构成 GPT 架构的核心组件。

4.6 实现 GPT 模型

本章开始时，我们对 GPT 架构进行了概览，并将其称为 DummyGPTModel。虽然 DummyGPTModel 的代码实现中展示了 GPT 模型的输入和输出，但其包含的构建块仍然是一个“黑盒”，因此我们使用了 DummyTransformerBlock 类和 DummyLayerNorm 类作为占位符。

在本节中，我们将把 DummyTransformerBlock 占位符和 DummyLayerNorm 占位符替换为之前编写的真正的 TransformerBlock 类和 LayerNorm 类，从而组装出一个完全可用且参数量为 1.24 亿的 GPT-2 模型。在第 5 章中，我们将对 GPT-2 模型进行预训练。在第 6 章中，我们将从 OpenAI 加载预训练的权重。

在用代码构建 GPT-2 模型之前，先来看看它的整体结构，如图 4-15 所示，其中包括了我们到目前为止讨论的所有概念。可以看到，Transformer 块在 GPT 模型架构中被多次重复。在参数量为 1.24 亿的 GPT-2 模型中，这个 Transformer 块被重复使用了 12 次，这可以通过 GPT_CONFIG_124M 字典中的 n_layers 条目进行指定。在最大规模的 GPT-2 模型（参数量为 15.42 亿）中，这个 Transformer 块被重复了 48 次。

疑问一:GPT 架构图里为什么好多步骤后面都跟着 dropout?为什么用得这么频繁?嵌入层之后也放了 dropout,可 embedding 不就是个查表吗,它有可学习的变化吗?

看起来"到处都是 dropout",其实真正的 dropout 只有三个位置:嵌入层之后一个、每个 Transformer 块内部两个(注意力之后、前馈层之后)。之所以显得多,是因为中间的 Transformer 块要重复 12 次,块内那两个 dropout 就跟着出现了 24 次。这是同一个结构被画了 12 遍的视觉效果,不是刻意堆了很多不同的 dropout。

dropout 的作用是防止过拟合:训练时随机把一部分神经元的输出置零,强迫网络不要过度依赖某几个特定神经元,从而学到更鲁棒的表示。它只在训练时生效,推理(实际生成文本)时会被关闭,所以不影响最终采样。GPT-2 里 dropout 率通常设得很小(如 0.1),甚至有些实现设为 0,并不是激进操作。

关于嵌入层之后为什么也放 dropout:embedding 确实是查表,而且这张表(词元嵌入矩阵、位置嵌入矩阵)确实是可训练的参数。但要修正一个因果误解——dropout 放这里不是因为这张表可训练。dropout 正则化的对象不是权重,而是流过网络的激活值(向量)。嵌入层的输出是后面 12 层 Transformer 的输入起点,如果不加 dropout,模型每次看到的输入向量都完全固定;加了之后,每次前向传播都随机遮掉输入向量的一部分维度,相当于给输入加噪声,迫使后续网络在"输入被部分破坏"时也能工作,从而更鲁棒。所以 dropout 作用在数据流上,不作用在权重上,和"前面是不是可学习变换"没有直接绑定关系。

(补充:真正直接作用在可训练参数上、防止参数过拟合的手段叫权重衰减 / L2 正则化,和 dropout 是两回事。)

---

疑问二:最终输出为什么是 4 × 50257?词表里只有 4 个单词吗?我理解应该覆盖全部词表、每个词都有概率,这样才能采样。另外,训练时是不是不止学了"Every effort moves you → forward"这一个模式,而是把输入本身也拆成了多个子模式一起训练?

先拆开这两个维度。50257 正是整个词表的大小(GPT-2 的 BPE 词表就是 50257 个词元)。所以模型输出的不是 4 个单词,而是对词表里全部 50257 个词元都给出一个分数(logits),经过 softmax 变成 50257 个概率,从中采样下一个词——这正符合"应该覆盖全部词表、每个词都有概率"的直觉。

4 对应的是输入序列的长度。“Every effort moves you” 被分成 4 个词元,模型对每个位置都预测了"下一个词元的分布":

• 第 1 行:看到 “Every” 后,预测下一个词
• 第 2 行:看到 “Every effort” 后,预测下一个词
• 第 3 行:看到 “Every effort moves” 后,预测下一个词
• 第 4 行:看到 “Every effort moves you” 后,预测下一个词

所以张量形状是 [序列长度 4, 词表 50257]。生成文本时只关心最后一行(第 4 行),因为它是"看完全部输入后"对真正下一个词的预测,这里预测出来的就是 “forward”。

关于训练:理解完全正确。喂进 “Every effort moves you” 时,模型不是只学"看完四个词后输出 forward"这一个模式,而是同时在学:

• Every → effort
• Every effort → moves
• Every effort moves → you
• Every effort moves you → forward

这就是 4 × 50257 张量的意义:4 行对应 4 个位置,每行都和该位置真正的下一个词算损失,四个损失一起反向传播,所以一句话能提供 4 条训练信号,不浪费。

这种做法叫自监督(self-supervised),也就是 causal language modeling。它不需要人工标注,文本本身就是标签——每个词的标准答案就是它后面那个真实存在的词。输入和标签只是错开一位的关系:输入是 [Every, effort, moves, you],目标是 [effort, moves, you, forward]。这也是大语言模型能用海量无标注文本训练的原因:任意一段文字都能自动切成无数个"前文 → 下一个词"的样本。

最终 Transformer 块的输出会经过最后一步的层归一化处理，然后传递到线性输出层。这个层会将 Transformer 的输出映射到一个高维空间（在本例中为 50 257 维，对应模型的词汇表大小），以预测序列中的下一个词元。

现在，让我们来实现图 4-15 中的架构，如代码清单 4-7 所示。

class GPTModel(nn.Module):
    def __init__(self, cfg):
        super().__init__()
        self.tok_emb = nn.Embedding(cfg["vocab_size"], cfg["emb_dim"])
        self.pos_emb = nn.Embedding(cfg["context_length"], cfg["emb_dim"])
        self.drop_emb = nn.Dropout(cfg["drop_rate"])
        
        self.trf_blocks = nn.Sequential(
            *[TransformerBlock(cfg) for _ in range(cfg["n_layers"])])
        
        self.final_norm = LayerNorm(cfg["emb_dim"])
        self.out_head = nn.Linear(
            cfg["emb_dim"], cfg["vocab_size"], bias=False
        )

    def forward(self, in_idx):
        batch_size, seq_len = in_idx.shape
        tok_embeds = self.tok_emb(in_idx)
        pos_embeds = self.pos_emb(torch.arange(seq_len, device=in_idx.device))
        x = tok_embeds + pos_embeds  # Shape [batch_size, num_tokens, emb_size]
        x = self.drop_emb(x)
        x = self.trf_blocks(x)
        x = self.final_norm(x)
        logits = self.out_head(x)
        return logits

得益于 TransformerBlock 类，GPTModel 类显得相对小巧且紧凑。

这个 GPTModel 类的__init__构造函数通过 Python 字典 cfg 传递的配置来初始化词元嵌入层和位置嵌入层。这些嵌入层负责将输入的词元索引转换为稠密向量，并添加位置信息（参见第 2 章）。

接下来，__init__方法会创建一个 TransformerBlock 模块的顺序栈，其层数与 cfg 中指定的层数相同。Transformer 块之后会应用一个 LayerNorm 层，将 Transformer 块的输出标准化，以稳定学习过程。最后，定义一个无偏置的线性输出头，将 Transformer 的输出投射到分词 器的词汇空间，为词汇中的每个词元生成分数（logits）。

foward 方法首先接收一批输入的词元索引，然后计算它们的嵌入表示，接着应用位置嵌入，将序列通过一系列 Transformer 块传递，并对最终输出进行归一化处理。最后一步是计算 logits，这些 logits 代表了下一个词元的非归一化概率。我们将在 4.7 节中把这些 logits 转换为词元和文本输出。

现在，使用传入了 cfg 参数的 GPT_CONFIG_124M 字典初始化参数量为 1.24 亿的 GPT 模型，并向其输入我们之前创建的批次文本数据：

torch.manual_seed(123)
model = GPTModel(GPT_CONFIG_124M)

out = model(batch)
# txt1 = "Every effort moves you"
# txt2 = "Every day holds a"
print("Input batch:\n", batch)
print("\nOutput shape:", out.shape)
print(out)

这段代码先打印输入批次的内容，然后打印输出张量：

Input batch:
 tensor([[6109, 3626, 6100,  345],
        [6109, 1110, 6622,  257]])

Output shape: torch.Size([2, 4, 50257])
tensor([[[ 0.1381,  0.0077, -0.1963,  ..., -0.0222, -0.1060,  0.1717],
         [ 0.3865, -0.8408, -0.6564,  ..., -0.5163,  0.2369, -0.3357],
         [ 0.6989, -0.1829, -0.1631,  ...,  0.1472, -0.6504, -0.0056],
         [-0.4290,  0.1669, -0.1258,  ...,  1.1579,  0.5303, -0.5549]],

        [[ 0.1094, -0.2894, -0.1467,  ..., -0.0557,  0.2911, -0.2824],
         [ 0.0882, -0.3552, -0.3527,  ...,  1.2930,  0.0053,  0.1898],
         [ 0.6091,  0.4702, -0.4094,  ...,  0.7688,  0.3787, -0.1974],
         [-0.0612, -0.0737,  0.4751,  ...,  1.2463, -0.3834,  0.0609]]],
       grad_fn=<UnsafeViewBackward0>)

可以看到，输出张量的形状为[2, 4, 50257]，因为传入了两个输入文本，每个文本有 4 个词元。最后一个维度 50257 相当于分词器的词汇量。在 4.7 节中，你将看到如何将这 50 257 维输出向量逐一转换回词元。

在继续编写将模型输出转换为文本的函数之前，让我们花一些时间了解一下模型架构及其规模。通过 numel()（“number of elements”的缩写）方法可以统计模型参数张量的总参数量：

total_params = sum(p.numel() for p in model.parameters())
print(f"Total number of parameters: {total_params:,}")

Total number of parameters: 163,009,536

现在，如果仔细看，你可能会注意到一个差异。前面我们提到要初始化一个参数量为 1.24 亿的 GPT 模型，那为什么上面代码实际输出的参数量是 1.63 亿呢？

原因在于原始 GPT-2 架构中使用了一个叫作权重共享（weight tying）的概念。也就是说，原始 GPT-2 架构是将词元嵌入层作为输出层重复使用的。为了弄清楚这意味着什么，让我们来看看前面通过 GPTModel 在模型上初始化的词元嵌入层和线性输出层的形状：

print("Token embedding layer shape:", model.tok_emb.weight.shape)
print("Output layer shape:", model.out_head.weight.shape)

从打印输出可以看出，这两个层的权重张量具有相同的形状：

Token embedding layer shape: torch.Size([50257, 768])
Output layer shape: torch.Size([50257, 768])

由于分词器词汇表中有 50 257 个条目，因此词元嵌入层和输出层非常庞大。根据权重共享的概念，我们需要从总的 GPT-2 模型参数计数中减去输出层的参数量：

total_params_gpt2 =  total_params - sum(p.numel() for p in model.out_head.parameters())
print(f"Number of trainable parameters considering weight tying: {total_params_gpt2:,}")

Number of trainable parameters considering weight tying: 124,412,160

如你所见，现在模型只有 1.24 亿个参数，与 GPT-2 模型的原始大小相当。

权重共享可以减少模型的总体内存占用和计算复杂度。不过，根据我的经验，使用单独的词元嵌入层和输出层可以获得更好的训练效果和模型性能。因此，我们在 GPTModel 实现中使用了单独的层。这一点在现代大语言模型中同样适用。不过，我们将在第 6 章中再次回顾并实现权重共享的概念，到时会加载来自 OpenAI 的预训练权重。

---

练习 4.1 前馈模块和注意模块的参数量
计算前馈模块和注意力模块所包含的参数量，并进行对比。

from gpt import TransformerBlock

GPT_CONFIG_124M = {
    "vocab_size": 50257,
    "context_length": 1024,
    "emb_dim": 768,
    "n_heads": 12,
    "n_layers": 12,
    "drop_rate": 0.1,
    "qkv_bias": False
}

block = TransformerBlock(GPT_CONFIG_124M)
print(block)

TransformerBlock(
  (att): MultiHeadAttention(
    (W_query): Linear(in_features=768, out_features=768, bias=False)
    (W_key): Linear(in_features=768, out_features=768, bias=False)
    (W_value): Linear(in_features=768, out_features=768, bias=False)
    (out_proj): Linear(in_features=768, out_features=768, bias=True)
    (dropout): Dropout(p=0.1, inplace=False)
  )
  (ff): FeedForward(
    (layers): Sequential(
      (0): Linear(in_features=768, out_features=3072, bias=True)
      (1): GELU()
      (2): Linear(in_features=3072, out_features=768, bias=True)
    )
  )
  (norm1): LayerNorm()
  (norm2): LayerNorm()
  (drop_shortcut): Dropout(p=0.1, inplace=False)
)

total_params = sum(p.numel() for p in block.ff.parameters())
print(f"Total number of parameters in feed forward module: {total_params:,}")

Total number of parameters in feed forward module: 4,722,432

问：FeedForward 模块的参数量为什么是 4,722,432,而 3072×768×2 = 4,718,592 对不上?

答:差的部分是两个 Linear 层的偏置(bias)。你只算了权重矩阵,漏掉了 bias。

两个 Linear 层都是 bias=True:

• Layer 0 Linear(768, 3072):权重 768×3072 = 2,359,296,偏置 3072
• Layer 2 Linear(3072, 768):权重 3072×768 = 2,359,296,偏置 768

合计:

• 权重:2,359,296 × 2 = 4,718,592
• 偏置:3072 + 768 = 3,840
• 总计:4,722,432

差值 4,722,432 − 4,718,592 = 3,840,正好是两个偏置向量之和。

问:GELU 有参数吗?

答:没有。GELU 是固定的逐元素激活函数,不含任何可学习的权重或偏置,numel() 为 0。FeedForward 的参数全部来自那两个 Linear 层。

total_params = sum(p.numel() for p in block.att.parameters())
print(f"Total number of parameters in attention module: {total_params:,}")

Total number of parameters in attention module: 2,360,064

FeedForward 模块:

• 第 1 个 Linear 层:768 输入 × (4×768) 输出 + (4×768) 偏置 = 2,362,368
• 第 2 个 Linear 层:(4×768) 输入 × 768 输出 + 768 偏置 = 2,360,064
• 合计:2,362,368 + 2,360,064 = 4,722,432

Attention 模块:

• W_query:768 输入 × 768 输出 = 589,824(无偏置)
• W_key:768 输入 × 768 输出 = 589,824(无偏置)
• W_value:768 输入 × 768 输出 = 589,824(无偏置)
• out_proj:768 输入 × 768 输出 + 768 偏置 = 590,592
• 合计:3×589,824 + 590,592 = 2,360,064

要点:是否有偏置由各层的 bias 设置决定。W_query/W_key/W_value 用 bias=False,所以没有偏置项;而 FeedForward 的两层和 out_proj 用 bias=True,需额外加上偏置数量。

---

GPT 整体前向传播 shape 流程

符号说明：

• batch_size：一次输入的样本数量
• seq_len：每个样本的 token 数量
• emb_dim：每个 token 向量的维度（真实值 768）
• vocab_size：词表大小（真实值 50257）
• n_heads：注意力头数量（真实值 12）
• n_layers：Transformer 块层数（真实值 12）
• head_dim：每个头的维度 = emb_dim / n_heads

矩阵乘法规则：[..., m, k] @ [..., k, n] → [..., m, n]，最后两维相乘，前面维度当批次。

in_idx                          [batch_size, seq_len]
   │
   ├─ tok_emb: 查表(等价 one-hot @ 词嵌入表)
   │  [batch_size, seq_len, vocab_size] @ [vocab_size, emb_dim]
   │→ [batch_size, seq_len, emb_dim]
   │
   ├─ pos_emb: 位置查表
   │→ [seq_len, emb_dim]
   │
   ├─ 相加(广播,逐元素)         → [batch_size, seq_len, emb_dim]
   ├─ drop_emb (dropout)          → [batch_size, seq_len, emb_dim]
   │
   ├─ trf_blocks (循环 n_layers 次) ─┐
   │   每个块内部：                   │
   │   norm1 (LayerNorm)            │ [batch_size, seq_len, emb_dim]
   │   ┌ 注意力：                    │
   │   │ Q: [batch_size, seq_len, emb_dim] @ [emb_dim, emb_dim]
   │   │  → [batch_size, seq_len, emb_dim]
   │   │ K: 同上 → [batch_size, seq_len, emb_dim]
   │   │ V: 同上 → [batch_size, seq_len, emb_dim]
   │   │ 拆成多头 → [batch_size, n_heads, seq_len, head_dim]
   │   │ 注意力分数 = Q @ K的转置：
   │   │   [batch_size, n_heads, seq_len, head_dim]
   │   │ @ [batch_size, n_heads, head_dim, seq_len]
   │   │   → [batch_size, n_heads, seq_len, seq_len]
   │   │ softmax → [batch_size, n_heads, seq_len, seq_len]
   │   │ 上下文 = 分数 @ V：
   │   │   [batch_size, n_heads, seq_len, seq_len]
   │   │ @ [batch_size, n_heads, seq_len, head_dim]
   │   │   → [batch_size, n_heads, seq_len, head_dim]
   │   │ 合并多头 → [batch_size, seq_len, emb_dim]
   │   └ 输出投影：
   │     [batch_size, seq_len, emb_dim] @ [emb_dim, emb_dim]
   │→ [batch_size, seq_len, emb_dim]
   │   + shortcut(残差,逐元素)      │ [batch_size, seq_len, emb_dim]
   │   norm2 (LayerNorm)            │ [batch_size, seq_len, emb_dim]
   │   ┌ 前馈网络 FFN：
   │   │ [batch_size, seq_len, emb_dim] @ [emb_dim, 4*emb_dim]
   │   │→ [batch_size, seq_len, 4*emb_dim]
   │   │ GELU 激活 → [batch_size, seq_len, 4*emb_dim]
   │   │ [batch_size, seq_len, 4*emb_dim] @ [4*emb_dim, emb_dim]
   │   │  → [batch_size, seq_len, emb_dim]
   │   └
   │   + shortcut(残差,逐元素)      │ [batch_size, seq_len, emb_dim]
   │◄────────────────────────────┘ 出块仍是 [batch_size, seq_len, emb_dim]
   │
   ├─ final_norm (LayerNorm)      → [batch_size, seq_len, emb_dim]
   │
   ├─ out_head：
   │  [batch_size, seq_len, emb_dim] @ [emb_dim, vocab_size]
   │→ [batch_size, seq_len, vocab_size]
   │
   ▼
logits                            [batch_size, seq_len, vocab_size]

in_idx [batch, seq_len]:每个格子一个整数,只表示"第几号词",没有语义。

tok_emb [batch, seq_len, emb_dim]:多出的 emb_dim 这一串数,就是这个词的"语义坐标"——用 emb_dim 个数值刻画它的意思。

pos_emb [seq_len, emb_dim]:没有 batch 维,因为"第几个位置"和是哪句话无关;这 emb_dim 个数表示"第 i 个座位"的位置含义。

相加后 [batch, seq_len, emb_dim]:同一串 emb_dim 里现在同时编码了"是什么词 + 在第几位"。

Q/K/V 各 [batch, seq_len, emb_dim]:形状和输入一模一样,但内容分成三种角色——每个词的 emb_dim 分别被改写成"我要问什么(Q)/ 我的标签(K)/ 我的内容(V)"。

拆多头 [batch, n_heads, seq_len, head_dim]:把 emb_dim 切成 n_heads 段,每段 head_dim。多出的 n_heads 维就是"同一个词被分给几个头,各自从一个角度理解"。

分数 Q@Kᵀ [batch, n_heads, seq_len, seq_len]:最后两维从 (seq_len, head_dim) 变成 (seq_len, seq_len),含义就是"seq_len 个词 × seq_len 个词"的关系表——第 (i,j) 格 = 词 i 对词 j 的关注度。head_dim 被点积吃掉了(它只是用来算相似度的中间量)。

softmax 后 [batch, n_heads, seq_len, seq_len]:形状不变,但每一行(某个词那一行)加起来=1,变成"这个词把注意力按比例分给其他词"。

上下文 分数@V [batch, n_heads, seq_len, head_dim]:后两维从 (seq_len, seq_len) 回到 (seq_len, head_dim),含义是"每个词按关注比例把别人的 V 吸收进来",于是又变回"每个词一个 head_dim 向量",但这个向量已经带上下文了。

合并多头 [batch, seq_len, emb_dim]:n_heads × head_dim 拼回 emb_dim,含义是"把这个词在各角度下的理解拼成一个完整的、理解过上下文的它自己"。

输出投影 [batch, seq_len, emb_dim]:形状不变,把刚拼起来的多头结果重新混合成给下一层用的统一表示。

+残差 [batch, seq_len, emb_dim]:形状不变,值 = 进块前的自己 + 这层新学到的,信息不丢。

FFN 升维 [batch, seq_len, 4*emb_dim]:emb_dim 撑大到 4 倍,含义是"把每个词摊到更大空间做深加工"。

FFN 降维 [batch, seq_len, emb_dim]:加工完压回 emb_dim 规格,好继续堆下一层。

final_norm [batch, seq_len, emb_dim]:形状不变,这是每个词过完所有层后的最终理解。

out_head → logits [batch, seq_len, vocab_size]:emb_dim 换成 vocab_size,含义是"把每个位置的理解,翻译成对词表里每个词的打分",分最高的就是预测的下一个词。

主线:emb_dim 维一路承载"词的理解",中间借 seq_len×seq_len 做词间交流,最后换成 vocab_size 落到"预测哪个词"。

关于开头 [batch_size, seq_len] 是否写错

没写错，[batch_size, seq_len] 是正确的。区别在于 token id 和 one-hot 是两种表示：

• 真实输入 in_idx 存的是整数 token id（如 [[0, 3]]），形状就是 [batch_size, seq_len]，每个位置只是一个整数编号，没有 vocab_size 维度。
• [batch_size, seq_len, vocab_size] 是 one-hot 编码后的形状，把每个整数 id 展开成长度为 vocab_size、只有一个位置是 1 的向量。

顺序如下：

真实输入 in_idx        [batch_size, seq_len]              ← 整数 id
(讲解用)展开成 one-hot  [batch_size, seq_len, vocab_size]
one-hot @ 词嵌入表      @ [vocab_size, emb_dim]
结果 tok_embeds         [batch_size, seq_len, emb_dim]

图中 tok_emb 那一步的 [..., vocab_size] 是为了用矩阵乘法解释「查表」临时构造出来的，不是真实输入。实际上 PyTorch 的 nn.Embedding 不会构造 one-hot，而是直接按 id 去表里取对应行，one-hot 乘法只是数学上等价的解释方式。

问题的汇总

1. content_length 限制住了（当前限制是512），为什么会有"第 513 个词"？直接截断不行吗？
2. 训练时只见过 512，是因为显存不够吗？
3. 那为什么不直接训 1T 上下文，一步到位？
4. 用了 RoPE 为什么还有 content_length 限制？
5. 注意力的 O(n²) 到底是怎么算出来的？

1、content_length是人为设定的值，正确的因果关系：

算力 + 显存 + 数据  →  决定训练时敢用多长（比如 512 / 4k）
位置编码的外推能力  →  决定训练完之后，能不能撑到比训练更长的位置

训练用 512，但想部署时读 2048 的长文档、多轮对话、RAG 拼接的一堆文档。
2、截断能"不崩"，但不能"读懂"
input[-content_length:] 是工程上最常用的兜底。
截断的本质 = 放弃能力，把前面的内容直接扔掉：3000 字合同问"第一条和最后一条矛盾吗" → 看不到第一条，无法回答。正因为很多高价值任务不能截断，业界才要死磕"扩展能力"这条路。
3、为什么训练时不敢用很长？—— 算力/显存是主因

1. 注意力的 O(n²) 复杂度 ← 注意力矩阵是 n×n，序列翻倍，显存和计算量翻 4 倍（下面有计算）
2. 激活值 / KV 显存：随长度线性增长
3. 长文本训练数据稀缺：高质量、连贯的超长文档本来就少
4. 边际收益递减：大部分文本的有效依赖没那么长

4、“为什么不直接训 1T 上下文，一步到位？”
被 O(n²) 按死了， O(n²) 到底怎么算出来的？
第 1 步 —— 算单个元素的成本。Q @ Kᵀ：[seq_len, head_dim] @ [head_dim, seq_len] → [seq_len, seq_len]。矩阵乘法规则 [a,b] @ [b,c] → [a,c]，中间维 b 相等，被消掉（求和）。

Q 的第 i 行:   [q1, q2, ..., q_d]     ← 长度 = head_dim (记作 d)
Kᵀ 的第 j 列:  [k1, k2, ..., k_d]     ← 长度 = head_dim (d)

score[i][j] = q1·k1 + q2·k2 + ... + q_d·k_d
              └────────── d 次乘加 ──────────┘

算一个元素 = d 次乘加。

第 2 步 —— 算有多少个元素。输出是 n×n，共 n × n 个格子，每个都要重复上面：

┌─────────────────────────────────────┐
│ score 矩阵  (n 行 × n 列)            │
│    j=1  j=2  j=3  ...  j=n           │
│ i=1 [■]  [■]  [■]  ...  [■]          │ 每个 [■] 都要
│ i=2 [■]  [■]  [■]  ...  [■]          │   算 d 次乘加
│ i=3 [■]  [■]  [■]  ...  [■]          │
│ ...                                  │
│ i=n [■]  [■]  [■]  ...  [■]          │
│    一共 n × n 个格子                 │
└─────────────────────────────────────┘

第 3 步 —— 两步相乘 = 总运算量：

总运算量 = 格子数量 × 每个格子的成本
        = (n × n)  ×   d
        =  n² · d  →  O(n²)    ← d(head_dim) 是常数，扔掉

也可以这么算：
第1行(1个向量) × 第1列 → score[1][1] → d 次乘加
第1行(1个向量) × 第2列 → score[1][2] → d 次乘加
…
第1行(1个向量) × 第n列 → score[1][n] → d 次乘加
一整行 = n × d，那么一共有n行，所以是nnd

直观感受"平方"的可怕：

n = 512:   512²  = 262,144      每个 ×64 ≈ 1678 万次运算
n = 1024:  1024² = 1,048,576    ≈ 6710 万次   (n 翻倍 → ×4)
n = 2048:  2048² = 4,194,304    ≈ 2.7 亿次    (再 ×4)

n 每翻一倍，运算量翻 4 倍——这就是"平方"的含义。

别忘了：注意力里有两次 n² 乘法：

① 算分数:   Q @ Kᵀ    → [n,d]@[d,n] = [n,n]    运算量 n²·d
② 加权求和: 分数 @ V   → [n,n]@[n,d] = [n,d]    运算量 n²·d
                          └─ 这里中间维是 n！

第 ② 步中间被消掉的维度是 n，所以 n × d × n = n²·d，同样 O(n²)。

“用了 RoPE，为什么还有 content_length 限制？”

RoPE 不是万能药，限制依然在，两层原因：

一、RoPE 自己也会外推衰减。RoPE 靠不同频率的旋转编码位置，训练只见过 0~4096 的旋转角度，硬推到 100k 时，低频维度的旋转角度进入训练时从没出现过的区间，模型对这些角度没学过，注意力照样乱。所以 RoPE 只是比正弦编码外推得远、崩得温柔，不是不崩。想真正拉长还得加：

• 位置插值 PI：把 100k 的位置"压缩"回 4096 的角度范围
• NTK-aware scaling / YaRN：按频率分段调整，兼顾远近
• 而且通常还要再微调才稳

二、就算 RoPE 外推完美，KV cache 的显存墙还在。这是最根本的。RoPE 只解决"位置编码怎么表示远距离"，完全没解决 O(n²) 计算和 KV cache 显存：上下文开到 1M，KV cache 就要几十上百 GB；推理速度随长度暴跌。

所以厂商设 content_length（128k、200k）是综合权衡：位置编码撑得住 + 显存装得下 + 速度能接受 + 效果不明显掉。位置编码只是约束之一。

修正后的完整图景

训练时长度上限  ←  主要是算力 / 显存 / 数据（O(n²) 是元凶）
                    位置编码不是主因

推理时能否超过训练长度  ←  这里位置编码才登场
                          正弦编码：几乎不能，直接崩
                          RoPE：能撑远一点，但也会衰减
                          RoPE + PI/NTK + 微调：能大幅拉长

最终 content_length     ←  上面所有约束的综合权衡
                          （位置编码能力 + KV 显存 + 速度 + 效果）

一句话总结

• content_length 是"我们敢开多大"的旋钮，不是物理墙
• 截断 = “认输”（别崩），外推/长上下文技术 = “打赢”（能读）
• 训练长度上限的核心元凶是注意力 O(n²)：n² 个格子 × 每格 d 次运算，d 是常数，剩 n²
• 所以不能"一步到位训 1T"，物理上做不到
• RoPE 缓解了位置编码的外推，但既会自身衰减、又没解决 KV 显存墙，所以 content_length 依然存在

最后，计算一下 GPTModel 对象中 1.63 亿个参数的内存需求：

# Calculate the total size in bytes (assuming float32, 4 bytes per parameter)
total_size_bytes = total_params * 4

# Convert to megabytes
total_size_mb = total_size_bytes / (1024 * 1024)

print(f"Total size of the model: {total_size_mb:.2f} MB")

Total size of the model: 621.83 MB

总之，通过计算 GPTModel 对象中 1.63 亿个参数的内存需求，并假设每个参数是占用 4 字节的 32 位浮点数，我们发现模型的总大小为 621.83 MB，这表明即使是相对较小的大语言模型也需要相对较大的存储容量。

现在，我们已经实现了 GPTModel 架构，并看到它输出了形状为[batch_size, num_tokens, vocab_size]的数值张量。接下来，我们将编写代码，将这些输出张量转换为文本。

---

练习 4.2 初始化更大的 GPT 模型

我们已经初始化了一个参数量为 1.24 亿的 GPT 模型，即“GPT-2 small”。在不修改代码 的情况下，只需更新配置文件，即可使用 GPTModel 类实现 GPT-2 medium（具有 1024 维嵌入、24 个 Transformer 块和 16 个多头注意力头）、GPT-2 large（具有 1280 维嵌入、36 个 Transformer 块和 20 个多头注意力头）和 GPT-2 xl（具有 1600 维嵌入、48 个 Transformer 块和 25 个多头注意力头）。同时，计算每个 GPT 模型的参数总数。

GPT_CONFIG_124M = {
    "vocab_size": 50257,
    "context_length": 1024,
    "emb_dim": 768,
    "n_heads": 12,
    "n_layers": 12,
    "drop_rate": 0.1,
    "qkv_bias": False
}


def get_config(base_config, model_name="gpt2-small"):
    GPT_CONFIG = base_config.copy()

    if model_name == "gpt2-small":
        GPT_CONFIG["emb_dim"] = 768
        GPT_CONFIG["n_layers"] = 12
        GPT_CONFIG["n_heads"] = 12

    elif model_name == "gpt2-medium":
        GPT_CONFIG["emb_dim"] = 1024
        GPT_CONFIG["n_layers"] = 24
        GPT_CONFIG["n_heads"] = 16

    elif model_name == "gpt2-large":
        GPT_CONFIG["emb_dim"] = 1280
        GPT_CONFIG["n_layers"] = 36
        GPT_CONFIG["n_heads"] = 20

    elif model_name == "gpt2-xl":
        GPT_CONFIG["emb_dim"] = 1600
        GPT_CONFIG["n_layers"] = 48
        GPT_CONFIG["n_heads"] = 25

    else:
        raise ValueError(f"Incorrect model name {model_name}")

    return GPT_CONFIG


def calculate_size(model): # based on chapter code
    
    total_params = sum(p.numel() for p in model.parameters())
    print(f"Total number of parameters: {total_params:,}")

    total_params_gpt2 =  total_params - sum(p.numel() for p in model.out_head.parameters())
    print(f"Number of trainable parameters considering weight tying: {total_params_gpt2:,}")
    
    # Calculate the total size in bytes (assuming float32, 4 bytes per parameter)
    total_size_bytes = total_params * 4
    
    # Convert to megabytes
    total_size_mb = total_size_bytes / (1024 * 1024)
    
    print(f"Total size of the model: {total_size_mb:.2f} MB")


from gpt import GPTModel


for model_abbrev in ("small", "medium", "large", "xl"):
    model_name = f"gpt2-{model_abbrev}"
    CONFIG = get_config(GPT_CONFIG_124M, model_name=model_name)
    model = GPTModel(CONFIG)
    print(f"\n\n{model_name}:")
    calculate_size(model)

gpt2-small:
Total number of parameters: 163,009,536
Number of trainable parameters considering weight tying: 124,412,160
Total size of the model: 621.83 MB


gpt2-medium:
Total number of parameters: 406,212,608
Number of trainable parameters considering weight tying: 354,749,440
Total size of the model: 1549.58 MB


gpt2-large:
Total number of parameters: 838,220,800
Number of trainable parameters considering weight tying: 773,891,840
Total size of the model: 3197.56 MB


gpt2-xl:
Total number of parameters: 1,637,792,000
Number of trainable parameters considering weight tying: 1,557,380,800
Total size of the model: 6247.68 MB

---

4.7 生成文本

接下来，我们将编写代码，将 GPT 模型的张量输出转换成文本。在开始之前，让我们简要回顾一下像大语言模型这样的生成模型是如何逐词（词元）生成文本的。

图 4-16 展示了 GPT 模型如何根据输入上下文（比如“Hello, I am”）逐步生成文本的过程。 每次迭代时，输入上下文都会增加，从而使模型生成连贯且符合上下文的文本。在第 6 次迭代时，模型生成了完整的句子“Hello, I am a model ready to help.”我们已经看到当前的 GPTModel 实现输出的张量形状为[batch_size, num_token, vocab_size]。现在的问题是：GPT 模型如何将这些输出张量转化为生成的文本？

GPT 模型将输出张量转化为生成文本的过程涉及多个步骤，如图 4-17 所示。这些步骤包括解码输出张量、根据概率分布选择词元，以及将这些词元转换为人类可读的文本。

图 4-17 中展示的下一词元生成过程说明了 GPT 模型如何在给定输入的情况下生成下一个词元。在每一步中，模型输出一个矩阵，其中的向量表示有可能的下一个词元。将与下一个词元对应的向量提取出来，并通过 softmax 函数转换为概率分布。在包含这些概率分数的向量中，找到最高值的索引，这个索引对应于词元 ID。然后将这个词元 ID 解码为文本，生成序列中的下一个词元。最后，将这个词元附加到之前的输入中，形成新的输入序列，供下一次迭代使用。这个逐步的过程使得模型能够按顺序生成文本，从最初的输入上下文中构建连贯的短语和句子。

问：GPT 输出的向量是一个个对应输入向量的吗？输入 4 个词为什么输出 4 个向量，每个又代表什么？

答：是的，一一对应。GPT 并行处理，输入几个 token 就输出几个向量。每个向量代表"看到当前位置及之前所有词后，对下一个词的预测"。以 “Hello, I am” 为例：

• 对应 “Hello” 的向量 → 预测 “Hello” 之后接什么
• 对应 “,” 的向量 → 预测 “Hello,” 之后接什么
• 对应 “I” 的向量 → 预测 “Hello, I” 之后接什么
• 对应 “am” 的向量 → 预测 “Hello, I am” 之后接什么

每个位置都在预测自己后面的词，所以输入 4 个就输出 4 个。

问：既然有多个向量，为什么推理时只用最后一个？输入 1000 个是不是前 999 个都丢弃，前面计算白做了？

答：推理（生成文本）时你只关心"整句话之后的下一个词"。前面向量预测的都是你已经知道的内容（就是你输入的词），没价值，只有最后一个向量有用。所以输入 1000 个 token，只用第 1000 个向量预测第 1001 个 token，前 999 个的预测结果这一步确实被丢弃。

但前面的计算不是白做：被丢弃的只是"下一个词预测"这个结论，而每个位置的内部计算（注意力里的 Key/Value）是最后一个位置做预测时必须依赖的——预测最后一个词要"看"前面所有词。

问：那每生成一个新词都要把全部 token 重算一遍吗？

答：不用，工程上靠 KV Cache 复用。第一次（prefill）并行处理全部输入，取最后一个向量得到新词；之后每一步（decode）前面 token 的 Key/Value 已缓存，只需计算新加入的那 1 个 token。所以除了第一次，后续每步实际只算 1 个新 token。

一句话总结：
输入 N 个 token → 输出 N 个向量（并行预测每个位置的下一个词）。推理时只采用最后一个向量的预测，前面预测丢弃但内部计算通过 KV Cache 复用；训练时则会用上全部 N 个预测信号，一次前向传播学习每个位置，效率很高。

在实际操作中，我们会多次重复这一过程（参见图 4-16），直到生成预定数量的词元。在代码中，可以按照代码清单 4-8 所示的步骤实现文本生成过程。

def generate_text_simple(model, idx, max_new_tokens, context_size):
    # idx 是当前上下文中索引的 (batch, n_tokens) 数组
    for _ in range(max_new_tokens):
        # 如果当前上下文超过支持的上下文大小，则进行裁剪
        # 例如，如果 LLM 只支持 5 个 token，而上下文大小为 10，
        # 那么只使用最后 5 个 token 作为上下文
        idx_cond = idx[:, -context_size:]

        # 获取预测结果
        with torch.no_grad():
            logits = model(idx_cond)

        # 只关注最后一个时间步
        # (batch, n_tokens, vocab_size) 变为 (batch, vocab_size)
        logits = logits[:, -1, :]

        # 应用 softmax 获取概率
        probas = torch.softmax(logits, dim=-1)  # (batch, vocab_size)

        # 获取概率值最高的词表条目的索引
        idx_next = torch.argmax(probas, dim=-1, keepdim=True)  # (batch, 1)

        # 将采样得到的索引追加到运行序列中
        idx = torch.cat((idx, idx_next), dim=1)  # (batch, n_tokens+1)

    return idx

这段代码展示了使用 PyTorch 为语言模型生成循环的简单实现。它首先循环指定次数以生成新词元，然后将当前上下文裁剪到模型的最大上下文大小，接下来进行预测计算，并根据最高概率选择下一个词元。

在编写 generate_text_simple 函数时，我们使用 softmax 函数将 logits 转换为概率分布，并通过 torch.argmax 确定最大值的位置。softmax 函数是单调的，这意味着它在转换为输出时保持了输入的顺序。因此，实际上 softmax 步骤是冗余的，因为 softmax 输出张量中最高分的位置与 logits 张量中的位置是相同的。换句话说，可以直接对 logits 张量应用 torch.argmax 函数，得到相同的结果。不过，我们提供了转换代码来展示将 logits 转换为概率的完整过程，这有助于更好地理解模型如何生成最有可能的下一个词元，这一过程被称为贪心解码。

当我们在第 5 章中实现 GPT 训练代码时，将使用额外的采样技术来调整 softmax 输出，从而避免模型总是选择最有可能的词元。这将为生成的文本带来可变性和创造力。

图 4-18 进一步展示了如何通过 generate_text_simple 函数逐步生成一个词元 ID 并将其 附加到上下文中的过程。（每次迭代的词元 ID 生成过程已在图 4-17 中详细描述。）我们通过迭代的方式生成了词元 ID。例如，在第一轮迭代中，模型获取了与“Hello, I am”对应的词元，以此预测下一个词元（“a”，ID 为 257），并将其添加到输入中。这个过程会重复进行，直到模型经过 6 轮迭代生成完整的句子“Hello, I am a model ready to help.”

现在让我们尝试使用"Hello, I am"上下文作为模型输入来调用 generate_text_simple 函数。首先，将输入上下文编码为词元 ID：

start_context = "Hello, I am"

encoded = tokenizer.encode(start_context)
print("encoded:", encoded)

encoded_tensor = torch.tensor(encoded).unsqueeze(0)
print("encoded_tensor.shape:", encoded_tensor.shape)

encoded: [15496, 11, 314, 716]
encoded_tensor.shape: torch.Size([1, 4])

encoded = tokenizer.encode(start_context) # 假设得到 [15496, 11, 314, 716]
encoded_tensor = torch.tensor(encoded) # shape: [4] 一维
encoded_tensor = encoded_tensor.unsqueeze(0) # shape: [1, 4] 二维
unsqueeze(0) 在第 0 个位置插入一个新维度,把形状从 [4] 变成 [1, 4]。

接下来，将模型设置为.eval()模式，这将禁用诸如 dropout 等只在训练期间使用的随机组件。 然后对编码后的输入张量使用 generate_text_simple 函数：

model.eval() # disable dropout

out = generate_text_simple(
    model=model,
    idx=encoded_tensor, 
    max_new_tokens=6, 
    context_size=GPT_CONFIG_124M["context_length"]
)

print("Output:", out)
print("Output length:", len(out[0]))

# 把 token ID 解码成文字
# out 的形状是 [1, N](batch 维度 + token 序列),out.squeeze(0) 去掉 batch 维度变成一维 [N],和之前 unsqueeze(0) 正好相反。
# .tolist() 把张量转成 Python 的普通列表,因为 tokenizer.decode 通常接收的是整数列表。
decoded_text = tokenizer.decode(out.squeeze(0).tolist())
print("Decoded text:", decoded_text)

Output: tensor([[15496,    11,   314,   716, 27018, 24086, 47843, 30961, 42348,  7267]])
Output length: 10
Decoded text: Hello, I am Featureiman Byeswickattribute argue

可以看到，模型生成了无意义的内容，完全不像连贯的文本“Hello, I am a model ready to help.” 发生了什么？模型不能生成连贯文本的原因是我们还没有对其进行训练。到目前为止，我们只是实现了 GPT 架构，并用初始随机权重初始化了 GPT 模型实例。模型训练是一个重要的主题，我们将在第 5 章中进行深入探讨。

---

练习 4.3 使用单独的 dropout 参数

在本章开头，我们在 GPT_CONFIG_124M 字典中定义了一个全局的 drop_rate 设置来控制 GPTModel 架构中各个位置的 dropout 率。请修改代码，为模型架构中的不同 dropout 层指定不同的 dropout 值。（提示：模型中有 3 个不同的 dropout 层：嵌入层、快捷连接层和多头注意力模块。）

GPT_CONFIG_124M = {
    "vocab_size": 50257,
    "context_length": 1024,
    "emb_dim": 768,
    "n_heads": 12,
    "n_layers": 12,
    "drop_rate_emb": 0.1,        # NEW: dropout for embedding layers
    "drop_rate_attn": 0.1,       # NEW: dropout for multi-head attention  
    "drop_rate_shortcut": 0.1,   # NEW: dropout for shortcut connections  
    "qkv_bias": False
}

import torch.nn as nn
from gpt import MultiHeadAttention, LayerNorm, FeedForward


class TransformerBlock(nn.Module):
    def __init__(self, cfg):
        super().__init__()
        self.att = MultiHeadAttention(
            d_in=cfg["emb_dim"],
            d_out=cfg["emb_dim"],
            context_length=cfg["context_length"],
            num_heads=cfg["n_heads"], 
            dropout=cfg["drop_rate_attn"], # NEW: dropout for multi-head attention
            qkv_bias=cfg["qkv_bias"])
        self.ff = FeedForward(cfg)
        self.norm1 = LayerNorm(cfg["emb_dim"])
        self.norm2 = LayerNorm(cfg["emb_dim"])
        self.drop_shortcut = nn.Dropout(cfg["drop_rate_shortcut"])

    def forward(self, x):
        # Shortcut connection for attention block
        shortcut = x
        x = self.norm1(x)
        x = self.att(x)  # Shape [batch_size, num_tokens, emb_size]
        x = self.drop_shortcut(x)
        x = x + shortcut  # Add the original input back

        # Shortcut connection for feed-forward block
        shortcut = x
        x = self.norm2(x)
        x = self.ff(x)
        x = self.drop_shortcut(x)
        x = x + shortcut  # Add the original input back

        return x


class GPTModel(nn.Module):
    def __init__(self, cfg):
        super().__init__()
        self.tok_emb = nn.Embedding(cfg["vocab_size"], cfg["emb_dim"])
        self.pos_emb = nn.Embedding(cfg["context_length"], cfg["emb_dim"])
        self.drop_emb = nn.Dropout(cfg["drop_rate_emb"]) # NEW: dropout for embedding layers

        self.trf_blocks = nn.Sequential(
            *[TransformerBlock(cfg) for _ in range(cfg["n_layers"])])

        self.final_norm = LayerNorm(cfg["emb_dim"])
        self.out_head = nn.Linear(cfg["emb_dim"], cfg["vocab_size"], bias=False)

    def forward(self, in_idx):
        batch_size, seq_len = in_idx.shape
        tok_embeds = self.tok_emb(in_idx)
        pos_embeds = self.pos_emb(torch.arange(seq_len, device=in_idx.device))
        x = tok_embeds + pos_embeds  # Shape [batch_size, num_tokens, emb_size]
        x = self.drop_emb(x)
        x = self.trf_blocks(x)
        x = self.final_norm(x)
        logits = self.out_head(x)
        return logits

import torch

torch.manual_seed(123)
model = GPTModel(GPT_CONFIG_124M)

---

4.8 小结

• 层归一化可以确保每个层的输出具有一致的均值和方差，从而稳定训练过程。
• 快捷连接是通过将一层的输出直接传递到更深层来跳过一个或多个层的连接，它能帮助缓解在训练深度神经网络（如大语言模型）时遇到的梯度消失问题。
• 作为 GPT 模型的核心模块组件，Transformer 块融合了掩码多头注意力模块和使用 GELU激活函数的全连接前馈神经网络。
• GPT 模型是具有许多重复 Transformer 块的大语言模型，这些 Transformer 块有数百万到数十亿个参数。
• GPT 模型具有多种规模，比如参数量分别为 1.24 亿、3.45 亿、7.62 亿和 15.4 亿的模型，我们可以使用相同的 GPTModel Python 类来实现它们。
• 类 GPT 大语言模型的文本生成能力涉及根据给定的输入上下文来逐个预测词元，然后将输出张量解码为人类可读的文本。
• 在没有训练的情况下，GPT 模型生成的文本是不连贯的，这显示了模型训练对于生成连贯文本的重要性。