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

femtoGPT简介

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

femtoGPT的主要特点包括:

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

技术实现

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中展示了多个训练阶段的输出样本,展示了模型从最初的随机输出到逐渐学习语言结构的过程。以下是一些有趣的例子:

1. 初期训练(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

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

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

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

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

3. 后期训练(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

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

4. 最新更新(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.

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

项目意义与展望

femtoGPT项目虽然在性能和实用性方面无法与大型商业模型相提并论,但它在教育和研究方面具有重要意义:

1. 深入理解:通过从零实现GPT模型,帮助开发者和研究者深入理解大型语言模型的工作原理。
2. Rust语言在AI领域的应用:展示了Rust语言在机器学习和人工智能领域的潜力。
3. 开源贡献:为开源AI社区提供了一个可学习、可改进的基础项目。
4. 硬件亲和:通过支持CPU和OpenCL GPU,降低了实验门槛。

未来,femtoGPT项目可能会在以下方面继续发展:

• 优化性能:在保持代码可读性的同时,提高计算效率。
• 扩展功能:实现更多GPT相关的技术,如微调、提示工程等。
• 社区合作:吸引更多贡献者,共同改进和扩展项目。

结语

femtoGPT作为一个教育和研究导向的项目,为我们提供了一个难得的机会,让我们能够从最基础的层面理解和实现GPT模型。虽然它可能永远无法取代那些训练有素的大型语言模型,但它的价值在于启发和教育。对于那些对AI和机器学习充满热情,希望深入了解这些技术的人来说,femtoGPT无疑是一个极好的起点。

如果您对femtoGPT项目感兴趣,可以访问其GitHub仓库了解更多信息,或者加入项目的Discord服务器参与讨论。无论您是经验丰富的开发者,还是刚刚开始学习AI的新手,femtoGPT都为您提供了一个绝佳的学习和实验平台。让我们一起探索GPT模型的奥秘,见证AI技术的不断进步吧! 🚀🤖