消费级 GPU 上从零训练大语言模型实战指南

从零开始训练基础模型：在 RTX 3090 上的实践

在完成了 Sebastian Raschka 的著作《Build a Large Language Model (from Scratch)》的主要内容后，我想做一个实验：能否在自己的硬件上训练一个基础模型？

这本书向你展示了如何训练你的 LLM，在一个小数据集上做了基础训练运行，然后我们转向下载 OpenAI 的"预烹饪"权重。考虑到并非每个读者都有足够的硬件从零开始真正训练，这是有道理的。在本系列的最开始，我对微调 LLM 时得到的数字做了一些天真的缩放计算，得出结论认为在合理的时间内这是不可能的。

但是我用 RTX 3090 在书的小型训练运行中获得的速度让我想，也许——仅仅是也许！——实际上可能在自己的硬件上训练这个大小的模型——约 163M 参数。也许不是在小笔记本上，但至少可以在一台相当高端的"游戏" PC 上。

另外，Andrej Karpathy 最近推出了 nanochat，"$100 能买到的最好的 ChatGPT"。他在主页上提到他训练了一个名为 d32 的模型，有 32 个 Transformer 层，1.9B 参数，花费约 $800。他更小的 20 层 d20 模型有 561M 参数，他说在 8x H100 GPU 节点上应该能在约四小时内训练完成，这样的节点每小时成本约 $24——因此总价格约为 $100。

nanochat 更有趣的是它是用 PyTorch 构建的；最初我有印象说它是基于他纯 C/CUDA 的 llm.c，这我想象会带来巨大的加速。但不是——他使用的是我在本系列中一直使用的同一个技术栈！

Karpathy 的模型都比 163M 参数更大，所以看起来这肯定是可行的。显然，我作为 AI 开发人员的经验远不如他丰富，他使用的机器更大（8 个 GPU，每个 GPU 都有超过 3 倍多的 VRAM），但他也在四小时内包括了训练分词器和指令微调的时间——而他更小的模型也比我的大三倍多。所以这一切都应该有帮助。

这篇文章的结构比我的 LLM 从零开始系列中的其他文章要少一些，因为它本质上是我在完成这个项目时所记笔记的整理版本。

但为了不埋没要点：使用 Hugging Face 的 FineWeb 系列数据集，我能够在仅仅超过 48 小时的时间内在自己的硬件上训练一个 GPT-2 小型基础模型，使其达到几乎与原始模型一样好的水平！基础模型：不只是大型 AI 实验室的专属。

这是完整的故事。

模型配置

对于这个项目，我想使用 Raschka 在《LLM from Scratch》书中呈现的完全相同的模型代码——我的副本在这里。自 GPT-2 以来，LLM 的架构有许多改进，但现在最好保持简单。

但仍有一些设置需要决定。我们一直在使用的模型的配置字典有这些参数：

vocab_size。这由分词器确定，我想使用 GPT-2 的，所以它需要是 50257。

context_length。GPT-2 有 1,024 个词元的上下文长度，所以我会坚持这个。

emb_dim, n_heads, n_layers——这些定义了我们训练的不同 GPT-2 模型类的哪一个，我想坚持最小的 gpt2-small，所以它们分别为 768、12 和 12。

drop_rate。对我来说，上面链接的"架构改进"文章中最令人惊讶的事情之一是 dropout 不再被那么广泛使用。然而，这似乎与自 GPT-2 以来流行的单轮次训练有关，所以我认为最好在这里坚持 0.1。

qkv_bias。从 Raschka 在书中所说的来看，这不会增加太多价值，尽管原始 GPT-2 使用了它，所以让我们将其设置为 False。

还有权重绑定的方面——原始 GPT-2 重复使用了其嵌入矩阵作为线性层的权重，该线性层将来自最后一个 Transformer 层的上下文向量投影到词汇空间以获得逻辑。

但是，我们一直在使用的代码中没有强制执行这一点——当我们在书中做小规模训练时，我们为这两个步骤使用了独立的权重。它唯一被"强制执行"的时候是当我们从 OpenAI 下载预训练权重时，我们将相同的值放入嵌入矩阵和最终输出头。

考虑到 Raschka 说总的来说最好避免权重绑定，而实际进行绑定会比不进行绑定更难，那么不进行绑定似乎是显而易见的选择。

那么，这对我们的模型意味着什么？

In [1]: big_train_params = {
   ...:     "vocab_size": 50257,
   ...:     "context_length": 1024,
   ...:     "emb_dim": 768,
   ...:     "n_heads": 12,
   ...:     "n_layers": 12,
   ...:     "drop_rate": 0.1,
   ...:     "qkv_bias": False
   ...: }

