llm.c：比 PyTorch 快 7% 的 LLM 推理引擎

Andrej Karpathy（@karpathy）

llm.c 的第 24 天：现在，我们已经可以直接用约 3000 行 C/CUDA 代码，在多 GPU 上以 bfloat16 和 flash attention 训练模型，而且速度非常快！🚀

我们的运行速度比 PyTorch nightly 快约 7%，并且没有任何附加条件。也就是说，这个基线包含了所有现代化的标准配置：混合精度训练、torch compile、flash attention，以及手动填充词表。（之前的比较会带有诸如“仅限推理”或“仅限 fp32”等附加条件。）

与当前的 PyTorch 稳定版 2.3.0 相比，llm.c 实际上快了约 46%。

我做这些对比只是为了说明“llm.c 很快”，并不是想贬低 PyTorch。PyTorch 能以完全通用的方式实现如此快的训练，真的非常了不起：你几乎可以随意构建并运行任意神经网络，还能把它们跑在大量不同的平台上。

在我看来，这两个项目的目标、优点和缺点各不相同，甚至可以说是互补的。事实上，我最初开发 llm.c 时，考虑的是用它配合即将推出的教学视频，解释 PyTorch 在底层替你完成了哪些工作。

过去约一周半，我们是这样走到今天的——新增：

• ✅ 混合精度训练（bfloat16）
• ✅ 大量 kernel 优化，例如 FusedClassifier。与当前的 torch.compile 不同，它不会把归一化后的 logits 实际生成并存入内存
• ✅ flash attention（目前来自 cudnn）
• ✅ Packed128 数据结构，强制 A100 使用 128-bit load（LDG.128）和 store（STS.128）指令

现在也可以进行多 GPU 训练了——新增：

• ✅ 使用 MPI+NCCL 实现的首个多 GPU 训练版本
• ✅ 使用 NVIDIA Nsight Compute 对完整训练过程进行性能分析
• ✅ ZeRO 第一阶段（optimizer state sharding）的 PR 即将合并

我们目前“仍然只有”3000 行 C/CUDA 代码。它确实变得没那么简单了，但相比约 300 万行代码，还是要好一些。

我们还把 fp32 代码库拆分到了单独的文件中。它将只包含纯 CUDA kernel，不使用 cublas、cudnn 等。我认为，这会非常适合作为一门 CUDA 课程的最终项目。

你可以从 gpt2.c 的纯 CPU 实现开始，在课程结束时看看：仅使用 kernel、不依赖任何外部库，能在 GPU 上把它优化到多快。

我们现在的目标，是打造一套可靠、整洁、经过测试、最小化、足够稳健且充分优化的 LLM 技术栈，直接使用 C/CUDA 复现 GPT-2 系列从 124M 到 1.6B 的所有模型规模。

更多详细信息请参阅：“State of the Union [May 3, 2024]”。

2024 年 5 月 3 日下午 6:22 · 83.92 万次浏览 · 1726 · 106 · 6500 · 2000

Joby（@Joby_Fi）
2024 年 5 月 3 日

读完这些，我才意识到，原来我对 ML 一窍不通。

121 · 19 · 416K

Michael Malek（@mikegmalek）
2024 年 5 月 3 日

对于 fused cross-entropy，可以通过以下方式降低峰值内存占用：

1. 对 logit 计算进行分块，避免生成完整的 logits tensor；
2. 在训练期间，直接用 logits 的 grad 覆盖 logits。

特别是在 vocab size >> dim 时，可以节省大量内存！

GitHub - mgmalek/efficient_cross_entropy

410 · 90 · 25K

sankalp（@dejavucoder）
2024 年 5 月 3 日

623.7K

加入对话

阅读另外 169 条回复