Andrej Karpathy：从零复现 GPT-2

在这节 Zero to Hero 课程中，我将复现 1.24 亿参数的 GPT-2，这是 GPT-2 系列中最小的模型。这个系列也让我们观察模型规模变化时不同任务的表现。

GPT-2 是一个模型系列，而不是单个模型。论文中的变体从这里复现的小模型一直到大得多的模型；随着规模增大，下游任务得分通常会提高。

按照发布的代码仓库说明，论文最初的表格算错了参数总数。我使用修正后的 1.24 亿参数描述最小模型，而不是把表格中的数字当作精确值。

模型有 12 层 Transformer，宽度为 768。我假设观众已通过前一节视频熟悉部分架构；这次的目标是实现并训练它。

目标是在留出的文本上得到较低验证损失。我估计，使用课程中的云端硬件复现这个小模型可能需要约一小时和十美元；这些数字仅对应当时的示例。

OpenAI 发布了 GPT-2 的权重，但论文对一些训练细节交代不够清楚。我也会参考 GPT-3 论文中的优化设置，因为它的架构变化不算很大。

我先从终点开始：加载已发布的检查点并从中采样。原始实现使用 TensorFlow；Hugging Face 已把权重转换为 PyTorch 实现，更便于这次演示。

已发布权重提供了具体目标。训练过程的说明则不那么完整，因此要区分匹配模型架构与精确重建原始训练过程。

在 Transformers 接口中，模型名 `gpt2` 指 1.24 亿参数的检查点；更大的 15 亿参数模型使用 `gpt2-xl`。我提取状态字典来检查原始张量，而不是把库当作黑箱。

模型的状态字典让我们检查参数名称、张量和形状。token 嵌入表有 50,257 行、768 列：词表中的每个 token 都有一个学习得到的 768 维表示。

嵌入表的每一行对应 GPT-2 词表中的一个 token，它不一定是一个完整单词。学习得到的向量是网络对该 token 的表示。

GPT-2 还为最多 1,024 个位置学习位置嵌入。加入位置信息后，模型可以同时利用 token 的身份和它出现的位置。

位置嵌入矩阵图将位置放在一条轴上、通道放在另一条轴上。这些模式从随机初始化的表中学习得到，因此它们在优化过程中的出现值得关注。

画出位置嵌入后，可以看到一些通道随位置平滑变化，形状类似波。这些值是学习得到的参数，因此可见结构是在优化中形成的；至于每个通道具体做什么，则更难解释。

一些位置通道在特定位置范围内响应更强。这些曲线并不完美且带有噪声，单靠可视化并不能建立完整解释。

我暂时把这些曲线的锯齿状形态解释为小模型训练尚不充分的迹象。这是对权重图的推测，并非对训练是否充分的直接测量。

这与原始 Transformer 中固定的正弦位置编码不同。GPT-2 从头学习位置表示，其中一些呈现出类似正弦波的模式。

其他权重矩阵也能看到结构，但单凭图片无法说明网络学到了什么。研究这些模式可能有助于机制可解释性工作。

采样提示是“Hello, I’m a language model,”，流水线生成五条续写。即使固定随机种子，库或代码变化也可能使样本与旧示例不同。

从已发布模型采样可以得到连贯的续写。这确认我们能够加载检查点；下一步是构建自己的模块并匹配它的行为。

在实现自己的版本之前，可以通过 Transformers 库对检查点采样。这为后续移植提供了一个可运行的参考。

能够把已发布权重复制进自己的实现，说明它属于相同的架构类型。但从头训练仍会得到不同权重，也可能使用不同数据。

GPT-2 调整了层归一化相对于原始 Transformer 的位置，并增加最后的归一化。复现现有检查点时，这些位置细节很重要。

GPT-2 是仅解码器架构：它省略原始 Transformer 图中的编码器以及对编码器输出的交叉注意力。因果自注意力与前馈模块构成主要的重复堆叠。

模块名称按照检查点来安排。匹配结构和命名有助于逐个比较及转移张量。

我把模型组织为 token 和位置嵌入、Transformer 模块列表、最终归一化以及语言模型输出头。小型检查点的列表中有十二个模块。

