从零打造LLM：第22讲深入模型训练实践

LLM 从零构建（48）

NSLU2 异地备份项目（13）

微调 LLM（10）

复古语言模型（2）

训练主机 poppy（1）

软件开发工具（1）

:: (Bloggable a) => a -> IO ()

David Friedman 的 Substack

创业极客精神

某种意义上的“魔法”

kaleidic.ai newsletter

听测试山羊的话！

Simon Willison 的 Weblog

一些确信程度各不相同的观点

从零编写 LLM，第 22 部分——终于开始训练我们的 LLM！

这篇文章为我阅读 Sebastian Raschka 的《Build a Large Language Model (from Scratch)》第 5 章所做的笔记收尾。理解交叉熵损失（cross entropy loss）和困惑度（perplexity）是这一章最让我头疼的部分；剩下的 28 页更多只是把各个部分拼装起来，再运行代码看看会发生什么。

这篇文章这么短，甚至让人觉得有点虎头蛇尾。前面 22 篇写了那么多，现在似乎真的没有太多可说的了——但这掩盖了一个事实：这一部分大概是整本书里实践起来最令人兴奋的内容。此前 140 页里，我们极其细致地构建了这么多组件，学到了这么多东西，却一直没有多少成果可供展示——突然之间，我们拥有了一套可以直接拿训练集全速开跑的代码库，而我们的模型开始和我们说话了！

我使用书中的示例数据集训练了模型，也就是 Edith Wharton 的《The Verdict》中大约 20,000 个字符的内容，然后让它预测“Every effort moves you”之后的 token。得到的结果是：

Every effort moves you in," was down surprise a was one of lo "I quote.

对于一个只用这么少的数据、训练时间仅略多于十秒的模型来说，已经相当不错了。

下一步是按照书中的说明，从 OpenAI 下载最初那个拥有 124M 参数的 GPT-2 版本的权重，然后把它们加载到我的模型中。使用这些权重，对同一个 prompt 进行生成，我得到了：

Every effort moves you as far as the hand can go until the end of your turn unless something interrupts your control flow. As you may observe I

这实在太酷了。它的连贯程度足以让人相信，这是一段游戏说明。

接下来我不会详细复述本章剩余的内容——正如我所说，它本质上只是把我们此前介绍过的各个部分组装起来，尽管最终效果非常出色。这篇文章只会简要记录几件我觉得有意思的事。

随机性与设置 seed

对于所有跟着这本书动手实践的人，我确实非常建议你亲手输入全部代码，并且自己运行一遍——这真的会帮助你记住各部分是如何组合在一起的。

不过，我发现这样做会遇到一个小问题：书中有不少示例，其代码输出会受到随机性的影响。例如，在开始训练之前查看模型处理某段示例文本时的损失，或者让模型在训练过程中生成样本。

理论上，因为 Raschka 在这些代码前面都调用了 torch.manual_seed，所以你得到的结果应该和书里的输出完全一致。然而，我们到了这个阶段，代码量已经相当大了——比如，我们有不少在前面章节中创建的辅助函数，其中一些也会使用随机性。

这意味着，如果想得到与书中完全相同的结果，你必须确保所有使用随机性的代码，都严格按照 Raschka 为这本书运行代码时的顺序执行。事实证明，要做到这一点出人意料地困难！

我的直觉是，这其实并没有那么重要。只要你看到的损失值和书中的结果处于同一个大致范围，生成的内容在训练前同样差不多不连贯，并且以感觉相近的速度逐渐变得更加连贯，那就没问题。最需要注意的，可能是训练正式开始时的表现：你应该像书中一样，看到训练集上的损失持续稳定下降；同样，验证集损失也应该像书里那样，很早便进入平台期。

过去在简单的神经网络中实现反向传播时，我通常会把梯度乘以一个很小的数，也就是学习率（learning rate），再从各自对应的参数中减去这个结果，从而得到更新后的参数——这就是经典的随机梯度下降（stochastic gradient descent）。

严肃一些的 ML 会使用 optimiser；我在微调 LLM 时接触过它们，上周写 RNN 代码时也用过一个。你不再亲自更新参数，而是通过调用 optimiser 的 step 函数，让它替你完成更新。在大多数教材中，AdamW 似乎是默认的 optimiser；不过，如果我的 AI X/Twitter 信息流可信的话，实际应用中最流行的似乎是 Muon。

我并不真正了解 optimiser 的具体工作原理，以后还得深入研究。不过，目前我脑海中有一个高度简化的整体认识：它们会随时间动态调整学习率，让你在刚开始时更容易沿着梯度向下迈出较大的“步子”，之后再逐渐缩小步幅。我认为它们有时也能避开损失地形中的局部最小值。我曾在某处读到过一个很形象的比喻（遗憾的是找不到出处了）：简单梯度下降就像让一个球从山坡上滚下去，而（某些？）optimiser 会为这个球提供一点动量，只要总体坡度仍然向下，它就能凭惯性越过一小段上坡。

总之，以后还需要进一步研究。

在实践中使用 AdamW 时，你会在训练循环开始时初始化它，并指定学习率——我猜它和旧代码中的学习率类似，是梯度的缩放因子——以及权重衰减（weight decay，:shrug:）。你还需要把它将要管理的参数提供给它。

在训练循环中，每个输入 batch 开始时，先调用 optimizer.zero_grad()，让它把自己管理的梯度清零；然后让数据通过模型并计算损失；接着调用 loss.backward() 得到梯度；最后只需调用 optimizer.step()，它就会完成参数更新。

还是那句话，以后我想更深入地研究 optimiser 的工作原理。但就目前而言，我认为知道这些已经够用了。

速度与训练成本

