前言
闲来无事逛掘金,发现一个比较感兴趣的项目 :llm.c -手搓大模型,核心居然是一个大神kaparthy用C写的LLM框架。 在我印象中,不都是python吗,牛*
我们先来翻译下项目的readme,方便后面本地实践
简介
以一种简单,纯c/cuda的方式来训练语言大模型
llm.c
使用简单、纯粹的C/CUDA进行LLM训练。无需安装245MB的PyTorch或107MB的cPython。在单个文件train_gpt2.c中,使用CPU和fp32训练GPT-2大约需要1000行简洁代码,而在GPU上训练则需要大约2000行代码(增加了CUDA内核),参见train_gpt2.cu。这段代码可以立即编译和运行,与PyTorch的参考实现完全匹配,并且其速度大致与(编译后的)PyTorch(fp32,无flash attention)相当。我选择GPT-2作为首个实例,因为它是LLMs的鼻祖,是现代技术堆栈首次整合在一起的时候。
目前,我们正在努力:
- 进一步优化CUDA实现以匹敌/超越PyTorch的速度
- 将精度从fp32降低到混合精度训练
- 添加多GPU训练,从DDP开始
- 重现GPT-2的训练运行(添加数据,评估)
- 更多的现代架构,例如Llama 2、Gemma、Mistral等
我希望这个仓库只维护C和CUDA代码。非常欢迎将此仓库移植到其他语言,但这应该在单独的仓库中进行,然后我很乐意在""notable forks"部分下面链接到它们,就像我在llama2.c notable forks中所做的那样。
快速开始(GPC)
如果你的电脑有gpu模块,那么可以运行如下指令:
bash
pip install -r requirements.txt
python prepro_tinyshakespeare.py
python train_gpt2.py
make train_gpt2cu
./train_gpt2cu
- 使用pip安装requirements.txt中列出的依赖项;
- 执行prepro_tinyshakespeare.py对tinyshakespeare数据集进行预处理和分词;
- 执行train_gpt2.py下载并保存GPT-2(124M)权重;
- 执行make train_gpt2cu在C/CUDA中从这些权重初始化,并在tineshakespeare上训练一个epoch,使用AdamW优化器(批大小为4,上下文长度为1024,共74步),评估验证损失,并采样一些文本。
快速开始(CPU)
运行以下命令:
- 使用pip安装requirements.txt中列出的依赖项;
- 执行prepro_tinyshakespeare.py对tinyshakespeare数据集进行预处理和分词;
- 执行train_gpt2.py下载并保存GPT-2(124M)权重;
- 执行make train_gpt2在C中从这些权重初始化,并在tineshakespeare上训练40步,使用AdamW优化器(批大小为4,上下文长度仅为64),评估验证损失,并采样一些文本。
总的来说,除非你有一个强大的CPU(并且在启动命令中可以增加OMP线程数),否则你不会在CPU上训练大型语言模型(LLMs)走得很远,但这可能是一个不错的演示/参考。
训练说明
下载并标记一个数据集。tinyshakespeare 数据集是下载和标记最快的:
bash
python prepro_tinyshakespeare.py
这将打印:
bash
Saved 32768 tokens to data/tiny_shakespeare_val.bin
Saved 305260 tokens to data/tiny_shakespeare_train.bin
.bin 文件是使用 GPT-2 分词器指示标记 ID 的 int32 数字的原始字节流。或者,您也可以使用 prepro_tinystories.py
标记 TinyStories 数据集。 原则上,我们现在可以准备训练模型了。然而,基线 CPU/fp32 参考代码效率太低,目前还不实际从头开始训练这些模型。相反,我们将使用由 OpenAI 发布的 GPT-2 权重进行初始化,然后进行微调。为此,我们需要下载 GPT-2 权重并将其保存为可以在 C 中加载的检查点:
bash
python train_gpt2.py
这段代码来自nanoGPT,是PyTorch中一个简单的GPT-2参考实现。这个脚本将下载GPT-2(124M)模型,对单个数据批次进行10次迭代的过拟合训练,运行几个生成步骤,最重要的是它将保存三个文件:1)包含用于在C中加载的原始模型权重的gpt2_124M.bin
文件,2)包含更多调试状态的gpt2_124M_debug_state.bin
文件:输入、目标、logits和损失(用于调试和单元测试),最后3)gpt2_tokenizer.bin
文件,用于存储GPT-2分词器的词汇表,将标记ID转换为UTF-8编码的字符串片段。现在我们可以使用这些模型权重初始化,并在原始C中继续训练。首先编译代码:
bash
make train_gpt2
您可以查看 Makefile
及其注释。它将尝试自动检测您的系统上是否有 OpenMP,这对于以非常低的代码复杂性成本加快代码速度非常有帮助。有些人似乎在 Ubuntu 上编译时遇到问题,请参阅 Issue 19,TLDR 您可能需要修改 CFLAGS
:
ini
# 首先尝试这个
CFLAGS="-Ofast -fno-finite-math-only -Wno-unused-result -march=native" make train_gpt2
# 然后尝试这个
CFLAGS="-O3 -Wno-unused-result -march=native" make train_gpt2
一旦编译完成 train_gpt2
,您可以运行它:
bash
OMP_NUM_THREADS=8 ./train_gpt2
您应该根据您的 CPU 核心数量调整线程数。该程序将加载模型权重、标记,运行几次 Adam lr 1e-4 的微调循环,然后从模型生成一个样本。文件(我认为)非常可读,您应该看一看。简而言之,所有层的前向和后向传递都有实现,并且它们被串联在一起形成一个大的、手动的前向/后向/更新循环。在我的 MacBook Pro(Apple Silicon M3 Max)上,输出如下:
yaml
[GPT-2]
max_seq_len: 1024
vocab_size: 50257
num_layers: 12
num_heads: 12
channels: 768
num_parameters: 124439808
train dataset num_batches: 1192
val dataset num_batches: 128
num_activations: 73323776
val loss 5.252026
step 0: train loss 5.356189 (took 1452.121000 ms)
step 1: train loss 4.301069 (took 1288.673000 ms)
step 2: train loss 4.623322 (took 1369.394000 ms)
step 3: train loss 4.600470 (took 1290.761000 ms)
... (trunctated) ...
step 39: train loss 3.970751 (took 1323.779000 ms)
val loss 4.107781
generating:
---
Come Running Away,
Greater conquer
With the Imperial blood
the heaviest host of the gods
into this wondrous world beyond.
I will not back thee, for how sweet after birth
Netflix against repounder,
will not
flourish against the earlocks of
Allay
我喜欢 Netflix 出现的地方,很明显,模型中仍然保留了一些与训练数据相关的特征或模式。。我没有尝试调整微调超参数,因此很可能可以大幅改进。我还注意到,稍微不同的平台(例如 MacOS / Linux)会(遗憾地)给出略有不同的结果,因此也许不要期望得到上述精确的数字或生成。另请注意,如果在生成中看到标记 ID 而不是文本,那可能是因为您的代码已过时,因为 Tokenizer 解码是在 2024 年 4 月 14 日添加的。git pull
更新,然后重新运行 python train_gpt2.py
,它现在还会保存分词器,C 可以读取并用于打印文本而不是标记 ID。
测试
我也附上了一个简单的单元测试,以确保我们的C代码与PyTorch代码一致。请编译并运行以下命令:
bash
make test_gpt2
./test_gpt2
这个测试现在会加载gpt2_124M_debug_state.bin
文件,执行一次前向传播,将logits和损失与PyTorch参考实现进行比较,然后进行10次使用Adam优化器的训练迭代,并确保损失与PyTorch匹配。
教程
我在这里附上了一个非常小的教程,位于 doc/layernorm/layernorm.md。这是一个简单的、逐步指导如何实现GPT-2模型中一个单独层的文档,即LayerNorm层。这是一个很好的起点,以理解如何在C语言中实现这些层。
CUDA
完整的训练循环也在一个纯CUDA文件中实现,但内核的优化正在进行中。目前,我们大致与PyTorch的速度相匹配。我们组织代码的方式是在dev/cuda
文件夹中有一个不断增加复杂度的内核集合,详见dev/cuda/README.md。然后,我们将最佳内核复制粘贴到单个训练文件train_gpt2cu.cu
中的主要训练循环中。 正确性。首先,我们可以进行10次训练迭代,验证我们的代码是否与PyTorch完全匹配并重现数字:
bash
make test_gpt2cu
./test_gpt2cu
这将打印overall okay: 1
。因此,前向激活、反向梯度以及10次迭代的单独损失值都完全匹配。 训练。要在单个CUDA文件中训练GPT-2,请运行训练脚本:
bash
make train_gpt2cu
./train_gpt2cu
这将加载tiny_shakespeare
数据集的验证和训练拆分。在默认设置下,B=4,T=1024,有8个验证批次和74个训练批次。脚本当前配置为进行一次学习率为1e-4的微调周期,并在此过程中评估验证性能并生成样本,例如:
yaml
step 1/74: train loss 4.367631 (80.639749 ms)
step 2/74: train loss 4.031242 (77.378867 ms)
step 3/74: train loss 4.034144 (77.315861 ms)
step 4/74: train loss 3.859865 (77.357575 ms)
...
step 72/74: train loss 3.085081 (78.850895 ms)
step 73/74: train loss 3.668018 (78.197064 ms)
step 74/74: train loss 3.467508 (78.009975 ms)
val loss 3.516490
generating:
---
?Where will you go?
I take you wherefore I can, myself, and must.
I cast off my beak, that I may look him up on the point;
For on his rock shall he be opencast.
<|endoftext|>My little nephew:
Keep on with me, my
这在我的A100上大约需要10秒钟。在PyTorch脚本中,这个训练循环大约是80毫秒/迭代,所以我们在这里略优于PyTorch。然而,这是使用了稍旧版本的PyTorch(我使用的是2.1.0),而且我们还没有包括FlashAttention或PyTorch的scaled_dot_product_attention融合操作。 我们可以将其与原始的PyTorch进行比较,如下所示,我们打开torch.compile
和使用TensorCores,这些TensorCores使用tf32类型:
bash
python train_gpt2.py --write_tensors 0 --sequence_length 1024 --batch_size 4 --compile 1 --tensorcores 1
编译(第一次迭代)大约需要27秒,但在我的A100上,此后每次迭代目前运行时间约为80毫秒。
实验/扫描
现在.cu脚本中已经有了基本的argparse和日志功能,我们可以进行第一次学习率扫描。这目前还是相当手动的,但只是为了记录一个在TinyStories上使用4个GPU的机器上扫描学习率的例子过程。运行一个shell脚本sweep.sh
(在你当然chmod u+x sweep.sh
之后):
bash
#!/bin/bash
learning_rates=(3e-5 1e-4 3e-4 1e-3)
for i in {0..3}; do
export CUDA_VISIBLE_DEVICES=$i
screen -dmS "tr$i" bash -c "./train_gpt2cu -i data/TinyStories -v 250 -s 250 -g 144 -l ${learning_rates[$i]} -o stories$i.log"
done
# you can bring these down with
# screen -ls | grep -E "tr[0-3]" | cut -d. -f1 | xargs -I {} screen -X -S {} quit
这个例子打开了4个屏幕会话,并运行四个具有不同学习率(LR)的命令。这将写入日志文件stories$i.log
,其中包含所有的损失值,你可以根据需要在Python中绘制这些损失值。这里有一个快速的例子脚本,用于在Jupyter笔记本中绘制损失值,显然稍后可以变得更加复杂:
python
import matplotlib.pyplot as plt
%matplotlib inline
def parse_log(logfile):
# look for lines like e.g. "s:100 tel:1.6952", step 100, val 1.6952
val_steps, val_losses = [], []
with open(logfile, "r") as f:
lines = f.readlines()
for line in lines:
if "tel" in line:
parts = line.split()
step = parts[0].split(":")[1]
loss = parts[1].split(":")[1]
val_steps.append(int(step))
val_losses.append(float(loss))
return val_steps, val_losses
results = [parse_log(f"stories{i}.log") for i in range(0, 4)]
for i, (val_steps, val_losses) in enumerate(results):
plt.plot(val_steps, val_losses, label="run {}".format(i))
plt.xlabel("steps")
plt.ylabel("loss")
plt.legend()
仓库管理理念
关于我希望这个仓库成为的样子,我还有一些话要说:
首先,我希望llm.c
成为一个教育的平台。例如,我们的dev/cuda
文件夹是所有手动手写并记录得非常好的内核库的地方,从非常简单的内核一直到更复杂/更快的内核。如果你有一个新的内核,具有各种不同的权衡,欢迎你在这里贡献。
话虽如此,我也希望llm.c
能够非常快,甚至实际上用于训练网络。例如,作为开始,我们应该能够复现大型GPT-2(1.6B)的训练运行。这要求我们整合所有最快的内核,包括使用cuBLAS、cuBLASLt、CUTLASS、cuDNN等库。我还认为这样做具有教育意义,可以建立专家的上限基准和测量单位,例如你可以说你手动编写的内核速度是cuBLAS的80%等。然后你可以选择进行超快速运行,或者你可以选择"拖放"任何你希望使用的手动内核,并用它们运行。
然而,作为一个约束条件,我希望保持根文件夹中的主要llm.c
简单且易读。如果有PR例如提高了性能2%,但它"花费"了500行复杂的C代码,可能还有一个奇特的第三方依赖,我可能会拒绝这个PR,因为这种复杂性不值得。在这个意义上,如果这意味着我们可以保持在大约2000行易读的代码,最小化奇特的依赖,我会同意只有例如PyTorch速度的90%。作为一个具体的例子 - 在根训练循环中默认使用cuBLAS进行矩阵乘法是一个不费脑筋的决定:它使主线代码更快,它是一行可解释的代码,而且是一个非常常见的依赖。在这旁边,我们可以在dev/cuda
中有可以与cuBLAS竞争的手动实现。
最后,我对项目根文件夹中的复杂性会更加敏感,该文件夹包含项目的主/默认文件。相比之下,dev/
文件夹更像是一个我们开发内核或类库的 scratch space,并分享有用或相关或教育性的代码,这些代码可以在本地复杂一些。
显著分支
- Mojo版本:由GitHub用户dorjeduck创建的llm.mojo,这是一个用Mojo语言编写的项目。你可以通过链接访问他的代码。
- C#版本:由GitHub用户azret创建的llm.cs,这是一个用C#语言编写的项目。你可以通过链接访问他的代码。
讨论
组织开发的方式:
- 遇到与仓库相关的具体问题?使用Issues。
- 有代码要贡献?打开一个PR。
- 讨论仓库,提问等?查看Discussions。
- 需要更快的响应?我在Zero to Hero Discord频道上创建了一个新的
#llmc
频道。
开源许可证
MIT
基础知识
Cpython
是用c语言实现的python解释器,即负责将python源代码转换为字节码,并通过解释器执行这些字节码来实现代码的运行;
fp32
即32位浮点数,是计算机中用于表示实数的一种数据类型。它遵循IEEE 754标准,具有1位符号位、8位指数位和23位尾数位。在深度学习和科学计算中,FP32因其较高的精度和适中的存储需求而广泛使用。
FP32是一种常用的浮点数格式,它在保证足够精度的同时,也考虑到了存储和计算效率。
Flash Attention
是一个高效的算法优化,特别适用于需要处理大量数据的深度学习模型,如Transformer,在保持模型性能的同时降低了计算资源的消耗。
是一种针对Transformer模型中自注意力机制的优化方法。该算法旨在解决长序列处理时的高内存消耗和计算成本问题。通过将输入分块并在每个块上执行注意力操作,Flash Attention显著减少了对高带宽内存(HBM)的读写次数。具体来说,这个算法利用了SRAM(快速缓存)比HBM更快的特性,将输入块加载到SRAM中执行计算,并最小化了两者之间的数据传输。这种方法不仅提升了计算速度,还有助于减少内存占用,使得训练更大或更深的Transformer模型成为可能。
ddp
DDP是一种用于AI模型训练的技术,它允许多个GPU同时处理不同的数据批次,以提高训练速度和效率。这种方法在处理大型数据集和复杂模型时尤其有用,因为它可以显著减少训练时间。具体来说
cuda
CUDA ( Compute Unified Device Architecture )是 NVIDIA 推出的一种并行计算平台和编程模型。它利用NVIDIA的图形处理器(GPU)强大的性能显著地加速计算应用。通过使用CUDA,开发者可以使用扩展的C/C++编程语言(称为CUDA C/C++)来编写运行在GPU上的代码,从而进行高性能计算。
现代框架:
- Llama 2:Llama 2 是由Meta AI(前身为Facebook AI Research)开发的大语言模型。这是基于原始Llama模型的进一步改进,专门设计用于处理各种语言任务,如文本生成、翻译和问答等。Llama 2 在性能和效率方面都进行了优化,能够适应更大的数据集并提供更准确的结果。
- Gemma:Gemma 是由爱丁堡大学的研究人员开发的一个开源库,旨在构建模块化且高效的Transformer模型。Gemma的目标是简化定制Transformer架构以针对特定任务的设计和训练过程。Gemma强调灵活性和可重用性,使研究者和开发人员更容易尝试不同的模型配置。
- Mistral:Mistral 是另一种基于Transformer的架构,专为自然语言处理任务而设计。
Transformer
Transformer 是一种革命性的深度学习模型架构,专门用于处理自然语言处理( NLP )任务。
Transformer模型的核心创新在于引入了自注意力机制(self-attention mechanism),这一机制允许模型在处理序列数据时能够关注到输入数据的不同部分,从而更好地理解语言的上下文关系
gpt-2是什么?为什么会被称为语言大模型的鼻祖呢?
GPT-2 被称为大语言模型( LLM )的鼻祖,因为它是第一个以现代形式成功实现并广泛影响了后续模型设计的大规模语言模型。
GPT-2在大型语言模型的发展中扮演了重要角色,其原因主要包括:
- 技术基础:GPT-2是基于Transformer架构的自回归模型。它使用了Transformer的解码器模块,并且在设计上进行了创新,如去掉了原始Transformer解码器中的编码器-解码器注意力层,这使得它专注于生成自然语言文本。
- 模型规模:GPT-2的成功训练和发布展示了大规模神经网络语言模型的潜力,为后来更多更大语言模型的发展奠定了基础。
- 影响力:GPT-2发布后,其在语言模型设计上的创新被广泛采纳和借鉴。例如,它的自回归方式和对解码器的使用影响了后续的语言模型设计思想。
- 资源提供:OpenAI提供了GPT-2的权重,使得研究者可以在其基础上进行进一步的研究和微调。这种开放性的做法促进了社区对于大规模语言模型的研究和应用开发。
综上所述,GPT-2不仅因其技术突破而被认为是LLM的鼻祖,也因为它对整个AI领域产生了深远影响,并为未来语言模型的发展铺平了道路。
GPU训练中起什么作用?
GPU在训练中的作用主要体现在以下几个方面:
- 加速计算:深度学习模型训练过程中需要进行大量的矩阵运算,GPU通过其众多的内核可以同时处理这些运算,大大加快了计算速度。
- 提高并行性:GPU的设计使其能够同时处理数千个独立的线程,这在执行深度学习中的并行操作时非常有用,如同时处理多个数据样本。
- 优化内存带宽:GPU拥有高内存带宽,可以快速地将数据传输到计算单元,这对于数据密集型的深度学习任务至关重要。
总的来说,GPU在深度学习和神经网络的训练中起着至关重要的作用,它通过并行处理能力、高效的内存带宽和浮点运算性能,显著提高了模型训练的效率和速度。这使得研究人员和工程师能够在更短的时间内训练出更复杂、更精确的模型,推动了人工智能技术的发展。
tinyshakespeare数据集
这个数据集是从莎士比亚的作品中提取的,包含了大约4万行文字。它通常用于训练和测试自然语言处理模型,特别是在文本生成领域。通过训练,模型能够学习到莎士比亚的语言模式和表达方式,从而生成风格相似的新文本。
什么是experiments / sweeps?
实验/扫描是深度学习中常见的实验管理技术,用于系统地探索模型架构、超参数设置和训练流程的不同组合。通过实验/扫描,可以更有效地发现最佳配置,并加速模型的优化过程。
在实验/扫描中,通常会定义一个参数空间,包括模型架构参数、优化器参数、学习率、批量大小等。然后,系统会自动执行一系列实验,每个实验使用参数空间中的一个组合。通过比较实验结果,可以找到最佳配置,从而改进模型性能。
实验/扫描可以使用各种工具和框架来实现,例如:
- Ray Tune:Ray Tune是一个用于分布式超参数优化和自动机器学习的开源库,提供了灵活的实验/扫描功能。
- Optuna:Optuna是一个用于超参数优化的自动化框架,支持并行和分布式优化。
- TensorBoard Hparams:TensorBoard的超参数插件可以帮助您跟踪和可视化实验/扫描的结果。 通过实验/扫描,您可以更快地找到最佳模型配置,提高模型性能并加速实验迭代过程。