In [2]: from gpt import GPTModel

In [3]: model = GPTModel(big_train_params)

In [4]: sum(p.numel() for p in model.parameters())
Out[4]: 163009536

这与我们在阅读这本书时得到的一致；163M 参数。我们能训练它吗？

数据选择

似乎每个 AI 项目都从"我们可以使用什么数据？"这个问题开始。

关于 GPT-2 的原始报告《Language Models are Unsupervised Multitask Learners》令人恼火地缺乏细节。但是，它确实说他们在"800 万份文档上训练，总共 40 GB 的文本"。现在，根据 OpenAI，可以合理地假设典型英文文本中大约每个词元 4 个字符。所以 40 GB 的文本约为 100 亿个词元。这些数据基本上是通过抓取从 Reddit 链接的页面获得的，这些页面在 Reddit 上有超过三个赞成票，所以质量相当高。我们能获得类似的东西吗？

方便地，Hugging Face 托管了一个称为 FineWeb 的大型数据集，它有一个 100 亿词元"样本"数据集，从完整的 18.5 万亿词元中随机选择。所以这个样本感觉起来在数量级上是对的。而在阅读更多关于 Karpathy 的 nanochat 时，我注意到它使用 FineWeb-Edu，这是 FineWeb 的一个版本，包含"仅最教育性的网页"。

我写了一个脚本来下载这两个，然后启动了它。每一个花费大约 20 分钟（我书房里的网络很慢，我获得的速度 < 5MB/s）；FineWeb 的 10B 样本占用约 29 GiB，FineWeb-Edu 的约 27 GiB。

数据加载和探索

是时候看看它们了。Hugging Face 数据集的 load_dataset 函数加载你提供的所有文件，你可以告诉它如何将它们分成训练/验证/测试集。这个命令只是加载整个 FineWeb，说"将其全部视为训练分割"，现在已经足够了：

In [1]: from datasets import load_dataset

In [2]: fw = load_dataset(
   ...:     "parquet",
   ...:     data_files="./fineweb/sample/10BT/*.parquet",
   ...:     split="train"
   ...: )
Generating train split: 14868862 examples [01:53, 130852.34 examples/s]
Loading dataset shards: 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████| 102/102 [00:03<00:00, 31.90it/s]

哎哟。生成训练分割花了 1 分 53 秒。但是，这似乎是一次性成本——当我稍后在不同的 Python 会话中使用相同的代码再次访问它时，它只进行第二个"加载数据集分片"部分，花费三秒钟，而不是生成分割。想必它缓存了。

无论如何，让我们看看里面有什么：

In [3]: print(fw)
Dataset({
    features: ['text', 'id', 'dump', 'url', 'date', 'file_path', 'language', 'language_score', 'token_count'],
    num_rows: 14868862
})

太好了，所以我们有 14,868,862 行，每一行都有各种信息。检查第一个的文本：

In [7]: print(fw[0]["text"][:500])
|Viewing Single Post From: Spoilers for the Week of February 11th|
|Lil||Feb 1 2013, 09:58 AM|
Don't care about Chloe/Taniel/Jen-Jen. Don't care about Sami, really, but hoping
that we get some good "SAMANTHA GENE!!" Marlena Death-Stares out of it. And
"newfound" feelings. Please. If only.
STEFANO!! STEFANO, STEFANO, STEFANO!!!! :cheer:
|Spoilers for the Week of February 11th · DAYS: News, Spoilers & Discussion|

好吧，对于 FineWeb 来说，这看起来不是特别"精细"，但我想这比 Karpathy 在他最近与 Dwarkesh Patel 的采访中谈论的要好：

当你看一个前沿实验室的预训练数据集并查看一个随机互联网文档时，它完全是垃圾。我根本不知道这如何能起作用。它就像股票代码、符号一样，是来自互联网各个角落的大量垃圾

让我们看看 FineWeb-Edu。

In [8]: fw_edu = load_dataset(
   ...:     "parquet",
   ...:     data_files="./fineweb-edu/sample/10BT/*.parquet",
   ...:     split="train"
   ...: )
Generating train split: 9672101 examples [01:32, 104057.34 examples/s]
Loading dataset shards: 100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████| 98/98 [00:02<00:00, 48.62it/s]

