最简单最快的库

詹士 发自 凹非寺量子位 | 公众号 QbitAI

训练/微调中型GPT,最简单最快的库来了!

其名为:NanoGPT。

从名字就能看出是个“纳米武器”,据作者介绍,该库代码简单易读,2个仅300行代码的文件。

现已基于OpenWebText重现 GPT-2 (124M),在单个8XA100 40GB节点上,训练时间为38小时。

值得一提的是,该库发布者是前特斯拉AI总监,李飞飞高徒,Andrej Karpathy。此次发布的NanoGPT,正是他2年前MinGPT的升级版。

小破本也能炼GPT!300行代码搞定,GitHub获6k星,来自特斯拉前AI总监

目前,此项目在GitHub所获star已超6k,HackerNews上points也破千。

小破本也能炼GPT!300行代码搞定,GitHub获6k星,来自特斯拉前AI总监

毫无意外地,评论区一片“喜大普奔”。

有网友表示,这才是咱独立开发者喜闻乐见的AI工具。

小破本也能炼GPT!300行代码搞定,GitHub获6k星,来自特斯拉前AI总监

还有人对其一直开放传授分享知识的做法,表示感谢。

小破本也能炼GPT!300行代码搞定,GitHub获6k星,来自特斯拉前AI总监

那么,这个最简单最快的NanoGPT怎么用?

下面展开讲讲。

NanoGPT的打开方式

发布文件里面包含一个约300行的GPT模型定义(文件名:model.py),可以选择从OpenAI加载GPT-2权重。

还有一个训练模型PyTorch样板(文件名:train.py),同样也是300多行。

作者补充道,代码并不难,很容易就能满足大家需求——无论是从头开始训练新模型,还是基于预训练进行微调(目前可用的最大模型为1.3B参数的GPT-2)。

小破本也能炼GPT!300行代码搞定,GitHub获6k星,来自特斯拉前AI总监

△ 一个训练实例展示

上手前,需要提前准备好依赖项:

  • pytorch <3
  • numpy <3
  • pip install datasets for huggingface datasets <3 (如果你需要下载和预处理OpenWebText)
  • pip install tiktoken for OpenAI’s fast BPE code <3
  • pip install wandb for optional logging <3
  • pip install tqdm

先下载并标记OpenWebText数据集。

$ cd data/openwebtext
$ python prepare.py

这将创建一个train.bin和val.bin文件,将 GPT2 BPE token id放入一个序列中。

然后准备训练,目前脚本默认是尝试重现GPT-2,124M参数版本,但作者更鼓励大家阅读代码查看文件顶部的设置及路径

$ python train.py

如需使用 PyTorch 分布式数据并行 (DDP) 进行训练,请使用 torchrun 运行脚本。

比如,要在4个GPU节点上运行,代码如下:

$ torchrun –standalone –nproc_per_node=4 train.py

要从模型节点中采样,就需将一些检查点写入输入目录中。

$ python sample.py

据作者目前自己的测试,他在1 个 A100 40GB GPU 上训练一晚,损失约为 3.74。如果是在4个GPU上训练损失约为3.60。

如果在8个A100 40GB节点上进行约50万次迭代,时长约为1天,atim的训练降至约3.1,init随机概率是10.82,已将结果带到了baseline范围。

观察不同参数下训练/验证loss值如下:

小破本也能炼GPT!300行代码搞定,GitHub获6k星,来自特斯拉前AI总监

至于如何基于新文本微调GPT,作者也简介了方法。

先访问data/shakespeare,查看prepare.py。

下载小型shakespeare数据集并将其呈现为train.bin和val.bin文件(方法前文已介绍),几秒即可搞定。

运行一个微调示例,如下:

$ python train.py config/finetune_shakespeare.py

该操作将加载配置参数,覆盖config/finetune_shakespeare.py文件。

作者指出,一般情况下,基本操作就是从GPT-2检查点初始化init_from,再正常训练。

此外,如果手里只有macbook或一些“力量”不足的小破本,作者建议使用shakespeare数据集,然后在一个很小的网络上运行。

先渲染数据;

$ cd data/shakespeare
$ python prepare.py

再用一个较小的网络来运行训练脚本。

比如下面就创建了一个小得多的Transformer(4层,4个head,64嵌入大小),只在CPU运行,在作者自己的苹果AIR M1本上,每次迭代大约需要400毫秒。

$ cd ../..
$ python train.py –dataset=shakespeare –n_layer=4 –n_head=4 –n_embd=64 –device=cpu –compile=False –eval_iters=1 –block_size=64 –batch_size=8

关于NanoGPT的后续计划,Andrej Karpathy也在网上有所分享。

他将试图让NanoGPT更快复现其他GPT-2模型,然后将预训练扩展至更大规模的模型/数据集中,此外,他还计划改进下微调部分的文档。

转战教育和开源的特斯拉前AI总监

熟悉Karpathy的圈内人肯定知道,他此前是李飞飞高徒,也长期致力于让更多人接触了解神经网络和相关数据集。

2020年8月,他就曾发布NanoGPT前一代,MinGPT,同样旨在让GPT做到小巧、简洁、可解释,同样主打300行代码搞定。

Karpathy另一大身份是前特斯拉AI核心人物。

在马斯克麾下,他历任特斯拉高级AI主管、特斯拉自动驾驶AutoPilot负责人、特斯拉超算Dojo负责人、特斯拉擎天柱人形机器人负责人…

小破本也能炼GPT!300行代码搞定,GitHub获6k星,来自特斯拉前AI总监

2022年7月,Karpathy Andrej离职,在业内引发不小讨论。他表示,未来将花更多时间在AI、开源技术教育上,比如他做了一档AI课程,现还在更新中。

此番发布NanoGPT同时,Karpathy还下场安抚了下催更党——新视频正从0开始构建,计划2周内发布

小破本也能炼GPT!300行代码搞定,GitHub获6k星,来自特斯拉前AI总监

最后附上:

NanoGPT项目:https://github.com/karpathy/nanoGPT

Andrej Karpathy课程:https://karpathy.ai/zero-to-hero.html