NanoGPT Tutorial

来自WHY42


环境准备

本文所有操作均在MacBook Air(macOS 13.5.1,2020版,M1芯片)和OptiPlex 7080(Linux Mint 21.2)上测试验证。 您也可以在其他的系统上运行,只需要在安装conda时按照官方文档稍作改动即可。

安装Miniconda 和Python

在MacOS下,可以通过以下脚本安装[1]

$ mkdir -p ~/miniconda3
$ curl https://repo.anaconda.com/miniconda/Miniconda3-latest-MacOSX-arm64.sh -o ~/miniconda3/miniconda.sh
$ bash ~/miniconda3/miniconda.sh -b -u -p ~/miniconda3
$ rm -rf ~/miniconda3/miniconda.sh

安装完成后,可以使用conda命令来管理机器学习的Python环境了。默认系统会自动创建一个Python3.11的环境:

$ python --version
Python 3.11.5
$ whereis python
python: /Users/riguz/miniconda3/bin/python

安装PyTorch

由于nanoGPT是机遇PyTorch的,因此需要安装它。在Mac上运行时,可以采用nightly版本, 因为更新的版本会集成Mac自带GPU的支持,“可能”性能上可以得到提升。

$ conda install pytorch-nightly::pytorch torchvision torchaudio -c pytorch-nightly

下载nanoGPT