token 嵌入和位置嵌入都是可索引的表。语言模型头将最终的 768 维表示投影为词表大小的分数；在此实现中，输出头没有偏置。

残差流携带表示穿过网络。每个模块先对输入归一化，再经过注意力或前馈网络，并把各自的结果加回残差流。

干净的残差路径有助于优化，因为加法让梯度既流经跳跃连接，也流经计算分支。该分支随后可以贡献自己学到的更新。

前置归一化使残差路径能够直接穿过模块。注意力和多层感知机的计算分别为这条路径贡献更新。

注意力让不同位置交换信息，多层感知机则处理各位置的表示。这有助于直观理解反复执行的两种计算，但不能完整解释它们学到的一切。

前馈网络在两个线性投影之间应用 GELU。它先在内部扩展表示，再把宽度还原为残差流的宽度，然后执行相加。

前馈模块先扩展通道维度，再投影回原来的宽度。为尽量匹配 GPT-2，我使用它的近似 GELU 激活，而不假设所有 GELU 实现的数值完全相同。

近似 GELU 有历史原因：课程提到，在早期环境中计算精确的误差函数成本较高。这里使用近似版本，是因为目标是匹配 GPT-2。

在多头注意力中，投影产生查询、键和值。把它们拆成多个头后得到并行的注意力流，结果再拼接并投影回残差流。

我把 GELU 与 ReLU 的负半轴平坦区域作比较：后者的梯度可能为零。GPT-2 的选择是一种历史架构设计，较新的架构可能使用其他方案。

在因果注意力中，未来位置的分数会在 softmax 前被排除。这确保当前位置的预测不会依赖训练序列中后面的 token。

紧凑实现把头的维度当作额外的批次维度，使 PyTorch 能并行计算各个头。重塑和转置则让数据布局符合这些运算的需要。

因果掩码阻止当前位置关注未来 token。Softmax 把允许的注意力分数转为权重，再与值相乘，形成对可用信息的加权组合。

多个注意力头是对投影后通道的一种重排。每个头独立计算注意力，输出重新组合后再进行最后的投影。

模型配置除了 12 层和 768 个通道，还使用 12 个注意力头。GPT-2 切词器有 50,257 个词表项，包括字节级片段、合并结果和文本结束 token。

架构搭好后，我把 Hugging Face 的参数加载到自己的模块。名称和形状必须匹配；由于两个实现的存储方式不同，部分矩阵需要转置。

某些旧版 GPT-2 层使用的权重矩阵方向与 PyTorch 线性层不同。转置这些权重可以保持预期计算。

状态字典中的每个条目都不一定是学习参数。移植时会忽略可以重新构造的掩码缓冲区，同时复制实际训练权重并检查形状。

复制权重前，我检查参数键和张量维度。成功加载状态字典是一个有用的检查，而生成文本还可以提供端到端验证。

前向传播接收 token 索引，相加 token 嵌入和位置嵌入，让序列通过各个模块，并返回每个位置上针对词表的 logits。

如果输入的批次和时间维度分别为 B 与 T，输出形状就是 B 乘 T 乘词表大小。每个向量描述当前位置之后可能出现的 token 的分数。

生成时，我用 GPT-2 切词器编码文本前缀，并在一个批次中放入多个副本。模型和输入张量必须位于同一设备上。

提示词需要先转换成 token ID 才能交给网络。采样循环完成后，生成的 token ID 可以再解码为文本。

把前缀的多个副本放进一个批次，可以一次得到多条续写。由于下一个 token 是采样得到的，续写内容可能彼此不同。

采样循环运行模型，取最后一个位置的 logits，将其转换为概率并采样下一个 token；然后追加这个 token，重复到所需长度。

Top-k 采样把抽取范围限制在一组得分较高的候选中。这会改变续写分布，并在保留随机性的同时使样本更接近可能的文本。

每次迭代中，模型都会看到不断增长的前缀。新采样的 token 列会拼接到现有批次，供下一次迭代使用。

