1. Llm.c – LLM training in simple, pure C/CUDA (github.com)
项目概述
llm.c 是一个使用纯 C/CUDA 代码实现的大语言模型(LLM)训练框架,旨在避免 PyTorch 等大型依赖(245MB PyTorch, 107MB cPython)。其当前核心目标是预训练,特别是复现 GPT-2 和 GPT-3 系列模型。项目同时提供了一个并行的 PyTorch 参考实现(train_gpt2.py,基于作者的 nanoGPT 改进)。llm.c 在性能上略快于 PyTorch Nightly(约7%)。
代码结构
- 核心代码:项目维护纯 C 和 CUDA 代码。
train_gpt2.cu:主开发线,包含前沿的、高性能的 CUDA 实现。train_gpt2.c:约1000行的、简洁的 CPU 纯 fp32 参考实现,便于理解和学习。
- 开发与测试:
dev/目录包含用于开发、测试的内核库(如手动编写的 CUDA 内核)、数据处理脚本和可视化工具。 - PyTorch参考:
train_gpt2.py用于生成测试基准数据和作为对比参考。
快速启动 项目提供了多种入门方式:
- GPU 训练(fp32 单GPU):
- 运行
./dev/download_starter_pack.sh下载预训练模型、分词器和数据集。 - 执行
make train_gpt2fp32cu && ./train_gpt2fp32cu编译并运行。 - 这是一个早期版本的“快照”,更简单、更易移植。
- 运行
- CPU 训练:
- 同样先运行 starter pack 脚本。
- 执行
make train_gpt2编译,然后用OMP_NUM_THREADS=8 ./train_gpt2运行。 - 此为更早期的参考实现,适用于在没有 GPU 的环境下进行演示或小规模微调。
- 手动准备数据:可以通过运行
pip install -r requirements.txt和python dev/data/tinyshakespeare.py等命令从头生成所需数据。
数据集与测试
- 数据集:
dev/data/中的 Python 脚本负责下载、分词并保存数据为易于 C 读取的.bin文件(包含头信息和 uint16 的 token 流)。 - 单元测试:
- CPU 测试:
make test_gpt2 && ./test_gpt2,加载 PyTorch 生成的状态文件,验证 C 代码的前向传播和训练损失与 PyTorch 一致。 - GPU 测试:
make test_gpt2cu && ./test_gpt2cu(支持 fp32 和混合精度路径)。
- CPU 测试:
多GPU与多节点训练
- 多GPU训练:需安装 MPI 和 NCCL,然后通过
mpirun -np <GPU数量> ./train_gpt2cu或使用./scripts/下的脚本运行。 - 多节点训练:支持三种初始化 NCCL 的方式:
- 使用 OpenMPI 交换 NCCL ID。
- 使用共享文件系统。
- 使用 TCP 套接字。
项目提供了对应的脚本(如
run_gpt2_124M_mpi.sh),并需根据 Slurm 环境(是否支持 PMIx)选择合适方法。
关键功能与集成
- Flash Attention:项目支持通过 cuDNN 使用 Flash Attention(自2024年5月起),但默认禁用。可通过
make train_gpt2cu USE_CUDNN=1启用,需安装 cuDNN。 - cuBLAS:矩阵乘法默认使用 cuBLAS,以实现高性能。
- 教程:
doc/layernorm/layernorm.md提供了一个分步实现 GPT-2 LayerNorm 层的小教程。
项目理念
- 教育:
dev/cuda目录旨在成为一个内核库,包含从简单到复杂的、有详细文档的手写内核,用于教学。 - 性能:同时也追求极致性能,计划复现大型 GPT-2(1.6B)训练,会集成 cuBLAS、CUTLASS、cuDNN 等快速库作为性能上限参考。
- 简洁性约束:项目根目录的主代码文件保持简单易读。对于显著增加复杂度但性能提升有限的 PR(如增加500行复杂代码仅提速2%),可能会被拒绝。复杂实现可放在
dev/目录。
相关语言端口 项目鼓励其他语言的移植,但建议放在独立仓库。列出的端口包括:AMD支持、C#、CUDA C++、C++、WebGPU C++、Go、Java、Metal、Mojo、OpenCL、Rust、Swift、Zig、Habana Gaudi2、Nim等。主仓库会链接这些“notable forks”。
开发与许可
- 沟通渠道:开发协调通过 GitHub Discussions 和 Discord(
#llmcon Zero to Hero,#llmdotcon GPU MODE)进行。 - 问题与贡献:具体问题用 Issues,代码贡献用 PR。
- 许可证:MIT。