书中介绍了如何使用一本属于公版作品的书来训练模型，也就是 Edith Wharton 的《The Verdict》。在大多数人手头现有的硬件上进行完整训练，成本会高得惊人，所以我们这里只使用这个简短的示例进行训练；之后再学习如何下载和使用 OpenAI 为其 GPT-2 模型公开的权重。

不过，有件事让我有点意外。谈到在《The Verdict》上进行的训练时，Raschka 说，在 MacBook Air 上“完成大约需要 5 分钟”。

在我的机器上使用 RTX 3090 的 CUDA 运行时，它只花了不到 11 秒。

当然，这完全说得通——AI 训练通常在 GPU 或定制硬件上完成是有充分理由的，而 MacBook Air 大概是在 CPU 上训练。但即使只是这么简单的示例，两者之间的差距如此巨大，仍然让我有些惊讶！

书中提到，训练 Llama 2 可能花费了数十万美元。不过我必须承认，我确实很好奇：在自己的硬件上训练一个拥有 124M 参数的模型，究竟要花多少钱？又或者，在我进行微调实验时从 Lambda Labs 租用的、配备 8 块 80GiB A100 GPU 的机器上训练，又要花多少钱？

Andrej Karpathy 使用自己手写的 C/CUDA LLM 系统 llm.c，只花了 20 美元就训练出了一个 124M 参数的 GPT-2 模型。毫无疑问，它的效率比我们在这本书中使用的 PyTorch 代码更高。但如果能弄清楚这件事对我来说到底是否可行，确实会非常有趣！他使用的训练数据，是 FineWeb 集合中包含 10B token 的版本，可以免费获取。1

我想，读完这本书之后，我已经有了一个很适合继续做的项目：看看自己在本地每秒能够训练多少 token——这样就能估算出，在完整训练集上跑完一个 epoch 需要多长时间。我猜，这个时间可能会长到让我不愿意让桌面电脑一直被这项任务占用；但之后我可以结合微调时学到的经验，看看能否在 Lambda Labs 上把它运行起来。如果成本只是几十美元，甚至一百美元左右，我真的觉得这件事值得一试！

“记忆”、temperature 与 top-k sampling

这一章中有一点让我有些困惑——严格来说，这只是个吹毛求疵的小问题——那就是防止“记忆”（memorisation）的部分。我想，这是因为我对这个词的理解与书中的用法并不一致。

对我来说，记忆是模型在训练过程中做的事情——如果像我们这里这样，反复使用一个只有 20,000 字符的文件训练拥有 124M 参数的模型，那么无论如何，模型都会把它记下来，这是不可避免的。从这个意义上说，要减少记忆，唯一的方法就是增加训练数据量；即使如此，正如《纽约时报》起诉 OpenAI 的诉讼结果所表明的那样，仍然会有一些内容被模型记住。

书中所说的“记忆”，更接近我所谓的“鹦鹉学舌”（parroting）——模型因为始终选择概率最高的下一个词，所以只会重复它记住的内容。当然，避免这种情况极其重要！只是这样的表述方式让我有些困惑。

无论如何，这些技术都很巧妙。第一种方案非常直观：直接把经过 softmax 的 logits 当作概率分布，并从中采样。在此基础上，temperature 是一个很聪明的技巧：只要在进行 softmax 之前，将 logits 除以一个大于 1 的数，最终得到的分布就会更加平坦；反过来，除以一个小于 1 的数，则可以让分布变得更加“尖锐”。书中用来展示其工作方式的图非常棒，不过我还让 Claude 做了一个 temperature 交互实验网站，感觉它让我理解得更加清楚。

最后是 top-k 技术：只考虑概率最高的 k 个 token，然后再进行 temperature 和 softmax 计算。这是在前面方案之上增加的一项合理改进。它的代码非常巧妙：找出最大的 k 个 logits，获取其中最小的那个值，再把所有小于该值的 logit 替换成负无穷。让这些值通过 softmax 后，被替换的项会得到 0，概率分布则由剩余的项构成。

所以，这些内容非常出色，书中的解释也很清楚——只是至少按照我对这个词的理解，我不觉得它所做的事情可以被确切地称为防止“记忆”。

下载 OpenAI 权重

在本章结尾，我们从 OpenAI 的网站下载了他们最初发布的 GPT-2 模型权重，并将其加载到自己的模型中。

值得庆幸的是，下载权重的代码不需要亲手输入，因为可以直接从 GitHub 下载。而且在一个与之相关的特定场景中，我还要推翻自己之前关于亲手输入代码的建议：我强烈建议你也直接从 GitHub 复制 load_weights_into_gpt，也就是把下载的权重复制进我们自己模型的那个函数。我确实亲手把它完整输入了一遍，但不觉得自己从中获得了什么。

在阅读这一节时，我还注意到，自己在撰写这个系列的过程中一直犯了一个错误：我原以为所有 GPT-2 模型的 embedding 维度都是 768。事实证明，只有该系列中的 124M 模型是这样，更大的模型拥有更多 embedding 维度。这非常合理——我也已经更新了之前的文章，以反映这一点。

关于第 5 章剩下的内容，我真正想补充的也就这么多。正如开头所说，对于书中一个能产出如此惊人结果的部分，自己却只写了这么点内容，感觉简直有些让人失望！不过，现在我们已经拥有了一个可以工作的 LLM，也至少具备了相应的基础：如果资源足够，或许就能从零开始训练自己的模型。

接下来：用它进行文本分类。这会不会既快速又简单？还是会把我们带进另一个引人入胜的兔子洞？时间会给出答案……

这里是本系列下一篇文章的链接。

他的新项目 nanochat——一个可以从零开始训练的 chatbot——更加酷。↩

他的新项目 nanochat——一个可以从零开始训练的 chatbot——更加酷。↩