输出是连贯的，虽然具体样本与先前的例子不同。随机数生成和其他实现细节都可能让采样结果变化；这里的重要检查是加载和生成能够运行。

我加入设备检测，让脚本根据可用情况使用 CUDA、苹果的 MPS 后端或 CPU。不同选择的性能差异很大。

模型和输入都需要明确放到对应设备。如果只把其中之一移到加速器，同一计算中就会出现不兼容的张量。

接下来从推理转向训练。模型需要 token 序列批次和向后错开一个 token 的目标，使每个位置都能学习预测下一个 token。

为了初步调试，我对一个小型文本数据集切词。这是有意选择的简单起点，不是最终训练语料。

对于小批次，输入与目标是同一个 token 缓冲区中相邻的切片。由于缓冲区多包含一个 token，每个输入行的最后一个 token 也有对应目标。

玩具数据使用同一个 GPT-2 切词器编码。对于下一 token 预测，目标使用的是同一数据流向后错开一个位置的版本。

可以把连续缓冲区重塑为 B 行、每行 T 个 token。多读取一个 token 就能得到匹配的目标：输入与目标数组相差一个位置。

每一行代表一条训练序列。该行中的 token 可以关注自己的因果上下文中之前的位置，而不能关注未来目标。

例如，24 个 token 可以视为四行六列数组：前六个成为第一行，接下来六个成为第二行，依此类推。目标缓冲区从后一个位置开始切片，再变成相同形状。

模型返回 B 乘 T 乘 C 的 logits，其中 C 是词表大小。计算损失时可以把批次和时间维度展平，让每个预测对应其目标。

交叉熵衡量预测分布给实际下一个 token 分配了多大概率。我们把展平后的 logits 和目标传给 PyTorch 的损失函数。

如果初始概率接近均匀分布，可以根据词表大小估计初始损失。这为优化开始前提供一个粗略的合理性检查。

初始化时，接近词表均匀分布的预测会产生一个有代表性的基准损失。如果初始损失偏差很大，就值得检查模型或数据。

如果 50,257 个 token 的概率近似均匀，预期交叉熵约等于词表大小的对数。模型学到有用偏好之前，训练损失应当从接近这一量级开始。

在完整训练之前，我先尝试过拟合单个批次。这个小训练问题应当表现为损失下降；否则前向传播、梯度或优化器可能有错误。

优化循环依次清除旧梯度、计算损失、执行反向传播并调用优化器步骤。如果小型固定批次无法拟合，就应该能够逐步检查这些环节。

Adam 维护梯度的运行统计量来调整更新尺度。第一和第二矩的细节会影响优化，但第一次实验主要检查循环能否学习。

同一固定批次上的损失下降，说明系统能拟合这个小目标。但这并不能说明模型能够泛化到未见过的文本。

设备不匹配暴露了一个常见 PyTorch 错误：对张量调用 `.to(device)` 会返回移动后的张量，因此需要接收返回值。只移动模型并不会移动独立创建的输入缓冲区。

过拟合固定批次只是调试练习。要训练有用的模型，我们需要不断提供新批次，而不是反复优化同一批样本。

我编写一个简易数据加载器，在 token 缓冲区中向前移动，返回输入和错位目标。对于这个玩具数据集，到达末尾后就从头循环。

加载器维护当前位置，并按一个批次的输入 token 数向前推进。它有意保持简洁，使我们在引入大型数据集之前可以验证基本数据流。

输入批次最初位于 CPU，随后复制到选定的训练设备。对于小数据集，一个 epoch 的迭代次数取决于批次大小和序列长度。

玩具加载器无法再提供完整批次时会回到开头。生产环境的数据流水线则需要更仔细地考虑边界和打乱。

token 嵌入表与输出投影具有兼容的维度。GPT-2 将这两个权重绑定，使两个模块属性指向同一个底层张量。

权重绑定让读取 token 和给可能的输出 token 打分共用表示，也减少了相比独立表格所需的参数数量。

在这个配置中，共享嵌入矩阵约有 4,000 万个元素。将它与输出投影绑定，可以避免维护第二个同样大小的独立矩阵。