In [9]: print(fw_edu[0]["text"][:500])
The Independent Jane
For all the love, romance and scandal in Jane Austen's books, what they are
really about is freedom and independence. Independence of thought and the
freedom to choose.
Elizabeth's refusal of Mr. Collins offer of marriage showed an independence
seldom seen in heroines of the day. Her refusal of Mr. Darcy while triggered by
anger showed a level of independence that left him shocked and stunned.
The freedom she exhibited in finally accepting him in direct defiance of Lady Cath

这看起来好多了！

现在让我们看看按词元计的文档长度。有一个 token_count 列，但我不知道那是为哪个分词器的，所以为了安全起见，我们会自己计算。

对 FineWeb 10B 中的每一行进行词元化以检查需要多长时间？让我们对 14,868,862 个中的前 10,000 个进行词元化，然后看看那需要多长时间——然后我们可以计算出整个过程的估计时间。

In [25]: import tiktoken

In [26]: import time

In [27]: tokenizer = tiktoken.get_encoding("gpt2")

In [28]: start = time.time()
    ...: for entry in fw.select(range(10_000)):
    ...:     tokenizer.encode(entry["text"])
    ...: end = time.time()

In [29]: end - start
Out[29]: 1.4528205394744873

In [30]: fw
Out[30]:
Dataset({
    features: ['text', 'id', 'dump', 'url', 'date', 'file_path', 'language', 'language_score', 'token_count'],
    num_rows: 14868862
})

In [31]: (14868862 / 10_000) * 1.4528205394744873
Out[31]: 2160.1788112211702

2,160 秒或约 36 分钟。哎哟！

但经过一番挖掘后，我发现 tiktoken 分词器可以处理批处理（文档记录不完善，但它在源代码中）：

In [45]: text_batch = ["a", "b", "c"]

In [46]: tokenizer.encode_batch(text_batch)
Out[46]: [[64], [65], [66]]

另外，我们可以在整个 HF 数据集上映射一个函数，这可以用多个进程运行。所以，我们可以结合两者：

In [47]: import os

In [53]: def add_len(examples):
    ...:     texts = [t or "" for t in examples["text"]]
    ...:     tokens = tokenizer.encode_batch(texts, disallowed_special=())
    ...:     return {"tok_len": [len(t) for t in tokens]}
    ...:

In [54]: start = time.time()
    ...: fw_with_len = fw.map(
    ...:     add_len,
    ...:     batched=True,
    ...:     batch_size=1024,
    ...:     num_proc=os.cpu_count(),
    ...: )
    ...: end = time.time()
Map (num_proc=24): 100%|████████████████████████████████████████████████████████████████████████████████████████████| 14868862/14868862 [03:15<00:00, 75869.33 examples/s]

略超过三分钟，还不错！（上面命令计数从 47 跳到 53 的原因是在第一次运行中我在那里没有 disallowed_special=()——数据集中的一行包含 <|endoftext|>，分词器拒绝了它。现在我要玩得随意一点，暂时忽略这一点。）

现在让我们看看它是如何添加的：

In [56]: fw_with_len[0].keys()
Out[56]: dict_keys(['text', 'id', 'dump', 'url', 'date', 'file_path', 'language', 'language_score', 'token_count', 'tok_len'])

In [57]: fw_with_len[0]["tok_len"]
Out[57]: 142

In [58]: len(fw_with_len["tok_len"])
Out[58]: 14868862

In [59]: fw_with_len["tok_len"][0]
Out[59]: 142

很好！我们添加了一个 tok_len 列，包含每一行的 GPT-2 词元数量，我们可以提取相当于那些值的列表。让我们将它们绘制为直方图。

尝试直接进行——即，只是进行

ax.hist(fw_with_len["tok_len"], bins=bins)

...似乎让 MatPlotLib 非常不高兴，我的解释器因 OOM 崩溃——我认为它可能试图一次性将所有数据集——文本、ID 等——加载到 RAM。

所以我启动了一个新的并再次做了加载和用词元长度标注它的东西——奇怪的是，这一次映射只花了大约 10 秒钟！这很奇怪，我需要研究一下。也许之前的命令将 tok_len 列添加到磁盘上的文件？

为了解决内存问题，我从数据集将 tok_len 列转换为实际列表：

In [11]: lengths = [n for n in fw_with_len["tok_len"]]

这花费了十到二十秒。让我们再试一次绘图（这次完整代码）：

