femtoGPT: 纯Rust实现的极简生成式预训练Transformer

femtoGPT

femtoGPT简介

femtoGPT是由开发者Keyvan Kambakhsh创建的一个开源项目,旨在用纯Rust语言从零实现一个极简版的生成式预训练Transformer(GPT)模型。这个项目的目标是帮助人们更深入地理解大型语言模型的工作原理,同时也展示了Rust语言在机器学习领域的应用潜力。

femtoGPT Logo

femtoGPT的主要特点包括:

纯Rust实现:整个项目完全使用Rust语言编写,包括张量处理逻辑和GPT架构的训练/推理代码。
极简设计:项目采用最简单的方式实现了GPT模型的核心功能,便于理解和学习。
CPU和GPU支持:可以使用CPU或GPU进行模型训练和推理。
开源透明:代码完全开源,方便研究和改进。

技术实现

femtoGPT的实现借鉴了Andrej Karpathy的nanoGPT视频讲座中的架构设计。项目使用了最基本的Rust库来实现所需功能:

随机数生成: rand和rand-distr
数据序列化: serde和bincode(用于保存和加载已训练的模型)
并行计算: rayon

值得注意的是,femtoGPT几乎所有的基本操作(如矩阵乘法)都采用了最简单的实现方式。这虽然牺牲了一些性能,但大大提高了代码的可读性和可理解性。

GPU加速

femtoGPT的GPU实现基于OpenCL,这意味着它可以在NVIDIA和AMD显卡上运行,而无需安装庞大的CUDA工具包。这大大降低了使用门槛,使更多人能够尝试使用GPU进行模型训练。

梯度检查

为了确保模型训练的正确性,femtoGPT使用了梯度检查方法来验证梯度计算的准确性。尽管如此,开发者也坦言可能仍有一些层的实现存在问题。

使用方法

要使用femtoGPT,首先需要安装Rust工具链:

curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh

如果想使用GPU进行训练,还需要确保系统中正确安装了GPU驱动和OpenCL运行时。对于Debian系统,可以通过以下命令安装OpenCL运行时:

sudo apt install ocl-icd-opencl-dev

接下来,将要训练的文本数据放入dataset.txt文件中。需要注意的是,为了简化模型,建议使用的数据集中包含的唯一字符数量不要太多(例如,当前数据集只使用了65个不同的唯一字符)。

然后,运行训练命令即可开始模型训练。训练数据将保存在train_data目录中,允许随时停止和继续训练。

训练效果展示

femtoGPT的开发者在项目README中展示了多个训练阶段的输出样本,展示了模型从最初的随机输出到逐渐学习语言结构的过程。以下是一些有趣的例子:

初期训练(300k参数模型,基于莎士比亚数据集):

LIS:
Tore hend shater sorerds tougeng an herdofed seng he borind,
Ound ourere sthe, a sou so tousthe ashtherd, m se a man stousshan here hat mend serthe fo witownderstesther s ars at atheno sel theas,
thisth t are sorind bour win soutinds mater horengher

这个输出虽然看起来毫无意义,但已经开始生成一些容易发音的"单词"。

中期训练(相似规模模型,更多训练时间):

What like but wore pad wo me che nogns yous dares,
As supt it nind bupart 'the reed:
And hils not es

此时模型已经开始学习一些单词和标点符号规则。

后期训练(5层模型,基于Reddit数据集,词汇量500,32核CPU训练约10小时):

Prompt "I think":

I think it s not a scientific method of atheism
I think it s not a fairly simple possible to have a
I think that s a possible point, but I m not sure if you re not
I think that s a great idea, but I m not sure how these are you talking
I think it s a simple thing to do, but I m not sure how thes

这个阶段的输出已经开始呈现出连贯的句子结构,虽然内容仍然有些混乱。

最新更新(8层8头128嵌入维度模型,基于TinyStories数据集,词汇量1000):

Once upon a time, there was a little girl named Lily.
She loved to play with her toys and she had a lot of fun.
One day, Lily saw a big chicky playing with her toys.
She asked her mom, "Can I play with her toys?" Her mom said,
"Sure, Lily. But we have to clean the pales. Let's suet some candy, Lily."
Lily nodded and went to her mom. They played with the mots and staugning her toys.

这个输出已经能够生成相对连贯的短故事,虽然仍然包含一些不存在的词语。