绑定输入与输出嵌入的理由是两张表都把词表项与网络表示空间联系起来。实现必须在优化过程中保持这种共享。

共享张量必须真正指向同一存储，而不只是初始化时数值相等。检查指针可以看清这种区别。

同一个绑定张量同时参与输入和输出。反向传播时，两处使用产生的贡献都会更新这组共享参数。

权重绑定也是一种归纳偏置：它假设读取 token 所用的嵌入与把它作为输出打分所用的表示之间存在有用关系。

初始化也需要注意。已发布的 GPT-2 代码使用标准差为 0.02 的正态分布初始化权重，而论文对某些细节没有明说。

已发布代码用标准差 0.02 初始化普通权重，偏置初始化为零；位置嵌入的处理略有不同。为了教学实现，我对嵌入模块统一使用 0.02。

残差分支反复向残差流加入贡献。对某些残差投影的初始化做缩放，有助于控制层数累积时方差的增长。

叠加许多独立初始化的残差贡献，可能使激活方差随深度增长。调整初始化的目的是让初始残差流更加稳定。

我给相关投影层加上标记并使用修改后的初始化，同时让其他线性层和嵌入权重保持默认尺度。实现细节需要对照正在复现的架构检查。

在一个简化示例中，反复加入 100 个单位尺度的独立随机贡献，标准差会增长到约 10。把各贡献乘以数量平方根的倒数，可以抵消这种增长。

残差投影在初始化时采用与深度有关的缩放；实现中的其他部分继续使用选定的基础标准差。

每个 Transformer 模块通过注意力和 MLP 两次向残差流相加。因此，这个实现的缩放计算使用 Transformer 层数的两倍。

由于嵌入层和语言模型头共享张量，遍历模块初始化时可能两次遇到同一个张量。我在检查初始化代码时特别指出这一点。

现在已有模型、数据加载器和基本训练循环。下一个挑战是吞吐量：数值格式与 GPU 执行路径都可能显著影响训练速度。

GPU 可以并行执行大量运算，但也需要在计算单元和内存间移动数据。工作负载可能受到算术吞吐、内存传输或额外开销限制。

GPU 内存容量限制了可以同时保存的批次和激活大小。内存带宽和计算吞吐量则是另外两类限制。

float32 每个数使用 32 位。神经网络训练通常可以容忍许多运算使用较低精度，但仍需要处理数值稳定性。

浮点格式要在指数范围和有效数字精度之间取舍。合适的选择取决于计算中出现的数值以及硬件支持的运算。

NVIDIA Tensor Cores 可以加速矩阵乘法。TF32 在受支持的运算中对输入使用较低精度表示，同时以更宽的格式累加。

这个小模型的大量计算发生在矩阵乘法中，包括从 768 个通道投影到词表大小 logits 的大型运算。Tensor Cores 可以加速符合条件的矩阵运算。

Tensor Core 的精度模式同时涉及输入表示和累加器精度。更快的算术运算以输入精度有所损失为代价。

我测量训练耗时和每秒处理的 token 数，而不只看代码是否运行。准确计时需要考虑 GPU 的异步执行，例如在读取时间前同步。

在示例中，第一次迭代可能因为设置和缓冲区初始化而较慢。后续迭代能更有代表性地反映稳定运行时的吞吐量。

在比较优化之前，我增大批次以更好利用 GPU，然后记录时间和吞吐量。要得到有意义的比较，必须对相同工作负载进行基准测试。

批次大小会影响设备利用率和内存使用。实验改用更大批次后，每秒 token 数比单步耗时更适合比较。

如果批次放不进内存，我建议逐步减小，直到能够运行。更大的批次可以改善利用率，但需要结合优化时使用的有效批次大小理解。

启用 TF32 可以在精度有所取舍的情况下加速符合条件的运算，但不会把程序中的每个张量都变成更小的存储格式。实际加速幅度取决于工作负载。

PyTorch 提供设置，允许 float32 矩阵乘法采用较低精度。我启用它，再次运行相同的训练工作负载。

这项设置会影响符合条件的矩阵乘法，例如线性层中的矩阵乘法。它不意味着所有激活值和保存的参数都被转换成 TF32。

