刚刚Andrej Karpathy宣布使用纯c代码训练大模型项目又迎来重大更新，项目名称：
llm.c
由人工智能大神Andrej Karpathy发起的llm.c 项目，其团队在多GPU训练、混合精度计算和注意力机制优化等多个方面取得了重大突破
最新进展显示,llm.c 在单GPU上的训练速度已经超过了 PyTorch 最新版本约7%,且包含了诸如混合精度训练、PyTorch编译和闪电注意力等现代标准功能。相较于 PyTorch 2.3.0 稳定版本,llm.c 的训练速度甚至快了约46%
Karpathy指出,他的目的并非是贬低 PyTorch,而是证明 llm.c 确实拥有出色的速度表现。他认为两个项目的目标和优缺点有所不同,但又是互补的。实际上,他最初启动 llm.c 项目的目的是为了制作一些即将推出的教育视频,解释 PyTorch 在底层是如何运作的
在过去约1.5周的时间里,llm.c 项目在以下几个方面取得了重大进展:
✅ 实现了 bfloat16 混合精度训练
✅ 优化了多个内核,包括一个比 PyTorch 编译器更精简的融合分类器(不需要显式计算归一化的 logits)
✅ 集成了来自 cuDNN 的闪电注意力机制
✅ 引入了 Packed128 数据结构,强制利用 A100 GPU 的 128位加载(LDG.128)和存储(STS.128)指令,以最大化内存带宽
此外,llm.c 现在也支持多GPU训练:
✅ 实现了基于 MPI+NCCL 的第一版多GPU训练
✅ 针对整个训练运行进行 NVIDIA Nsight Compute 分析
✅ 即将合并 ZeRO 优化器状态切分(Stage 1)的 PR
令人惊叹的是,在实现上述所有优化的同时,llm.c 的代码量仍然只有约3000行 C/CUDA 代码。虽然复杂度有所增加,但远远小于 PyTorch 的约300万行代码。同时,llm.c 项目团队还将纯 fp32 代码分离到了单独的文件中,该文件只包含纯 CUDA 内核,没有 cuBLAS、cuDNN 等依赖,可以作为 CUDA 教程的绝佳实践项目
目前,llm.c 项目的目标是创建一个可靠、简洁、经过测试、精简、加固和充分优化的大型语言模型堆栈,能够在 C/CUDA 中复现 GPT-2 从 124M 到 16 亿参数的所有模型尺寸。按照目前的进度,卡帕蒂预计这个目标将在大约2周内实现
不过,项目中也出现了一些令人沮丧的情况。nvcc 编译时间已经从之前的2.4秒增加到4.3秒,与 PyTorch 的导入时间相当。更糟糕的是,启用闪电注意力会导致编译时间飙升至1分24秒,这是一个前所未有的较大拖延。团队正在积极探索如何克服这一问题
另一个积极的进展是,通过精心管理内存分配和使用方式,特别是在融合分类器中,llm.c 的峰值内存使用量已大幅降低。在批量大小为32、序列长度为1024的设置下,llm.c 的内存占用为16.6 GB,而 PyTorch 为37.2 GB(不过后者可能存在一些异常)
结语
此外,还有三个基于 llm.c 的分支项目值得关注:@gevtushenko 的 llm.cpp(使用 CUDA C++ Core Libraries)、@saimirbaci 的 llm.zig(Zig 语言实现)和 @joshcarp 的 llm.go(Go语言实现)