In [19]: import numpy as np
    ...: import matplotlib.pyplot as plt
    ...:
    ...: bins = np.arange(0, 2048 + 16, 16)
    ...:
    ...: plt.xkcd()
    ...: plt.rcParams['font.family'] = "xkcd"
    ...: fig = plt.figure(figsize=(10, 6))
    ...: ax = plt.gca()
    ...:
    ...: ax.hist(lengths, bins=bins)
    ...: ax.set_xlabel("TOKENIZED LENGTH (GPT-2 TOKENS)")
    ...: ax.set_ylabel("COUNT")
    ...: ax.set_title("FINEWEB DISTRIBUTION OF TOKENIZED LENGTHS")
    ...:
    ...: mean_len = float(np.mean(lengths))
    ...: median_len = float(np.median(lengths))
    ...: h_mean = ax.axvline(mean_len, linestyle="--", label=f"MEAN = {mean_len:.1f}")
    ...: h_med  = ax.axvline(median_len, linestyle=":",  label=f"MEDIAN = {median_len:.1f}")
    ...: ax.legend(handles=[h_mean, h_med])
    ...:
    ...: ax.grid(True, axis="y", alpha=0.3)
    ...: plt.tight_layout()
    ...: plt.savefig("fineweb-token-length-distribution.png")

这花费了大约 11 秒来运行，结果是这样：

这真的很有希望！大多数都少于我们的 1,024 词元序列长度。1 如果我们将数据集中的每一行作为独立训练样本呈现，必要时裁剪它们，也许我们不会丢失太多数据？让我们看看。

第一步，总共有多少词元？

In [20]: sum(lengths)
Out[20]: 10336315397

很好，约 10B，如预期。如果我们将它们裁剪到默认的 GPT-2 上下文长度 1,024，我们会有多少词元？

In [21]: sum(l if l < 1024 else 1024 for l in lengths)
Out[21]: 7354541756

哎哟，7.3B。这是相当大的减少：

In [22]: 7354541756 / 10336315397
Out[22]: 0.7115245107685639

所以通过那种裁剪我们损失了 29% 的词元。这是由于截断仅 16% 的序列：

In [26]: len([l for l in lengths if l > 1024])
Out[26]: 2438899

In [27]: len(lengths)
Out[27]: 14868862

In [28]: 2438899 / 14868862
Out[28]: 0.1640272806351959

数据处理策略

我感觉我们有两个选择：

选项 1： 裁剪所有输入序列——即，数据集中的每一行——使得每一行的长度不超过我们的 1,024 序列长度。然后我们可以用序列结束词元（这是标准做法）填充它们，使它们都是 1,024。这会让我们损失相当多的词元，但有一个大优势是很容易。

选项 2： 将语料库视为本质上一个长文档，在每一行之间有序列结束分隔符，然后将其分成 1,024 词元序列。这样做意味着我们会使用所有训练数据。但它会更复杂，特别是如果我们遇到内存限制。

在实验的这一点上，我会让两个选择都开放。我倾向于后者（我相信它更接近真正的 GPT-2 训练所做的），但我不确定。

无论如何，我们在这里进行范围界定，所以让我们继续。

训练轮数和词元数计算

查看完数据后，我对此更多地思考了。我之前一直在考虑训练的所有词元；我们会通过数据集中的 10B 词元工作，然后我们完成。

但在训练模型时，你通常会进行多个 epoch——你运行一次通过数据集，在进行时更新梯度，然后同样地再次运行，最终当验证损失开始上升时停止。

我认为因为我读过 LLM 现在通常只在一个 epoch 上训练，我有点将其内化认为我们只需要做一个。但在 2019 年 GPT-2 出现时情况并非如此。他们的数据更少——仅约 10B 词元，相比之下数据集巨大得多，如完整 FineWeb（不是我们一直在查看的 10B 的——完整的 18.5T 一个），所以他们会为某些数量的 epoch 训练它。

多少个？这是另一种情况，GPT-2 论文令人恼火地不够详细。这份报告在"复制 GPT-2"部分说 OpenAI 为其训练了 800k 次迭代，批大小为 512。代入序列长度 1024，这给了我们这么多词元：

现在，如果我们相信他们的数据集是 10B 词元，那么我们可以计算出那相当于多少个 epoch：

同一份报告说他们——即报告的作者——使得这种说法"大约总共 60 个 epoch 通过训练集"——我相信他们谈论的训练集可能略短于原始 GPT-2 的——GPT-2 作者没有发布他们自己的，称为"WebText"，所以报告的作者使用的是不同的试图复制它的一个，OpenWebText。

这听起来很昂贵；即使不知道我们可以训练多少词元每秒，40 多个 epoch 的 10B 词元每个听起来会花费很长时间。是否有任何其他比较点可能告诉我们应该训练多长时间？