如果工作负载受其他因素限制，包括在内存中移动全精度张量，实际收益可能小于理论算术加速。

float16 和 bfloat16 是两种较低精度格式。bfloat16 保留较宽的指数范围，但精度位数较少；在支持它的硬件上，这个范围可能使训练更方便。

指数控制可表示的范围，有效数字则控制范围内的精度。bfloat16 和 float16 对有限位数的分配方式不同。

使用 float16 时，梯度缩放可以帮助防止很小的梯度下溢。具体保护措施取决于精度格式和训练设置。

PyTorch autocast 为适合的运算选择较低精度，同时让其他运算保持较稳妥的格式。这是混合精度，并不意味着每个参数和中间结果都用同一种格式存储。

autocast 上下文围绕前向计算应用精度选择。我们仍要检查实际张量，而不能假设程序的所有部分使用同一种 dtype。

我检查权重和中间激活的 dtype 来说明这一区别。优化器可以保留全精度参数，同时让部分矩阵乘法使用较低精度。

autocast 放在前向传播和损失计算周围；反向传播与优化器步骤保持在该上下文之外。这遵循了课程中讨论的 PyTorch 指引。

在本次演示使用的 A100 上，采用 bfloat16 的 autocast 改善了运行速度。数值表现和加速效果取决于硬件与工作负载。

具体的 autocast 策略很重要：矩阵乘法可以使用较低精度，而对精度敏感的运算可以保留 float32。我不会假设每一层都会自动转换。

接着介绍 torch.compile。编译可以减少 Python 和内核启动开销，并融合运算，但调试与兼容性会影响它是否方便使用。

在本次演示中，模型编译完成后，测得的单步时间明显缩短。这个数字反映特定模型、批次和 A100 配置，并不保证普遍适用。

编译有启动成本，因此计时应区分第一次编译执行和后续步骤。课程中的速度比较针对的是这一特定训练工作负载。

没有融合时，中间结果可能要写入 GPU 内存，再为下一次运算读回来。融合后的内核可以让更多工作留在靠近计算单元的位置。

编译器可以同时看到一串逐元素运算，并生成组合内核。在即时执行模式中，各次独立调用可能分别产生内存传输与启动开销。

即使算术运算很快，内存带宽仍然重要。一系列简单的逐元素运算可能受反复移动中间数组限制，而不是受计算次数限制。

融合通过合并运算，避免部分往返全局内存的传输。收益取决于融合计算内部能保留多少工作和中间数据。

片上存储与更大的 GPU 内存之间的差异，解释了为什么移动中间张量可能很昂贵。在数据仍靠近计算单元时重复利用它，可以避免一部分往返传输。

Flash attention 重新组织注意力计算，以减少昂贵的内存传输。它可能执行更多算术运算，却因为避免物化和反复移动大型注意力矩阵而运行得更快。

Flash attention 的关键思路是安排内存访问。因此，只计算浮点运算次数可能无法正确预测哪种实现运行更快。

PyTorch 提供优化的缩放点积注意力。我用该运算替换手写的注意力序列，并指定因果行为，确保未来 token 仍被遮蔽。

优化后的注意力调用必须保留因果掩码。只有在计算仍符合语言模型所需行为时，更快的内核才有意义。

硬件内核通常更高效地处理与合适倍数对齐的维度。GPT-2 词表有 50,257 个 token，因此我在这次训练中把输出维度补齐到 50,304。

切词器仍只产生原有 ID；额外行是一种计算上的补齐选择。性能改善来自更适合 GPU 内核的形状，而不是增加了有用的新词表。

补齐展示了看似很小的维度选择如何影响内核效率。最佳对齐方式取决于硬件和具体运算。

原始切词器不会把新增词表项作为输入 ID 产生。它们仍会进入输出分布，因此模型可以学习给它们较低概率。

补齐带来的实测收益随 PyTorch 版本和所用内核变化。我说明演示中的设置，而不把某个百分比当作普遍结论。

加速工作让训练循环变得实用，但仅有速度还不足以复现结果。我们还需要合理选择优化器、学习率调度、批次大小和训练数据。

