100美元成本、8000行代码纯手搓克隆ChatGPT!
特斯拉前AI总监、OpenAI创始成员、宣布全职搞教育的AI大神Andrej Karpathy(卡帕西)沉寂了好久,终于终于终于来上新课了!
新作nanochat,被其本人称作是写得最“精神错乱”放飞自我的作品之一。
它是一个极简的、从零开始构建的全栈训练/推理pipeline,用最少量依赖的单一代码库实现了简易版ChatGPT。
只要你启动一台云GPU服务器,运行一个脚本,最快只要4小时,就能在类似ChatGPT的网页界面与自己训练的大语言模型对话。
整个项目约8000行代码,可实现以下功能:
基于全新Rust语言实现,训练分词器 (tokenizer)
在FineWeb数据集上预训练Transformer架构大语言模型,并通过多项指标评估CORE得分
在SmolTalk用户-助手对话数据集、多项选择题数据集、工具使用数据集上进行中期训练 (Midtrain)
执行指令微调 ( S FT) ,并在世界知识多项选择题数据集 (ARC-E/C) 、数学数据集 (GSM8K) 、代码数据集 (HumanEval) 上评估对话模型性能
可选在GSM8K数据集上通过“GRPO”算法对模型进行强化学习 ( RL ) 训练
在推理引擎中实现高效模型推理,支持KV缓存、简易预填充/解码流程、工具使用 (轻量级沙箱环境中的Python解释器) ,可通过CLI或类ChatGPT的WebUI与模型交互
生成单个Markdown格式报告卡,对整个训练推理流程进行总结,并加入“游戏化”呈现 (如用评分、进度等形式直观展示结果)
整体成本只需约100美元(在8×H100上训练4小时),就能训练复刻出一个可进行基础对话、创作故事诗歌、回答简单问题的简易版ChatGPT模型。
整体表现指标如下:
训练约12小时后,模型在CORE指标上的表现即可超越GPT-2。
若进一步将成本提升至约1000美元(训练约41.6小时),模型表现显著提升,能解决简单的数学/代码问题,还能做多项选择题。
举个具体的例子:一个深度为30的模型训练24小时后(相当于GPT-3 Small 125M的算力消耗,仅为GPT-3的千分之一),在MMLU数据集上可达到40多分,在ARC-Easy数据集上达70多分,在GSM8K数据集上达20多分。
卡帕西表示,他的目标是将这套完整的“强基线”技术栈整合为统一、极简、易读、可修改、易分发的代码库。
nanochat将成为LLM101n课程的压轴项目(该课程仍在开发中)。
我认为它还有潜力发展为一个研究工具框架或基准测试的工具,就像之前的nanoGPT一样。目前该项目远未完全优化(实际上存在大量可改进空间),但整体框架已足够完整,可以发布到GitHub上,后续所有模块都能在社区中进一步优化。
等来新作的网友也已彻底疯狂。项目刚发出来,GitHub Star数已飙到4.8k:
太酷了!跑一次这个项目,就把“机器学习工程师(ML Engineer)”放在我的简历上!
你发布的不只是代码,更是可被理解的智慧,价值爆炸,栓Q。
在评论区,卡帕西还解释了nanochat基本架构与Llama类似,但更简化一些,也借鉴了部分modded-nanoGPT的设计,整体是为此规模的模型找到一个稳健的基础架构。
以及这个项目基本上是完全手写的。
我确实尝试过用Claude或Codex之类的Agent来帮忙,但效果非常糟糕,几乎毫无帮助。可能是因为这个repo的结构偏离了它们训练数据的分布,所以它们根本“对不上号”。
话不多说,下面来看nanochat快速上手的详细指南。
100美元成本,能捏出的最好的ChatGPT
从比如Lambda GPU Cloud上启动了一台8卡H100的服务器,每小时要花大约24美元,所以接下来得争分夺秒了。
环境搭建
克隆项目:
目标是用100美元的成本训练出一个最好的类ChatGPT模型,称之为一次“速通(speedrun)”,可参考speedrun.sh这个脚本,它被设计成能在一台全新的服务器上直接从头到尾运行。
但接下来,卡帕西会逐步讲解其中的每一步。
首先需要确保安装了当下热门的uv项目管理器。安装uv,在.venv目录下创建一个新的虚拟环境,获取所有依赖项,然后激活该环境,这样当输入python时,使用的是虚拟环境中的Python,而不是系统自带的Python:
接下来,需要安装Rust/Cargo,以便编译自定义的Rust分词器。引入一个全新/自定义的分词器确实有点折腾,但遗憾的是,卡帕西觉得早期minbpe项目中的Python版本速度太慢,而huggingface的分词器又过于臃肿且令人困惑。
因此要专门为训练打造了自己的新分词器(经测试与Python版本效果一致),不过在推理时仍会使用OpenAI的tiktoken来保证效率。
现在就开始编译分词器吧:
训练分词器
接下来,需要获取预训练数据,这样才能:1)训练分词器;2)对模型进行预训练。
预训练数据就是大量网页的文本内容,这里将使用FineWeb-EDU数据集。
通常来说,可以直接用huggingface datasets.load_dataset(),但卡帕西不喜欢它过于臃肿笨重且掩盖了本应简单的逻辑,所以把整个数据集重新打包成了简单、完全打乱的分片,这样就能轻松高效地随意访问,并且把它的sample-100B版本重新上传为karpathy/fineweb-edu-100b-shuffle。
在这个页面上,你还可以预览数据集中的示例文本。每个分片是一个约0.25M个字符的简单parquet文件,压缩后(gzip格式)在磁盘上大约占100MB。总共有1822个分片,但训练深度为20的模型只需要其中240个。
现在就开始下载所有数据吧。虽然需要下载约24GB,但在云服务器上通常速度很快:
默认情况下,所有这些都会被下载到~/.cache/nanochat目录下。
下载完成后,开始训练分词器——它负责在字符串与符号码本(codebook)序列之间进行双向转换。默认情况下,训练的词汇表大小是2¹⁶= 65,536个tokens(这是个不错的数字),其中部分tokens会被保留作为特殊tokens(供后续聊天模式使用)。训练集包含2B字符,训练仅需约1分钟。
训练算法与OpenAI使用的完全一致(regex splitting, byte-level BPE)。想了解更多信息,可以看卡帕西关于tokenization技术的视频讲解。
训练完成后可以评估这个分词器:
评估结果显示,实现了约4.8的压缩比(即原始文本中平均4.8个字符压缩为1个token),还可以看到与GPT-2、GPT-4分词器的对比结果。
相比GPT-2(拥有50257个tokens),在压缩文本方面全面更优,仅在数学内容上稍逊一筹:
与GPT-4相比,表现并不突出,但需要考虑到GPT-4拥有更大的词汇表规模(100,277个tokens)。特别是在多语言处理方面GPT-4优势明显(由于FineWeb数据集高度侧重英语内容,这个结果很合理),同时在代码和数学领域也更胜一筹:
尽管如此,即使在词汇量较小的条件下,我们在FineWeb数据集上仍以微弱优势超越了GPT-4——因为这正是我们训练所用的数据集,所以我们的分词器能完美契合该文档分布(例如在英语文本压缩方面可能更具优势)。
预训练
在启动预训练之前,需要下载另一个被卡帕西称之为“评估包(eval bundle)”的文件。
在预训练过程中,脚本会定期评估CORE指标。你可以在DCLM论文中看到一些细节,本质上,它是一个很好的、标准化的、宽泛的指标,用于衡量模型在大量自动补全数据集上的表现好坏。
这些数据集包括HellaSwag、jeopardy、bigbench QA wikidata、ARC-Easy/Challenge、copa、commonsense qa、piqa、lambada、winograd、boolq等等(共22个)。
下载、解压该评估包,并将评估包目录放置到基础目录~/.cache/nanochat/eval_bundle下:
还建议(尽管这是可选的)再做一项设置:
配置wandb,以便在训练过程中查看美观的图表。前面uv已经安装好了wandb,但你仍需创建账户并登录:
现在我们可以启动预训练了!这是计算量最大的部分,要训练大语言模型(LLM),通过预测序列中的下一个token来压缩互联网网页文本,在此过程中,大语言模型会获取大量关于世界的知识:
在这里,通过scripts/base_train.py脚本在8块GPU上启动训练。我们正在训练一个有20层的Transformer。默认情况下,每块GPU在每次前向/反向传播时处理32行、每行2048个tokens的数据,优化器每一步总共处理32×2048=2¹⁹=524,288≈0.5M个tokens。
如果已经设置好了wandb,可以添加—run=speedrun(所有训练脚本都支持该参数)来设置运行名称并记录相关数据。
当你启动训练后,会看到类似这样的输出(为简洁起见,省略了大量内容):
可以看到,这个Transformer有1280个channels,注意力机制中有10个注意力头,每个头的dim=128。它大约有560M参数。为了符合Chinchilla scaling law的建议,这意味着我们需要用560M×20≈11.2B tokens来进行训练。
由于优化器的每一步处理524,288个tokens,这意味着11.2B/0.5M≈21400次迭代。
通过对每个token的估计FLOPs与总tokens数相乘,我们可以知道这将是一个计算量达约4e19 FLOPs的模型。
学习率会自动按1/sqrt(dim)自动缩放,因为更大的模型更偏好更小的学习率。
我们使用Muon来优化矩阵,使用AdamW来优化嵌入和反嵌入。在这个模型中,没有其他可训练的参数(比如偏置、rmsnorm参数等)。训练过程会定期报告“验证集bpb”,即验证数据集上每字节的位数。
每字节位数(bits per byte)是一个比典型的交叉熵损失更好的衡量指标,因为它通过每个token的字节数进一步归一化了每个token的损失,使得该指标与分词器无关。
所以,无论你使用的是词汇量小的分词器还是词汇量大的分词器,这个数值都是可比较的,而原始的交叉熵损失则不然。
注意,每一步大约耗时0.5秒,lrm是学习率衰减乘数(在训练接近尾声时,它会线性下降到0),报告的MFU(模型flops利用率)看起来很不错,几乎达到了一半,这意味着我们充分利用了可用的bfloat16计算能力。
现在,要等待大约3小时,直到4e19 FLOPs的计算量完成……在你的wandb图表中,你应该会看到类似这样的内容:
随着时间的推移,bpb下降是好的迹象(说明模型能更准确地预测下一个token)。此外,CORE分数在上升。
除了这些近似的指标,还可以更全面地评估模型:
可以看到,训练集/验证集的bpb达到了约0.81,CORE指标上升到了0.22。
作为对比,评估包中包含了GPT-2模型的CORE分数。具体来说,0.22的CORE分数略高于GPT-2 large(0.21),但略低于GPT-2 xl(即“标准”的GPT-2,为0.26)。
此时,这个模型就像一个高级的自动补全工具,所以我们可以运行一些提示词,来感受模型中存储的知识。base_loss.py文件会运行这些提示词。这些提示词包括:
补全后的文本如下:
所以,模型知道巴黎是法国的首都、Au代表金、星期六在星期五之后、“冷”是“热”的反义词,甚至还知道太阳系的行星。
不过,它对天空的颜色还不太确定,也不太会做简单的数学题。
对于一个花费72美元训练出来的模型来说,已经不算太差了。推理过程使用了一个自定义的Engine class,利用KV缓存来实现高效推理,同时还简单实现了两种常见的推理阶段:预填充和解码。
我们的Engine class还支持工具使用(比如Python解释器),这在GSM8K数据集上训练时会很有用(之后会详细介绍)。
训练中期
接下来是中期训练,这一步会在smol-SmolTalk数据集上进一步微调模型。
算法层面和预训练完全一致,但数据集变成了对话内容,而且模型会去适应那些用于构建多轮对话结构的新特殊token。现在,每次对话大致是这样的,大致遵循OpenAI的Harmony聊天格式:
像<|example|>这样显示的token是特殊token,遵循OpenAI特殊token的格式。中期训练阶段对模型的多种适配非常有用:
模型学习与多轮对话相关的特殊token (除了用于分隔文档的<|bos|>token,基础模型预训练期间没有这些token) 。
模型适应对话的数据分布,而非互联网文档的数据分布。
对我们来说非常重要的一点是,必须教会模型做多项选择题,因为在这么小的模型规模下,模型无法从随机的互联网数据中学会这一点。具体而言,模型必须学会将几个选项与几个字母 (如ABCD) 关联起来,然后输出正确选项的算法。通过混合10万道来自MMLU辅助训练集的多项选择题来实现这一点。需要明确的是,问题不在于模型没有相关知识,而在于它不理解多项选择题的运作方式,无法将知识展现出来。这很重要,因为许多常见的模型评估 (如MMLU) 都采用多项选择题的形式。
你可以教会模型使用各种工具。对我们来说,需要通过在特殊token <|python_start|>和<|python_end|>之间放入Python命令,来教会模型使用Python解释器。这对之后解决GSM8K问题会很有用。
在中期训练期间,你还可以针对许多其他适配进行训练,例如上下文长度扩展 (尚未探索) 。
中期训练混合数据默认是这样的:
然后按如下方式启动它:
这次运行只需要大约8分钟,比预训练的约3小时短得多。现在,模型已经是一个真正的聊天模型,能够扮演助手的角色回答用户的问题,可以对其进行评估:
得到了该阶段模型的以下结果:
可以看到:
世界知识:前三项(ARC-E/C和MMLU)都是多项选择题测试,用于衡量模型在各个领域的世界知识。由于有4个选项(A、B、C、D),随机猜测的正确率约为25%,所以模型已经表现得比随机猜测更好了。(对于这么小的模型来说,多项选择题是相当难的)
数学:GSM8K是小学水平的数学题。这里的基准性能是0%,因为模型必须写出实际的答案数字。目前我们的性能仍然不是很强,只解决了2%的问题。
代码:HumanEval是一个Python编码基准测试,同样,随机基准性能为0%。
ChatCORE:这是卡帕西尝试复制CORE分数对基础模型的评估方式,并将其扩展到聊天模型的成果。也就是说,将上述所有指标都减去基准性能,这样分数就在0到1之间(例如,随机模型得0分,而不是MMLU上的25%),然后报告所有任务的平均值。它是对当前模型实力的一个单一数字总结。
这些评估仍然相当不完整,还有很多其他可以衡量但尚未衡量的方面。
确实没有一个很好的图表来展示这一步,但这里有一个之前对另一个更大的模型进行中期训练的例子,只是为了让你了解在微调运行期间这些指标上升时的样子:
监督微调
中期训练之后是监督微调(SFT)阶段。
这是在对话数据上额外进行的一轮微调,理想情况下,你会精心挑选最优质的好数据,而且也会在这里进行安全训练(比如助手拒绝不当请求的训练)。
我们的模型甚至连天空的颜色都还不确定,所以目前在生物危害这类问题上可能还是安全的。这里会进行的一项领域适配是,SFT会拉伸数据行并对其进行填充,完全模拟测试时的格式。
换句话说,示例不再像预训练/中期训练时那样为了训练效率而被随机拼接成长行。修正这种领域不匹配的问题,是另一个小小的“拧紧螺丝”式的提升。我们可以运行SFT并重新评估:
这个过程同样只需运行约7分钟,你应该能观察到各项指标均有小幅提升:
终于,我们可以以用户身份与模型对话了!
其实在中期训练后就可以进行对话,但现在效果会更理想些。你可以通过终端窗口(方式1)或网页界面(方式2)与它交流:
chat_web脚本会使用FastAPI来提供Engine服务。要确保正确访问它,比如在Lambda上,使用你所在节点的公网IP,后面加上端口,例如http://209.20.xxx.xxx:8000/等等。
那看起来会很棒,大概是这样的:
它短期内还无法在物理或诗歌比赛中获胜,但话说回来——用这么少的预算能做到这个程度,看起来还是很酷的,而且这个项目还远远没到充分调优的地步。
强化学习
“速通”的最后一个阶段是强化学习。
基于人类反馈的强化学习(RLHF)是一种不错的方法,能提升几个百分点的性能,还能缓解很多因采样循环本身带来的模型缺陷——比如幻觉、无限循环等。
但以我们的规模,这些都不是主要考虑因素。话虽如此,在我们目前使用的所有数据集中,GSM8K是唯一一个有清晰、客观奖励函数的(数学题的正确答案)。
所以我们可以运行RL(/GRPO)脚本,通过交替进行采样和训练的简单强化学习循环,直接在答案上进行性能攀升:
在强化学习过程中,模型会遍历训练集中所有的GSM8K题目,对完成情况进行采样,然后我们会对这些采样结果进行奖励,并针对获得高奖励的样本进行训练。
我们使用的是高度简化的GRPO训练循环,比如,不使用信任区域(舍弃参考模型和KL正则化),采用在策略(舍弃PPO的比率+裁剪),使用GAPO风格的归一化(基于token级,而非序列级归一化),优势函数仅通过均值进行简单的奖励平移(舍弃用除以标准差的z分数归一化)。
所以最后得到的东西看起来更像是REINFORCE算法,但保留了GR(”组相对”)部分来计算奖励的优势值。在当前规模和任务简单度下,这种方法效果尚可。更多细节请参阅脚本。
目前强化学习默认是注释掉的,因为它还没有经过很好的调优,而且我们也没有完整通用的RLHF。
只针对GSM8K进行了强化学习,这也是为什么用-a标志将评估也限制在GSM8K上。由于强化学习就像通过吸管汲取监督信号,这个过程会运行相当长的时间。
例如,默认设置下运行约1.5小时后,效果如下所示:
成绩
最后卡帕西指出的是项目文件夹里出现的report.md文件。它包含了很多与运行相关的细节,最后还有一个不错的总结表格:
Characters:333,989
Lines:8,304
Files:44
Tokens(approx):83,497
Dependencies(uv.lock lines):2,004
总用时:3小时51分钟
需要注意的是,由于目前对强化学习(RL)的支持还不太完善,在计算总耗时时把它排除了。到监督微调(SFT)阶段为止,整个过程运行了3小时51分钟,总成本为(3+51/60)×24=92.4美元(如果加上强化学习,现在总时间会更接近5小时)。
甚至还剩下8美元可以买冰淇淋呢。
该你了
借助nanochat,你可以对任何部分进行调优。
更换分词器、修改任意数据、调整超参数、改进优化过程……有很多想法可以去尝试。你或许还想训练更大的模型。这个代码库的设置能让你轻松做到这一点。
只需使用—depth参数来更改层数,其他所有相关设置都会基于这个参数作为复杂度的单一调节项而自动调整。比如,通道数会增加,学习率会相应调整等。
原则上,仅通过改变深度,你就能探索出一整套nanochat的“迷你系列”模型。使用更大的深度并等待更长时间,理论上你应该能得到明显更好的结果。
你需要在base_train.py的预训练阶段传入深度参数。例如,要得到一个CORE指标约为0.25、性能接近GPT-2的模型,尝试depth=26是个不错的选择。
但训练更大模型时,需要调整设备最大批处理大小,比如从32降至16:
代码会察觉到这一变化并自动进行补偿,它会通过2次梯度累积循环来达到目标批处理量0.5M。要训练depth=30的模型,需要进一步降低设置:
依此类推。欢迎大家去阅读代码,卡帕西尽力让代码保持易读性,添加了注释,代码整洁且易于理解。
当然,你也可以把所有内容打包,去询问你喜欢的大语言模型,或者更简单的是,使用Devin/Cognition的DeepWiki来对这个代码仓库提问。只需把代码仓库的URL从github.com改成deepwiki.com即可,比如 nanochat DeepWiki。
就是这样,调优整个流程的任意部分,重新运行,然后享受其中的乐趣吧!
AI界超高人气专注于教育的大牛
卡帕西曾任特斯拉AI主管,之后去了OpenAI,去年2月从OpenAI离职。
他在整个AI界拥有超高的人气,很大一部分来自于他的课程。
包括他自己的早期博客文字分享和后来的一系列Youtube视频教程,他还与李飞飞合作开设的的斯坦福大学首个深度学习课程CS231n《卷积神经网络与视觉识别》。
今天的不少学者和创业者,都是跟着他入门的。
卡帕西对教育的热情,甚至可以追溯到学生时期在网上教大家玩魔方。
去年7月,从OpenAI离职的卡帕西突然官宣创业,搞了一家AI原生的新型学校——Eureka Labs。
怎么理解AI原生?
想象一下与费曼一起学习高质量教材,费曼会在每一步中1对1指导你。
不幸的是,即使每个学科都能找到一位像费曼这样的大师,他们也无法分身亲自辅导地球上的80亿人。
但AI可以,而且AI有无限的耐心,精通世界上所有的语言。
所以卡帕西要打造“教师+人工智能的共生”,可以在一个通用平台上运行整个课程。
如果我们成功了,任何人都将易于学习任何东西,扩大教育这个概念本身的“范围”和“程度”。
Eureka Labs首个产品,也是首门课程LLM101n。
手把手带你构建一个类似ChatGPT的故事生成大模型,以及配套的Web应用程序。
GitHub repo:https://github.com/karpathy/nanochat
详细指南:https://github.com/karpathy/nanochat/discussions/1
参考链接:https://x.com/karpathy/status/1977755427569111362