嗯，有一个我听说过的"Chinchilla 启发式"，它说你应该训练大约 20 个词元每模型参数。我花了一些时间阅读这来自哪里；最初它在 Google DeepMind 的"训练计算最优大型语言模型"中，这是一篇有趣的论文，读起来惊人地容易，有一些变得有点复杂的数学位（但不需要理解论文想说什么的很好的感受）。我推荐你看看。

它是在 2022 年写的，作者感觉人们在大幅扩展模型，但没有充分增加用于训练的词元数量。所以，他们训练了大量的模型，试图回答这个问题："考虑到特定的训练 FLOP 预算，训练词元与参数的最优平衡是什么，以确保你最有效地使用那些 FLOP？"。他们在反对某篇论文中采取的方法，另一个团队在那里训练了一个模型（称为 Gopher）所用词元少于他们认为最优的。

用于训练模型的 FLOP 数量与参数数量和训练所用词元数都是线性的，所以如果你获得你以前有的 FLOP 的 2 倍，你可以在相同的模型上训练两倍多的词元，或者你可以将其大小翻倍。哪个更好？他们的结论是你实际上应该同时将参数和词元按相同数量进行缩放——即，在 2 倍的情况下你会想要参数和词元都 2 倍，这将使你的 FLOP 翻倍并获得更好的性能。

正如你可能看到的，通过这样做他们间接计算出了为特定大小的模型训练的最优词元数。他们没有自己陈述"20x"启发式，但在论文的表 3 中这很清楚，他们给出了许多模型大小和每个的最优词元数。

现在，这个数字不是你需要为一个特定数量的参数获得最好的模型而训练的词元数；给定大小的模型总能训练更多，（希望）会变得更好。但它告诉你何时你已经训练了足够多的词元，以至于你可以通过训练一个比现在更大的模型来获得更好的结果。

他们隐含地假设模型可以尽可能大，这当然不是真的——在现实中，你将针对特定的模型大小，可以适应你的训练硬件大小（或更可能使用生产模型，可以适应你的计划推理硬件的大小）。

但有趣的是，查看 Karpathy 的 nanochat 项目的 README.md，他在 38B 词元上训练了他的 1.9B "d32" 模型——正好是 20 倍。而如果你查看同一个仓库中的 speedrun.sh 脚本，他明确说他正在为 d20 更小的模型训练 20 倍参数：

# The d20 model is 561M parameters.
# Chinchilla says #tokens = 20X #params, so we need 561e6 * 20 = 11.2B tokens.

如果 Andrej Karpathy 认为为 Chinchilla 最优性训练是正确的做法，那么我是谁来反对呢？;-)

更认真地说，也许数据集的更高质量使这成为合理的事情。从 GPT-2 论文，他们描述如何获得数据：

...我们创建了一个强调文档质量的新网络爬虫。为此，我们只爬取了被人类策划/过滤的网页。手动过滤完整网络爬虫异常昂贵，所以作为起点，我们爬取了 Reddit（社交媒体平台）的所有出站链接，该平台至少获得了 3 个业力。这可以被认为是其他用户发现该链接有趣、有教育意义或仅仅有趣的启发式指示符。

这是一个聪明的技巧，但我相信 FineWeb 比他们从中获得的 WebText 数据集过滤和改进得更加仔细。回到 2019 年，他们必须从零开始做所有事情——找到适当的方式获取数据、过滤它等等。现在我们可以只是从 Hugging Face 下载东西。所以也许 Chinchilla 最优是足够的。

无论如何，我们有 163,009,536 个参数，所以基于那个基础，让我们训练：

...词元。（从现在开始我只会使用 3.2B，但那是我的意思的实际数字。）

那非常酷！我们的 FineWeb 10B 样本中已经有比那更多的词元，所以我们可以做一个单个 epoch 的训练运行。

所以问题是——那即使在我的硬件上可行吗？

这一切都取决于我们能以多少个词元每秒训练。一个检查这个的好方法是写一个一次性"训练器"。我们可以用那来计算出 RTX 3090 的 24 GiB VRAM 上的最大批大小，然后通过一堆批次运行——每个进行前向和后向通过——看看我们得到多少。

这当然不会估计我们花在验证模型上的多少时间。但我的直觉告诉我，我们应该花不超过 5% 的训练时间运行验证，所以我们稍后可以做类似的测试，eval 模式，只进行前向通过没有梯度跟踪，并用那来计算训练集中应该有多少词元。

所以，让我们估计训练速度。这段代码在不同的批大小获得词元/秒的估计。希望它足够清楚不需要深入解释。一个大纲：

我们从 FineWeb 加载足够的 GPT-2 词元用于 NUM_BATCHES 的 MAX_BA