我查阅 GPT-3 论文中较明确的超参数。复现使用 AdamW 和梯度裁剪，并参考公开设置，同时承认 GPT-2 的细节并未全部说明。

梯度裁剪测量梯度的全局范数，在超过选定阈值时对其重新缩放。这可以限制异常大的更新，减少它们破坏训练稳定性的可能。

打印梯度范数可以观察裁剪是否发生。阈值控制更新使用的最大范数，并不保证得到某个特定损失。

固定学习率并非唯一选择。这个例子的调度先进行预热，然后按余弦曲线衰减到较低的最终学习率。

梯度裁剪示例把全局范数上限设为 1.0。报告的范数也是诊断信息：异常大的值或尖峰可能指向训练不稳定或特殊批次。

预热在最初的步骤中逐渐提高学习率。之后，余弦衰减在设定的训练范围内平滑降低学习率；这只是众多调度方案中的一种。

作为参考的小型 GPT-3 配置在论文中给出的峰值学习率为 6e-4。示例在选定调度期间将学习率降至峰值的百分之十左右。

调度器根据当前步数计算学习率。超过衰减范围后，实现可以让学习率保持在最低值，而不是无限继续余弦曲线。

这次复现的调度并非论文的精确副本：这里把衰减范围和总训练范围设为相同。这也是需要将其描述为采用选定设置的复现的原因之一。

权重衰减通常应用于矩阵权重，而不应用于偏置和其他一维参数。我将参数分组，让优化器分别处理。

GPT-3 论文还描述了训练期间逐渐增大批次。本实现跳过该过程，让 token 和步数计算更简单，同时解释为什么早期较小批次可能更高效。

在这个简单设置中，加载器会不放回地遍历数据直到 epoch 边界。随后在优化器中加入权重衰减，作为单独的正则化选择。

融合的 AdamW 实现可以合并原本需要启动许多小型 GPU 内核的工作。我在优化器和 CUDA 设备支持时启用它。

融合优化器运算可以减少为各参数张量启动许多独立内核的开销。收益取决于优化器支持情况与设备。

我把优化器设置放进一个接收相关超参数并报告参数分组的辅助函数，让训练循环更专注于各个步骤的顺序。

使用融合选项之前要先检查它是否可用。较旧的 PyTorch 版本或不同设备可能需要普通优化器路径。

学习率、批次大小、衰减和 Adam 设置彼此影响。在本次演示中，我以公开的选择作为起点，而不是声称每个数值都普遍最优。

目标批次可能大到无法放进 GPU 内存。梯度累积让我们把它拆成微批次，多次执行前向和反向传播，最后只进行一次优化器更新。

累积步数取决于期望的总 token 数和每个微批次的 token 数。这样可以保持有效批次大小，而不必一次分配全部内存。

参数缓冲区中的梯度会跨多次反向传播累积，直到被清零。只有预定微批次都贡献梯度后才执行优化器步骤。

由于批次损失通常取平均，反向传播前要把每个微批次的损失除以累积步数。随后将梯度相加，才对应有效批次上的平均值。

我用一个小型回归例子演示其算术关系。计算完整批次梯度，与累积适当缩放的单样本损失所得梯度，除数值舍入外应当一致。

回归例子揭示了一个常见错误：对未经缩放的单样本损失梯度求和，其尺度不同于这些损失平均值的梯度。

把每个损失除以微批次数量，就能恢复预期的平均值。各次反向传播会把相应贡献加到同一组梯度缓冲区中。

在这里展示的单 GPU 配置中，计算出的累积次数为 32。这个数字由本次运行的目标批次与微批次维度决定，不是通用规则。

分布式数据并行把训练扩展到多个 GPU。每个进程拥有模型副本并处理批次的不同部分；梯度经过同步，使副本一致更新。

torchrun 启动器会启动指定数量的进程，并设置用于识别进程的环境变量。这改变了脚本相对于单进程运行的启动方式。

启动器提供 rank、local rank 和 world size。每个进程选择自己的 GPU，rank 为零的进程可以承担日志和保存检查点等公共工作。