nanoGPT依赖于一些Python的软件包,在运行之前应该首先进行安装[2]。 这里直接通过pip进行安装即可, 为加快下载速度,可以选择使用国内的pip源,如清华镜像(-i https://pypi.tuna.tsinghua.edu.cn/simple)。

$ pip install torch numpy transformers datasets tiktoken wandb tqdm -i https://pypi.tuna.tsinghua.edu.cn/simple

接下来将nanoGPT克隆到本地:

$ git clone https://github.com/drriguz/nanoGPT.git

运行“莎士比亚”例子

为了测试nanoGPT是否能正确工作,我们可以运行它自带的莎士比亚的例子。 该例子通过一个两三万行的莎士比亚剧本作为训练语料,大概长这样:

First Citizen:
Before we proceed any further, hear me speak.

All:
Speak, speak.

First Citizen:
You are all resolved rather to die than to famish?

All:
Resolved. resolved.
...

运行这个例子分为三步:

  1. 准备数据,通过data/shakespeare_char/prepare.py脚本将语料加载到本地,并将数据切分为训练集和验证集,供训练使用
  2. 通过train.py进行训练
  3. 训练完成后,通过sample.py运行查看结果

首先,数据准备非常简单,直接运行prepare.py即可。 该脚本会将数据拆分成两部分,一部分用来进行训练,一部分用来做验证, 并将原文分别进行简单的编码,穷举所有字符并制作一个映射表通过序号进行映射,最后生成包含这些编码序号的文件train.bin和val.bin:

$ cd nanoGPT
$ python data/shakespeare_char/prepare.py
length of dataset in characters: 1,115,394
all the unique characters:
 !$&',-.3:;?ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz
vocab size: 65
train has 1,003,854 tokens
val has 111,540 tokens

然后,开始真正的训练过程。由于MacBook上没有支持CUDA的显卡,只能使用CPU推理:

# 如果在不支持CUDA的设备上运行会报错:
# $ python train.py config/train_shakespeare_char.py
# raise AssertionError("Torch not compiled with CUDA enabled")
python train.py config/train_shakespeare_char.py --device=cpu --compile=False --eval_iters=20 --log_interval=1 --block_size=64 --batch_size=12 --n_layer=4 --n_head=4 --n_embd=128 --max_iters=2000 --lr_decay_iters=2000 --dropout=0.0

...
step 2000: train loss 1.7640, val loss 1.8925
saving checkpoint to out-shakespeare-char
iter 2000: loss 1.6982, time 306.45ms, mfu 0.05%

可以看到,train loss(训练损失)会逐渐减少,说明训练起到效果了。 整个过程需要大概1分钟的时间,训练结束后,就可以运行查看效果了:

$ python sample.py --out_dir=out-shakespeare-char --device=cpu
Overriding: out_dir = out-shakespeare-char
Overriding: device = cpu
number of parameters: 0.80M
Loading meta from data/shakespeare_char/meta.pkl...

I by doth what letterd fain flowarrman,
Lotheefuly daught shouss blate thou his though'd that opt--
Hammine than you, not neme your down way.

ELANUS:
I would and murser wormen that more?
...

对于笔者来说,很难判断生成的结果到底有多好,就算拿莎士比亚的原文可能也品味不出来好坏。 但是,至少从形式上来看,生成的结果是令人满意的,至少看起来像那么回事。

至此,我们已经亲手训练并运行成功了nanoGPT,是不是很简单! 理论上如果通过GPU训练会更快一些,但是在MacBook上测试通过mps(--device=mps) 训练的时间与CPU几乎相当。

训练一个中文的nanoGPT模型

从莎士比亚的例子可以看出,nanoGPT训练的过程实际上十分简单,只需要一段文本,就可以从中学习到一些知识。 本质上,学习的是“概率”,即文字与文字之间的关系。 那么,如果我们用一些其他的语料, 是不是就可以得到想要的结果呢?

学习一下启蒙读物《三字经》?

看起来nanoGPT学习能力还不错,那我们不妨来让它学习一下中国小孩的启蒙教材《三字经》吧。 这里我们仿照shakespeare_char的做法, 把训练数据换成三字经:

人之初
性本善
性相近
...

我们期望nanoGPT学习了这些数据之后, 起码能生成像每句话三个字的“三字经”。 使用同样的参数训练一下试试:

$ python data/sanzijing_char/prepare.py
$ python train.py config/train_sanzijing_char.py --device=cpu --compile=False --eval_iters=20 --log_interval=1 --block_size=64 --batch_size=12 --n_layer=4 --n_head=4 --n_embd=128 --max_iters=2000 --lr_decay_iters=2000 --dropout=0.0

...
iter 1998: loss 0.0437, time 23.68ms, mfu 0.06%
iter 1999: loss 0.0535, time 23.63ms, mfu 0.06%
step 2000: train loss 0.0445, val loss 8.9908
iter 2000: loss 0.0524, time 331.35ms, mfu 0.05%

接下来运行一下试试:

$ python sample.py --out_dir=out-sanzijing-char --device=cpu
...

宇文周
与高齐
迨至隋
一土宇
不再传
失统绪
唐高祖
...

看出问题了么?尽管模型生成的的确像是三字经,这根本就是原封不动的三字经! 也就是说,模型并没有“生成”我们想要的内容, 更像是在背诵我们训练时的语料。 在机器学习中,这种现象叫做“过拟合(overfitting)”,当模型尝试“记住”训练数据而非从训练数据中学习规律时,就可能发生过拟合。

通常,出现这种问题是由于数据集过小导致的。我们的三字经总共就三百多条数据(机器学习几百条的数据量一般认为是很少,几千到几万可以算中等[3]),由于数据量太少而导致过拟合。 过拟合时,从训练的损失上也可以看出端倪:

# train loss不断收敛,但是跟val loss差别很大
step 2000: train loss 0.0445, val loss 8.9908

这样训练出来的数据,缺乏泛化能力。也就是说, 对于训练中出现的数据能够给出结果, 但是如果是没有出现过的,就不能得到结果了。为了验证这个结果, 我们可以修改sample.py运行时的start参数,该参数用来控制模型根据什么进行续写:

$ python sample.py --out_dir=out-sanzijing-char --device=cpu --start="养不"
...
养不教
父之过
教不严
...
$ python sample.py --out_dir=out-sanzijing-char --device=cpu --start="养要"
...
养要
记其事
五子者

可以看出,如果是出现过的字,模型能够准确“背出”内容;但对于没有见到过开头,就不知道怎么回答了。

不听不听,和尚念经

由于三字经实在是内容太少,通过这几百条数据很难训练成一个会写三字经的模型。 那么,只要提高内容,多给点训练数据,来个几万几十万条,不是就可以了么? 有什么东西可以找到丰富的语料呢?

经过一点调查,笔者发现佛经是个不错的选择。 一方面佛经很啰嗦,动不动就几十卷; 一方面又很晦涩,就是看也很难看懂,可能甚至比看莎士比亚的英文剧本更难理解。

于是,我们拿《大般若波罗蜜多经》这部也许是最长的佛经来训练一下试试。 训练数据共五万多行,合计约五百七十万字。足够大了!

$ python data/boluojing_char/prepare.py
$ python train.py config/train_boluojing_char.py --device=cpu --compile=False --eval_iters=20 --log_interval=1 --block_size=64 --batch_size=12 --n_layer=4 --n_head=4 --n_embd=128 --max_iters=2000 --lr_decay_iters=2000 --dropout=0.0
...
step 2000: train loss 1.5015, val loss 3.3029

这次训练后,train loss和val loss就比较接近了。 然后运行测试一下效果:

$ python sample.py --out_dir=out-boluojing-char --device=cpu
...
    “不也,善现。”
  “即四念住、四正断、四神足、五根、五力、七等觉支、八圣道支无二分故。”
  “世尊,预流果无生不二。”
  “善现,一切智智无二言定,无二为方便、无生为方便、无所得、无生为方便,回向一切智智,修习五眼、六神通。”
  “世尊,云何以五眼、六神通无二为方便、无生为方便、无所得为方便,回向一切智智、道相智,修习六神通,修习八圣道支?”
  “庆喜,云何菩萨摩诃萨行无忘失法无二无二为方便、无生为方便、无所得为方便,回向一切智智智,修习佛十力、四无所畏、四无碍解、大慈、大悲、大喜、大舍、十八佛不共法、一切法。”
  “庆喜,以四无所畏、四无碍解、大慈、大悲、大喜、大舍、十八佛不共法性空与外空、内空、空、大空、胜义空、有为空、无为空、毕竟空、无际空、散空、无变异空、本性空、自相空、共相空、一切法空、不可得空、无性空、自性空、无性自性空清净,外空乃至无性自性自性空清净故法界清净。何以故?若一切智智清净,若法界乃至不思议界清净,若道相智、一切相智清净,无二、无二分、无别、无断故。
  “善现,一切智智清净故一切相智清净,一切相智清净故无性空清净。何以故?若一切智智清净,若无
$ python sample.py --out_dir=out-boluojing-char --device=cpu --start="如是我闻:一时,佛在舍卫国祇树给孤独园,与"
...
如是我闻:一时,佛在舍卫国祇树给孤独园,与善男子、善女人圣过生人、无为智不染著,或不?”
  “善现,是善男子、善女人等,诸有情毕竟中,不获福聚无量无边有,无数无增语是菩萨摩诃萨不见余菩提,亦不见无所有,亦不见有唯有唯有有所得故人中,法界、法界、意识界,亦不见有见可得、不见识界及身触、意触为缘所生诸受乃至身触为缘所生诸受无染有性故,名为名四念住、四正断乃至十后、中际谓无所有性。
  “善现,如是,善现,如见诸佛如来、应、正等觉、正等觉亦无相,无所有不见。何以故?眼界性空法界、无愿故。色界等无所有故。眼识界及眼触为缘所生诸受无我亦无所有故,受、想、行、识;声、香、味、触、法处无所有故,当知四念住亦无所有;声、香、味、触、法处故,当知作意,当知作意;无明无所有,当知无所有故,善现,当知是为菩萨摩诃萨不作者,亦无所有;法亦无所有故,当知般若波罗蜜多,当知般若波罗蜜多亦无所有。
  “善现,菩萨摩诃萨若善,亦无所有故,当知般若波罗蜜多亦无所有,水、火、风、空、识界无所有故,菩萨摩诃萨行般若波罗蜜多亦无所有;一切菩萨摩诃萨所行般若波罗蜜多,净戒乃至般若波罗蜜多亦无所有;水、火、风、空、识界无所有不可得,声、香、味、触、法处菩萨摩诃萨无所
...

与之前的三字经不同,这次生成的“佛经”不仅看起来像那么回事, 其中很多句子都是训练的原文没有的,是nanoGPT所“臆造”出来的。

但是,对于生成的结果的好坏,难以评判。 也许这些句子是毫无意义的组合而已,但是对于普通人,缺乏相关的佛学知识, 也不知道是自己不理解,还是模型生成的不正确。 这样训练出来的模型大概率只能用来唬人了。

续写《水浒传》

如果把佛经换一下,换成小说,我们就很容易可以判断出生成的结果如何了。 以《水浒传》为例,全文共计四千多行,九十多万字。

$ python data/shuixu_char/prepare.py
$ python train.py config/train_shuixu_char.py --device=cpu --compile=False --eval_iters=20 --log_interval=1 --block_size=64 --batch_size=12 --n_layer=4 --n_head=4 --n_embd=128 --max_iters=2000 --lr_decay_iters=2000 --dropout=0.0
...
step 2000: train loss 4.2661, val loss 4.8593

训练2000步后,train loss下降到4.2661,相对于之前的训练来看这个数值偏高。 一直好奇如果武松来押送生辰纲会是什么效果,现在可以来生成一下试试了:

$ python sample.py --out_dir=out-shuixu-char --device=cpu --start="武松道:“不要问,只管用瓢舀来吃酒。贩枣的客人那里,买三五斤枣 来,一发算钱。”竟抢先要下"
...
武松道:“不要问,只管用瓢舀来吃酒。贩枣的客人那里,买三五斤枣来,一发算钱。”竟抢先要下山来,便道:“这个去说。”阮小二道:“你们相见他。”那妇人只见了一个头,就那里飞入房里,那里取出来酒来。李逵又去看时,那大汉庄上马上见他。背后坐在武松歇下。晁盖道:“你既是叔叔在县里,何不敢归。”只听得茶讨了。王婆娘道:“你且不是不得。”晁盖道:“小人便是这厮们。”店府道:“我自和道:“却有百姓知,有几个教小人说。”
且说当日是两个个都是火的门外里有。二人一日从人中间,出身前来到几个村坊库,自上把手来的,又会了船只一个,便走!李应也把武松,那汉,却是两个去的人出。李逵大听得山寨,便叫道:“小人莫不曾头,且不知!我也不可要来,你知去借了。”,只见了一十五句,早要行不得。”石秀叫道:“押司宋江,不知闲便拜让,他做甚为父。”王伦答道:“兀自去甚事!”酒保道:“小人都是个是这话。”山边山寨,将过过一路。四十数日戴宗,便向前面,下墙边一个不在房。李逵道:“你必是我。”王婆道:“这个是这句话,却不得。”原来到这那婆子里说道:“你却是不要说,昨夜不须!便是五阮哥哥哥娘的?”宋江道:“好好来去走,怎生要走有个。”吃了的也,不要走,只见得衣服却不要寻。杨志便禀道:“你不肯起甚么?”正说道:“嫂嫂两个土兵
...

很容易看出,这生成的是些啥!乱七八糟,前言不搭后语。 也许是训练的不够好?

尝试提升迭代次数

我们试试提高训练的步数,让其损失进一步收敛:

$ python train.py config/train_shuixu_char.py --device=mps --compile=False --eval_iters=20 --log_interval=1 --block_size=64 --batch_size=12 --n_layer=4 --n_head=4 --n_embd=128 --max_iters=20000 --lr_decay_iters=20000 --dropout=0.0
...
step 20000: train loss 3.1365, val loss 4.5919

经过10000步训练之后,损失降低了一些,再来试试效果:

$ python sample.py --out_dir=out-shuixu-char --device=cpu --start="武松道:“不要问,只管用瓢舀来吃酒。贩枣的客人那里,买三五斤枣 来,一发算钱。”竟抢先要下"
...
武松道:“不要问,只管用瓢舀来吃酒。贩枣的客人那里,买三五斤枣来,一发算钱。”竟抢先要下一刀来,和那妇人说道:“我只怕你武头陀直吃了去,不去取笑!”那妇人道:“你猜个不说,我自在在家里,只是如今来却和你好也,也吃他口骂我!我若是你去倒来,你也不还我!”那妇人便道:“郓哥与你说话。小人是个干娘,你却不知来,如何不来欺负他!”武松再笑起来劝他。武松把两口刀来筛了,说道:“我只道是这个罪人,老爷少也不也。”那妇人道:“一个真人真个是谁?”武松道:“三个酒肉吃了一杯,也热了,便去换些香味去里面。”武松道:“这个人姓甚么?”那妇人答道:“主人,我家道个不曾有一个。前日,买了酒果儿送酒归来与我吃了,却暗地下来?”那妇人道:“你这个道童,你来,对我说了。听我叫道:‘我和你押司说么?”李逵道:“也须是。好教他作事,我不得你,须每这些酒食与你吃了些枪棒,提条好去。”婆子道:“你若是他,不怕他,且我如何不叫他有些和他?却怎地戏弄我?”那汉喝道:“你是甚么官人便来。我写一封书,却要使人去店里卖炊饼出卖与你?”武松道:“恁地时,怎地?”那卖话正卿吃得动,只见那和尚入来,便叫道:“大郎已不要吃酒,快去取来,小人便问人家里取笑话。”那妇人便道:“叔叔恁地说谎。”那妇人道:“这个兄弟,又没酒吃,这几日

这个效果很难说好,很难衡量它与之前的生成结果的好坏,但是直观感觉就是都很烂。 看起来,我们的训练遇到瓶颈了。

提升nanoGPT参数并使用GPU训练

由于我们之前在训练时指定了使用较小的模型参数以便减少计算量(例如n_layer=4 --n_head=4 --n_embd=128), 如果把模型参数提升一下,是否能够得到效果上的提升呢? nanoGPT示例中给出了一个“baby GPT”参数值,相比于现有的参数只是稍微有些变化,比如层数从4层提高到了6层:

# baby GPT model :)
n_layer = 6
n_head = 6
n_embd = 384
dropout = 0.2

但是这一些参数的提高,带来的是计算量的极大增强。使用CPU推理时, 以前几分钟可以完成的,现在需要数小时才能完成,为了更快能够验证结果, 我们需要使用GPU来进行推理。 尽管我们没有GPU,但是不要慌,有很多免费的GPU训练平台可以使用,Google的Codelab就是之一。

在Codelab中,新建一个记事本,并将运行环境设置为GPU("代码执行程序“->"更改运行时类型")。 Google对于免费用户提供了Telsa T4的GPU选项,其显存为16G, 已经是个很不错的GPU了。

Codelab默认就支持Python环境,因此对于我们的训练来说开箱即用。不过首先,我们来确认一下GPU是能够正确检测到的:

!nvidia-smi
import torch
torch.cuda.is_available()

注意这是一段shell和python混排的命令,以“!”开头的为shell命令,其他为python脚本。 如果配置正确,就可以检测到T4 GPU,并且torch.cuda.is_available()返回结果为True。

Tue Dec 12 05:09:25 2023       
+-----------------------------------------------------------------------------+
| NVIDIA-SMI 525.105.17   Driver Version: 525.105.17   CUDA Version: 12.0     |
|-------------------------------+----------------------+----------------------+
| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |
| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |
|                               |                      |               MIG M. |
|===============================+======================+======================|
|   0  Tesla T4            Off  | 00000000:00:04.0 Off |                    0 |
| N/A   42C    P8    11W /  70W |      3MiB / 15360MiB |      0%      Default |
|                               |                      |                  N/A |
+-------------------------------+----------------------+----------------------+
                                                                               
+-----------------------------------------------------------------------------+
| Processes:                                                                  |
|  GPU   GI   CI        PID   Type   Process name                  GPU Memory |
|        ID   ID                                                   Usage      |
|=============================================================================|
|  No running processes found                                                 |
+---------------