顶级研究者开源项目展示纯 C 实现高效 LLM 训练,深度揭示模型训练细节,是学习 LLM 的黄金资源。
通知:你必须登录才能更改通知设置
加载时出错。请重新加载此页面。
加载时出错。请重新加载此页面。
karpathy 2024 年 7 月 11 日 维护者
在这篇文章中,我们将使用 llm.c 复现 GPT-2。这里说的是“真正的 GPT-2”,即 OpenAI 于 2019 年 2 月 14 日在博文《Better Language Models and their Implications》中发布的完整 1558M 参数版本。llm.c 直接使用 C/CUDA 实现这一过程(总计约 5,000 行代码),没有使用通常包含 Python 解释器以及 PyTorch/JAX、huggingface/transformers 等复杂度高得多的深度学习库的训练技术栈。2019 年,训练 GPT-2 是一个需要整个团队投入的复杂项目,也被视为一次大模型训练;但大约 5 年后的今天,得益于计算硬件(H100 GPU)、软件(CUDA、cuBLAS、cuDNN、FlashAttention)和数据(例如 FineWeb-Edu 数据集)的进步,我们可以在单个 8×H100 节点上用 24 小时、花费 672 美元复现这个模型,这相当不可思议。不过,其中也有一些需要说明的限制和注意事项:llm.c 仍未得到完美调优,也尚未充分稳定(我们偶尔仍会看到损失尖峰和异常的激活值范围),我们的评估也并不全面(例如,我们没有仔细评估多语言、代码和数学能力);其中很多工作仍在进行中。
独角兽。最自然的切入点大概就是独角兽。在 GPT-2 博文中,OpenAI 使用了下面这段独角兽提示词来测试 GPT-2。让我们看看我们的模型如何理解安第斯山脉中会说英语的独角兽:
提示词:一项令人震惊的发现表明,科学家在安第斯山脉一处偏远且此前从未被探索过的山谷中,发现了一个独角兽群。更令研究人员惊讶的是,这些独角兽能够说一口流利的英语。续写:这听起来可能像一场骗局,或者只是关于这些神话动物的一场梦,但这似乎可能是真实存在的独角兽语言案例。安第斯地区矿产丰富,生物多样性也非常高。根据研究人员的研究,这些独角兽的语言名为 Elveseo。为了发出这种语言,动物先用嘴唇形成一个辅音,再用牙齿发出一个元音,然后用嘴巴发出一个辅音,并从胸腔中发出“ow”的声音。研究合著者 Paul Marek 评论道:“问题在于,是什么导致了这种现象?它又揭示了哪些有关智能和交流的信息?”视频中的独角兽掌握了正确的说话方式,这些动物能够交流,并能用流利的英语交谈。
很有趣!:) 这个模型的输出相当连贯,从直观感受来看,大致达到了 GPT-2 的水平。你可以在这里找到 GPT-2 和 llm.c 模型各自生成的 20 个样本,也可以按照下文的说明生成更多样本。
训练。使用 llm.c 训练 GPT-2 非常简单,因为它是用 C/CUDA 编写的,所以不需要 minconda、Python、PyTorch 等。你需要一台配备 8×H100 GPU 的机器,我建议从 Lambda Labs 启动一台。不过,llm.c 对计算资源的适应性很强——如果你只有 1 块 GPU,仍然可以训练出自己的 GPT-2,只不过需要等待 8 天,而不是 1 天。如果你有 16 块 GPU(例如使用新的 Lambda 1 Click Clusters),就可以进行多节点训练,只需等待 12 小时。启动节点后,按照下面的完整步骤即可训练 GPT-2(从一台空白机器到开始执行训练步骤只需约一分钟):
# install cudnn so we can use FlashAttention and run fast (optional)
# https://developer.nvidia.com/cudnn-downloads
# for me, CUDA 12 (run `nvcc --version`) running on Linux x86_64 Ubuntu 22.04
wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/x86_64/cuda-keyring_1.1-1_all.deb
sudo dpkg -i cuda-keyring_1.1-1_all.deb
sudo apt-get update
sudo apt-get -y install libcudnn9-dev-cuda-12
# "install" cudnn-frontend to ~/
git clone https://github.com/NVIDIA/cudnn-frontend.git
# install MPI (optional, if you intend to use multiple GPUs)
# (you might also have to install NVIDIA NCCL if it doesn't come with your setup)
sudo apt -y install openmpi-bin openmpi-doc libopenmpi-dev
# download and enter llm.c repo
git clone https://github.com/karpathy/llm.c.git
cd llm.c
# download the "starter pack" (~1GB download)
# contains GPT2-124M weights (used in tests), tokenizer, eval data .bin s
./dev/download_starter_pack.sh
# download the training dataset (FineWeb-Edu 100B token) .bin data shards
# note: this is a total of 1001 data shards. If you only want to test things
# out and don't want to do an actual run, feel free to append the number of
# training shards to download (e.g. for just 10 shards: ./edu_fineweb.sh 10)
# the full dataset is ~200GB, we can store it here in dev/data directory.
cd dev/data
./edu_fineweb.sh
# compile (~1 min 1st time for cuDNN mostly, few sec from then on)
cd ../../
make train_gpt2cu USE_CUDNN=1
# and train! (wait 24 hours here)
mpirun -np 8 ./train_gpt2cu \
-i "dev/data/edu_fineweb100B/edu_fineweb_train_*.bin" \
-j "dev/data/edu_fineweb100B/edu_fineweb_val_*.bin" \
-o "log_gpt2_1558M" \
-v 250 -s 300000 -g 384 \
-h 1 \
-b 16 -t 1024 \
-d 1048576 \
-r 0 \
-z 1 \
-c 0.1 \
-k "cosine" \
-l 0.0006 \
-q 0.1 \
-u 700 \
-n 2000 \
-x 32000 \
-ge 1 \
-y 1 \
-e "d48"
稍后我会介绍这些参数。你会看到大量输出不断滚动,随后优化过程便会开始:
num_parameters: 1557686400 => bytes: 3115372800
allocated 2971 MiB for model parameters
batch_size B=16 * seq_len T=1024 * num_processes=8 and total_batch_size=1048576
=> setting grad_accum_steps=8
created directory: log_gpt2_1558M
allocating 40409 MiB for activations
val loss 11.129390
allocating 2971 MiB for parameter gradients
allocating 742 MiB for AdamW optimizer state m
allocating 742 MiB for AdamW optimizer state v
allocating 742 MiB for master copy of params
step 1/32000 | loss 11.133732 (+nanz)| norm 52.9732 (+nanz)| lr 8.57e-07 | 3056.36 ms | 42.6% bf16 MFU | 343080 tok/s
step 2/32000 | loss 10.539388 (+nanz)| norm 43.5996 (+nanz)| lr 1.71e-06 | 2747.19 ms | 47.4% bf16 MFU | 381690 tok/s
step 3/32000 | loss 9.894109 (+nanz)| norm 23.2229 (+nanz)| lr 2.57e-06 | 2753.25 ms | 47.3% bf16 MFU | 381259 tok/s
step 4/32000 | loss 9.566241 (+nanz)| norm 28.4920 (+nanz)| lr 3.43e-06 | 2741.47 ms | 47.5% bf16 MFU | 381690 tok/s
step 5/32000 | loss 9.482848 (+nanz)| norm 23.7817 (+nanz)| lr 4.29e-06 | 2752.07 ms | 47.3% bf16 MFU | 381507 tok/s
step 6/32000 | loss 9.332832 (+nanz)| norm 15.9113 (+nanz)| lr 5.14e-06 | 2751.01 ms | 47.3% bf16 MFU | 381431 tok/s
step 7/32000 | loss 9.165650 (+nanz)| norm 10.5941 (+nanz)| lr 6.00e-06 | 2753.03 ms | 47.3% bf16 MFU | 381327 tok/s
step 8/32000 | loss 9.132234 (+nanz)| norm 16.2733 (+nanz)| lr 6.86e-06 | 2748.91 ms | 47.3% bf16 MFU | 381348 tok/s
step 9/32000 | loss 9.097384 (+nanz)| norm 12.1342 (+nanz)| lr 7.71e-06 | 2748.73 ms | 47.3% bf16 MFU | 381367 tok/s
step 10/32000 | loss 9.072879 (+nanz)| norm 10.5923 (+nanz)| lr 8.57e-06 | 2749.40 ms | 47.3% bf16 MFU | 381369 tok/s
...
可以看到,每个训练步骤大约需要 2.75 秒,总共有 32,000 步,所以现在要等待约 24 小时。在每一步中,这次训练都会取出 FineWeb-EDU 中约 100 万个 token 的一批数据(这些数据来自互联网上的教育类网页),并更新模型的 15.58 亿个权重,使模型在预测序列中的下一个 token 时表现得稍微更好。训练结束时,我们总共将处理 32,000 * 1048576 = 33.6B 个 token。随着模型越来越擅长预测下一个 token,损失会不断下降。范数会稳定在 0.1~1 左右,学习率则会在最初几个步骤中逐渐预热。我们的模型 FLOPS 利用率(MFU)约为 50%,也就是说效率相当高。
现在等待 24 小时,让训练完成。之后,你可以使用 dev/vislog.ipynb Jupyter Notebook 将 main.log 日志文件可视化。为此,你还需要安装 Python 和 matplotlib,随后会看到如下结果:
评估。左侧显示的是 FineWeb-EDU 验证数据上的损失。如果直接运行 OpenAI 发布的 GPT-2,并在这个数据划分上评估其损失,就会得到红色水平线(损失为 2.83)。可以看到,我们的训练非常快就超过了它,大约在第 5,000 步。不过,这并不是公平的比较,因为 GPT-2 是在从未公开的 WebText 数据集上训练的,因此可能存在很大的分布偏移。例如,如果以 1e-4 的学习率对 OpenAI 模型微调 1,000 步,损失会迅速降至蓝线(损失为 2.61),因为模型正在快速适应新数据的统计特征。我喜欢把验证损失作为健全性检查,但要进行真正的比较,我们应该查看固定的第三方评估。其中一个表现稳定、曲线平滑、常用且经常被引用,同时还能提供早期信号的评估是 HellaSwag。它包含一些简单的常识场景,模型必须选出正确的后续内容。我们在右侧面板中评估 HellaSwag,可以看到大约在第 25K 步时超过了 GPT-2 模型(早于 GPT-2;据估计,GPT-2 使用了约 100B 个 token 进行训练。这可能与数据质量的提升有关,我们之前训练 124M 模型时也观察到了这一点)。绿线表示相同规模的 GPT-3 模型,其模型架构与 GPT-2 基本相同,仅有少量差异(上下文长度从 1024 增加到 2048),但使用了 300B 个 token 进行训练(即大约是我们这里训练 token 数量的 10 倍)。需要说明的是,即使 HellaSwag 也不是理想的单点比较指标,因为它测试的是简单英语和常识,并不测试多语言、数学或代码等能力。WebText 的数据混合中可能更侧重这些领域,而这些领域在某种程度上“占用”了模型容量;由于该数据集从未公开,我们无从得知。最后,总体而言,在 GPT-2 这种模型能力较低的情况下,很难设计出良好的评估,因为模型甚至可能无法理解选择题,其生成样本的质量也不足以在标准数学或代码评估中取得显著高于随机水平的成绩。
参数指南。现在让我们更详细地看看传给训练程序的参数。OpenAI 发布 GPT-2 时提供了模型权重,但几乎没有提供细节;而发布 GPT-3 时没有提供权重,却给出了大量细节。因此,在许多情况下,我们遵循 GPT-3 论文中的超参数,因为 GPT-2 论文提供的信息非常少:
mpirun -np 8 ./train_gpt2cu \ 是启动命令:我们使用 mpi 启动 8 个进程(每个进程在 1 块 GPU 上运行训练,因此在这个示例的 8XH100 节点上共使用 8 块 GPU)。如果你有 4 块 GPU,请使用 -np 4。如果你只有 1 块 GPU,可以跳过 mpi,也就是直接将它改为 ./train_gpt2cu。
-i -j 是训练集和验证集的 token 文件,之前已经通过 edu_fineweb.sh 下载。
-o 是用于写入日志和检查点的输出目录。
-v 250 表示每 250 步评估并记录一次验证损失。
-s 300000 表示每 300000 步采样一些 token。由于总步数会小于这个数字,这是一种有点取巧的关闭采样方式,我们只会在最后进行一次采样。
-g 384 将最终采样的 token 数量设置为 384。
-h 1 表示评估 HellaSwag 准确率。
-b 16 将微批次大小设置为 16。如果显存不足,请减小这个值,例如依次尝试 8、4、2,必要时一直降到 1。
-t 1024 将最大序列长度设置为 1024,与 GPT-2 相同。
-d 1048576 要求总批次大小为 2 的 20 次方,遵循 GPT-3 论文中的超参数表。代码会确保达到所需的总批次大小,并计算优化所需的梯度累积“内循环”步数。例如,在上面的例子中,我们有 8 块 GPU,每块处理 16 X 1024 个 token,因此每个微步(一次前向传播和反向传播)处理的 token 数量为 8 X 16 X 1024 = 131,072。因此,代码计算出需要进行 8 步梯度累积,才能在每一步达到所需的 1M 批次大小。也就是说,它会执行 8 次前向传播和反向传播,然后进行一次参数更新。
-r 0 将重计算设置为零。重计算是一种在计算量和显存占用之间进行权衡的方法。如果使用 -r 1,我们会在反向传播期间重新计算前向传播的一部分(GeLU)。这意味着不必缓存它,从而节省显存,代价是增加一些计算量。因此,如果显存不足,可以尝试 -r 1,或者 -r 2(同时重新计算 layernorm)。
-z 1 会在多块 GPU 之间启用 ZeRO-1(即优化器状态分片)。如果使用超过 1 块 GPU 训练,这个设置显然是正确选择,基本上应该始终开启。在单块 GPU 上,这个设置不起作用。
-c 0.1 将权重衰减设置为 0.1。与 GPT-2 完全相同,只有(二维)权重会应用衰减,而这个数值来自 GPT-3 论文。
-k "cosine" 设置余弦学习率调度。由于它本来就是默认值,所以这个参数有点多余。
-l 0.0006 将最大学习率设置为 6e-4。GPT-3 论文建议对这一模型规模使用 2e-4,但这里我们将其提高了三倍,看起来训练得更快,而且没有出现任何问题。不过目前还没有对此进行非常仔细的调优。
-q 0.1 表示在整个训练过程中,将学习率衰减到最大学习率的 10%,遵循 GPT-3 论文。
-u 700 表示在最初 700 次迭代中,将学习率从 0 逐步提升到最大学习率。按照 0.5M 的总批次大小计算,这相当于 350M 个 token,遵循 GPT-3 论文。
-n 2000 表示每 2000 步保存一次模型检查点。
-x 32000 表示总共训练 32K 步。我选择这个数字,是因为它比较整齐,而且刚好能控制在 24 小时以内。
-ge 1 启用最近刚合并的 CublasLt gelu 重计算设置(可选)。
-y 1 开启“恢复训练”标志。如果训练因任何原因崩溃或卡住,可以按 CTRL+C,然后重新运行此命令,程序会尝试恢复优化过程。llm.c 具有逐位确定性,因此得到的结果与未发生崩溃时完全相同。
-e "d48" 表示从头初始化一个深度为 48 层的 GPT-2 模型。
显存指南。大多数人可能面临的最大限制,是他们的 GPU 没有 80GB 显存。没关系,只要有耐心,你应该仍然可以运行上面的所有内容,只是速度会更慢。那么,如果模型装不下,应该调整哪些参数?最重要的是微批次大小 -b。尝试减小它,但尽量使用规整的数字,例如 16 -> 8 -> 4 -> 2 -> 1。接下来,还可以尝试调整重计算设置 -r:0(最快,但占用大量显存)、1(只会稍微变慢,却能大幅节省显存),或者 2(略慢,进一步节省一些显存)。接下来可以禁用 fp32 主权重,方法是设置 -w 0(默认值为 1)。这样我们就不会维护参数的 fp32 副本。根据此前几次运行的经验,这似乎没有问题,可能是因为我们使用了随机舍入。如果这样仍然装不下(应该不太可能吧?),可以尝试通过 -t 减小最大序列长度。默认值是 1024,可以降到 512、256 等;但这样会降低模型质量,因为你缩短了它的最大注意力跨度。
代码。当然,我觉得自己可能有所偏爱,但 llm.c 确实非常漂亮:
它只需要基础的 CUDA 依赖即可运行。
它是使用 C/CUDA 编写的直接、精简且易读的实现。llm.c 总共大约有 5,000 行 C/CUDA 代码。我们尽量主要使用 C,而不是 C++,以保持简单。神经网络训练只不过是在一个 float 数组上,对相同的简单算术运算(想想 +、-、*、/)执行一个 while 循环,实在不应该那么复杂。
它的编译和运行都非常快(只需几秒),因此你可以把更多时间用在迭代上,而不是等待上。
它只在启动时一次性分配全部 GPU 显存,此后训练期间的显存占用会精确保持不变。因此,一旦开始迭代,你就知道剩余训练过程不会有问题,也不会发生 OOM。
它具有逐位确定性。
它很高效,MFU 略低于约 50%。
主入口以及大部分代码位于 train_gpt2.cu 文件中。该文件用约 2,000 行代码包含了 GPT-2 模型定义和训练循环,并从 llmc 目录导入了一批辅助文件,其中包含各种实用工具和各个层的实现。cloc llmc 报告共有 23 个文件、3170 行代码,而目前 cloc train_gpt2.cu 的结果是 1353 行代码。
多节点训练。如果你属于拥有大量 GPU 的特权上层阶级,llm.c 支持多节点训练,而我见过有人使用 llm.c 训练时动用的 GPU 数量最多约为 500 块。到目前为止,我个人进行过的最大规模训练,是使用 Lambda 新推出的一键式集群功能,在 2 个节点上使用 16XH100 GPU。失业的坏处。Lambda 团队已经发布了详细说明,介绍如何在他们的一键式集群上训练 llm.c 模型。例如,使用每小时 2,300 美元的 512-GPU H100 集群,你也许能在大约 30 分钟内完成 GPT-2 的训练。你需要增大总批次大小(例如增至约 8M),并且可能需要稍微调整一下超参数。我还没有尝试过,但它大概可行,而且一定会非常酷 :)
在 PyTorch 中,相对可比的运行应该看起来像这样,使用我们的并行 PyTorch 实现:
torchrun --standalone --nproc_per_node=8 train_gpt2.py \
--input_bin "dev/data/edu_fineweb100B/edu_fineweb_train_*.bin" \
--input_val_bin "dev/data/edu_fineweb100B/edu_fineweb_val_*.bin" \
--write_tensors 0 \
--model d48 \
--batch_size 8 --sequence_length 1024 --total_batch_size 1048576 \
--dtype bfloat16 \
--compile 1 \
--tensorcores 1 \
--flash 1 \
--num_iterations 32000 \
--warmup_iters 700 \
--weight_decay 0.1 \
--overfit_single_batch 0 \
--learning_rate 0.0006 \
--zero_stage 1
PyTorch 代码旨在作为测试参考而非实际实现,所以训练循环在某些地方略有不同(例如数据加载器不会打乱分片等),但作为参考点仍可能有用。我也修改了默认词汇表大小从 50257 到 50304 以获得更高的效率,然后当前的 PyTorch nightly 给出:
step 16/32000 | train loss 8.903997 | norm 8.3474 | lr 1.37e-05 | (3381.88 ms | 310057 tok/s)
step 17/32000 | train loss 8.870140 | norm 3.7936 | lr 1.46e-05 | (3381.95 ms | 310051 tok/s)
step 18/32000 | train loss 8.875732 | norm 9.4993 | lr 1.54e-05 | (3393.09 ms | 309033 tok/s)
step 19/32000 | train loss 8.817432 | norm 2.8345 | lr 1.63e-05 | (3379.75 ms | 310253 tok/s)
step 20/32000 | train loss 8.798056 | norm 4.1234 | lr 1.71e-05 | (3386.53 ms | 309631 tok/s)
step 21/32000 | train loss 8.777574 | norm 2.8010 | lr 1.80e-05 | (3386.05 ms | 309675 tok/s)
...
现在我不能说我对 PyTorch 脚本的调优程度完全有信心,但可以做出以下观察。PyTorch 似乎消耗了更多的内存(这次运行约 80GB),而 llm.c 在 57GB(提升 29%)。内存很重要,因为它允许你提高批大小(例如 llm.c 这里可以达到 24 个微批),这会更快一些。其次,我们看到每次迭代约 3386 对 2750 毫秒,所以 llm.c 的步进快约 19%。这里的一些优势有已知来源,例如 llm.c 包含的优化如融合分类器启动反向传播,这是 torch.compile 目前据我所知不做的。但这个脚本也可能没有完全最优调整,不过无论如何我展示这个对比是为了 1) 其他人想看看、玩玩、比较、帮助调整,以及 2) 只是想说 llm.c 在 GPT-2/3 训练的具体情况下相当优化和快速。
一些可能有帮助的链接,供参考:
model_00032000.bin llm.c bin 模型文件
转换为 huggingface transformers GPT-2 模型,我上传在这里:karpathy/gpt2_1558M_final2_hf。
我现在也添加了一个训练 100K 步并达到 HellaSwag 57.7 的模型版本,以及 330K 步达到 62.7 的版本。
模型导出可以如下进行,例如:
python dev/eval/export_hf.py --input log_gpt2_128M/model_00032000.bin --output gpt2_1558M_export
这样你可以运行 Eleuther 评估框架,或运行 huggingface 采样管道获取模型样本:
# take model for spin
import torch
output = "./gpt2_1558M_final2_hf"
# set pytorch seeds
torch.manual_seed(42)
torch.cuda.manual_seed(42)
prompt = "In a shocking finding, scientist discovered a herd of unicorns living in a remote, previously unexplored valley, in the Andes Mountains. Even more surprising to the researchers was the fact that the unicorns spoke perfect English."
from transformers import AutoModelForCausalLM, AutoTokenizer
tokenizer = AutoTokenizer.from_pretrained(output)
model = AutoModelForCausalLM.from_pretrained(output, attn_implementation="flash_attention_2", torch_dtype=torch.bfloat16, device_map='cuda')
model.eval()
tokens = tokenizer.encode(prompt, return_tensors="pt")
tokens = tokens.to('cuda')
output = model.generate(tokens, max_new_tokens=500, pad_token_id=tokenizer.eos_token_id, do_sample=True, top_k=50, num_return_sequences=4)
samples = tokenizer.batch_decode(output)
for sample in samples:
print('-'*30)
print(sample)
另外,查看 dev/eval 获取如何运行 Eleuther 评估框架、来自 HuggingFace Open LLM Leaderboard 的评估等的说明。
我也尝试训练 GPT-2 的时间远长于 33B tokens。特别是,我将 -x 改为 400,000 以训练 420B tokens(比相同大小的 GPT-3 模型更多,后者训练了 300B)。这个模型运行在大约第 330,000 步时看起来很好:
这个模型在 HellaSwag 上戏剧性地击败了相同大小的 GPT-2 和 GPT-3(达到约 61%),但遗憾的是在那之后变得不稳定并爆炸。沿途有更多较小的尖峰,但代码配置为检测更简单的瞬间不稳定并跳过更新(我使用了标志 -sl 5.0 -sg 5.0),这有助于减轻和推迟问题。然而,我认为我们在初始化、激活范围和整体模型训练稳定性方面还不够谨慎,存在更深层的问题逐渐将模型推向不稳定,特别是对于更大的模型和长期训练。待续。如果你有稳定 LLM 模型训练的想法或建议,请在下面的讨论中贡献你的经验。
我能从 llm.c 中的模型采样吗? 有点可以,但效率低且有点怪异,如果你想提示模型就更不规范了。现在使用上面的 huggingface 路径。
我能与它聊天吗? 不能,这目前只是预训练,不是聊天微调。
你能用 fp8 训练吗? 不能,我们目前主要用 bf16 训练,但早期版本仍在大量开发中。
我有非 NVIDIA GPU,我能运行 llm.c 吗? 不能,llm.c 仅支持 C/CUDA,但存在不错的分支(见主 README)。例如有一个由 @anthonix 主动维护的 AMD 分支相当不错。
我也想链接到一篇关于在 llm.c 中训练 GPT-2 (124M) 模型的早期文章,其中包含一些与 llm.c 运行相关的更多信息。124M 是 GPT-2 小系列中的一个较小模型,仅 124M 参数,相比之下是 1558M 参数。
对 llm.c 的重大贡献来自现在感觉像是 llm.c 核心开发团队的人,除了我自己:
Lambda labs 赞助了用于 llm.c 开发的 GPU。这里的历史是我已经愉快地使用 Lambda 几年了,几个月前我诚恳地请求他们是否愿意不向我的账户收费用于 llm.c 开发工作,他们同意了。