主进程约定是一种避免重复输出的实用方法。每个 rank 仍然参与分布式计算和同步。

每个 GPU 的微批次大小只是总量的一部分。world size 和累积步数都参与决定一次更新的有效批次。

数据加载器必须给不同 rank 分配不同样本。否则多个 GPU 可能重复做相同工作，使期望的有效批次缺少应有的多样性。

如果日志不区分 rank，每个进程都可能打印相同消息。把公共输出限制在主进程，可以让日志更易读。

示例中，我为每个 GPU 启动一个进程。运行结束时需要有序清理分布式进程组，否则框架可能报告关闭警告。

数据加载器为每个进程设置不同起始偏移，并按所有 rank 总共消耗的数据量向前推进。这样每个进程就能得到数据流中自己的部分。

反向传播后，分布式训练通过 all-reduce 在各 rank 之间平均梯度。随后每个副本都可以对自己的参数应用相同更新。

跨进程同步梯度需要通信。分布式包装器处理了许多底层机制，但训练脚本仍须正确管理数据和日志。

进行梯度累积时，每个微批次都同步会浪费通信。我们可以先在本地累积，在一次更新的最后一个微批次完成时再同步。

最后一次累积反向传播后，all-reduce 的结果让每个进程得到平均梯度。进程内的缩放与跨进程的平均都必须符合预期优化目标。

计算吞吐量时应包含所有进程和完整有效批次。只统计单个进程的 token 数，会低估分布式训练的总工作量。

使用八个 GPU 时，示例把累积工作分配给各进程，使每个进程的累积次数低于单 GPU 设置。

单 GPU 与多 GPU 的损失不一定每一步都完全一致，因为数据加载顺序可能不同。比较时需要考虑每次运行实际看到哪些样本。

最后一个主要组成部分是数据。公开网页既有有价值的文本，也有垃圾内容和其他噪声，因此大型数据集通常要经过过滤与处理，而不是直接使用原始网页。

我介绍 FineWeb，并在实验中使用其教育子集。数据集和过滤选择与原始 GPT-2 训练数据不同，解读复现结果时必须考虑这一点。

准备脚本加载文档、将其切词并写成 token 分片。相比一个巨大的文件，分片更便于存储和读取大型数据集。

保存之前，文档会被转换为 token ID 数据流。在构造出的数据流中，保留的文本结束标记用于分隔文档。

加载器可以选择训练或验证分片。保留单独的验证部分，可以让我们在不用于参数更新的数据上评估模型。

我根据计划的 token 预算和有效批次大小计算优化步数。这些选择决定训练会消耗多少语料。

预热步数通过 token 数除以有效批次大小得到。以 token 表达调度，有助于比较不同批次大小的配置。

估计总运行时间是用测得的单步时间乘以计划步数得到的。这是针对所展示云端配置的估算。

验证使用留出的数据，并且不进行梯度更新。定期计算验证损失可以观察模型在训练批次之外的表现，定期采样则展示模型生成的文本。

每次验证时，加载器都会重置，模型在留出的批次上评估且不追踪梯度。对损失取平均可以降低报告值对单个批次的依赖。

我还可以在同一留出数据上评估已发布的 GPT-2 检查点。这为该数据集提供有用比较，同时仍需看到训练数据和流程之间的差异。

评估和采样的执行频率低于训练步骤，以限制额外开销。我还遇到了与编译有关的采样问题，调整了演示，并指出之后需要进一步处理。

为了增加一个基准，我使用 HellaSwag：这是一项选择续写的任务，旨在测试类似常识的预测能力。它是有用的信号，但不能完整衡量语言模型。

HellaSwag 的一道题提供一个共同上下文和四个可能的结尾，要求判断哪个结尾是合理续写。

随着系统改进，基准难度与模型分数也会变化。课程中的历史分数应放在当时的基准与受评模型背景下理解。

每个例子包含上下文和几个候选结尾。我把候选内容切词，并使用掩码，使评分关注续写部分，而非填充位置或共同上下文。

填充和掩码让不同长度的候选结尾能够在同一批次评估。对相关续写 token 取平均，可以避免仅仅因为 token 较少而偏向短候选。

脚本比较各候选续写的平均损失，选出损失最低的一项。我用同样的方法评估已发布的 GPT-2 检查点，为自己的训练建立参考。

模型分别给每个候选打分；它并不是在同一提示中同时看到四个选项后作选择。我们在模型外比较计算出的分数。

在这套评估中，脚本得到的已发布 1.24 亿参数 GPT-2 分数约为 29.5%，而四选一随机基线为 25%。这是后续训练的参考值。

配置好验证、基准评估和定期生成后，我启动更大规模的训练。不同指标回答不同问题：预测损失、基准表现和样本文本的定性表现。

进行分布式基准评估时，各进程评估一部分样本，再合并正确数与总数。最终准确率根据合并后的计数计算。

训练脚本定期记录验证和基准结果并保存样本。这些记录让我们能在隔夜训练结束后检查运行情况。

第二天早上，我检查结果。复现模型在训练中明显改善，并且在这套设置下，HellaSwag 分数与已发布的小型 GPT-2 检查点相比表现不错。

在留出的 FineWeb-Edu 数据上，新模型的验证损失低于已发布检查点。由于后者使用不同分布的数据训练，我把它当作交叉检查，而不是完全公平的比较。

这种比较存在限制：训练语料和 token 预算不同，基准数据污染或评估代码差异也可能影响结果。因此不能把它视为对原始训练过程的精确重放。

初次运行处理约 100 亿个 token。它在展示的 HellaSwag 评估中超过小型 GPT-2 参考模型，但数据集差异和基准可能重叠，使训练效率比较变得复杂。

基准表现可能受到训练数据和评估流程差异的影响。比较单个分数有参考价值，却不能证明两次训练经历等价。

我注意到训练损失有周期性模式，怀疑这可能与 FineWeb-Edu 样本的顺序有关。我并不确定原因，并把它标记为仓库中需要调查的问题。

我还查看生成样本，作为另一项定性检查。样本可以流畅或有趣，但这并不能证明事实准确，也不能证明模型具备广泛能力。

较长的隔夜运行对样本大约遍历四次，处理约 400 亿个 token。其 HellaSwag 分数进一步提高，接近但没有达到图中展示的小型 GPT-3 参考值。

演示中的加载器在各次遍历之间重复文档顺序。打乱文档顺序可以减少模型反复看到同样偶然文档边界的情况。

一个数据细节是文档如何通过结束标记连接成数据流。如果每个 epoch 都重复同一顺序，偶然的相邻关系也会被保留下来，因此值得考虑打乱文档顺序。

较长运行在展示的 HellaSwag 评估中达到约 33.24%。我还查看其续写，认为它们比短运行的样本更连贯，同时承认样本只是定性检查。

改变模型宽度、深度或上下文长度，可以让实现接近 GPT 系列的其他成员。更大的配置也会改变计算与内存需求。

如果要更贴近 GPT-3，还需要把上下文长度从 1,024 改为 2,048 个 token。若要保持每步 token 总数不变，就需要相应调整每个设备的批次维度。

只保存模型参数就足以进行某些后续评估，但要忠实恢复中断的训练，还需要更多状态：优化器统计量、当前步数、数据位置和随机数状态都可能重要。

这次训练得到的是预测文档续写的基础语言模型。要把它变成对话助手，还需要进一步工作，例如使用对话数据进行监督微调；本课止于预训练。

我介绍相关的 C 和 CUDA 实现 llm.c，它可以高效训练类似模型。它的细节和结果不会自动与本次 PyTorch 演示完全一致。

较底层的实现可能与 PyTorch 做出不同工程取舍。比较速度或复现结果时，我建议查看它自己的代码和文档。

本课串联了架构、数据、精度、高效内核、分布式训练和评估。build-nanogpt 仓库记录了这些修改，读者可以检查实现并关注后续勘误。

视频结束时列出了尚待解决的问题：损失的异常周期性、此版本中编译对生成和评估的干扰，以及需要更好地打乱数据。后续修正记录在代码仓库中。