nanoGPT用红楼梦数据从头训练babyGPT-12.32M实现任意问答

1. 引入

大神karpathy从openai离职后，创办了AI教育公司Eureka Labs（参考1），同时也创办了知名的nanoGPT项目。
目前，使用nanoGPT（参考2），你可以在几分钟内训练出一个babyGPT（甚至于在CPU上就可以运行训练过程），也可以在8XA100 40GB的GPU上，从头开始训练了GPT2-124M（大约4天时间）。

但是，nanoGPT提供的训练代码，是用英文的莎士比亚书籍来训练的。如果想用中文语料进行训练，需要怎么操作呢？

2. 训练babyGPT-12.32M

本文使用的具体代码和数据，均已开源，详见参考3的链接。

1. 环境配置

本文使用python3.11，服务器是ubuntu 12.04，GPU为1张A800，CUDA版本是12.4。
需要安装如下依赖：

pip install torch numpy transformers datasets tiktoken wandb tqdm

2. 数据准备

• 上传《红楼梦》文本到data/hongloumeng_char_local/input.txt
• 去掉原始代码prepare.py中的网络下载数据的逻辑
• 最终完成的代码见：
  • https://github.com/ybdesire/baby_gpt_hongloumeng/blob/main/data/hongloumeng_char_local/prepare.py

运行如下命令

python data/hongloumeng_char_local/prepare.py

会在数据目录生成三个文件：‘train.bin’, ‘val.bin’, ‘meta.pkl’。分别为训练集、测试集和“汉字/字符到整数值的映射关系”。

同时，也会输出本次处理得到的词表等数据情况：

vocab size: 4,435
train has 791,045 tokens
val has 87,894 tokens

可以看到，《红楼梦》使用了4435个不同的汉字/字符，训练集有79w的tokens，测试集有87k的tokens。

3. 修改训练配置

在原始给定的`config/train_xxx.py’中，修改如下几个字段

out_dir = 'out-hongloumeng-char-local' # 训练得到的模型保存的目录
wandb_log = False # 不使用wandb
dataset = 'hongloumeng_char_local' # 数据集所在的目录

其他配置见：

• https://github.com/ybdesire/baby_gpt_hongloumeng/blob/main/config/train_hongloumeng_char_local.py

4. 训练模型

运行如下命令，即可训练模型

python train.py config/train_hongloumeng_char_local.py

训练过程中，输出如下log：

tokens per iteration will be: 16,384
found vocab_size = 4435 (inside data/hongloumeng_char_local/meta.pkl)
Initializing a new model from scratch
number of parameters: 12.32M
num decayed parameter tensors: 26, with 12,418,176 parameters
num non-decayed parameter tensors: 13, with 4,992 parameters
using fused AdamW: True
compiling the model... (takes a ~minute)
step 0: train loss 8.4504, val loss 8.4506
iter 0: loss 8.4499, time 16598.15ms, mfu -100.00%
iter 10: loss 7.6026, time 29.95ms, mfu 14.21%
iter 20: loss 6.8748, time 30.21ms, mfu 14.19%
iter 30: loss 6.0137, time 13.99ms, mfu 15.82%
iter 40: loss 5.6637, time 30.12ms, mfu 15.65%
... ...
iter 4890: loss 1.4662, time 30.37ms, mfu 16.18%
iter 4900: loss 1.4400, time 15.54ms, mfu 17.30%
iter 4910: loss 1.4483, time 26.49ms, mfu 17.17%
iter 4920: loss 1.4610, time 30.31ms, mfu 16.86%
iter 4930: loss 1.4936, time 25.88ms, mfu 16.82%
iter 4940: loss 1.4505, time 30.22ms, mfu 16.54%
iter 4950: loss 1.4753, time 30.46ms, mfu 16.29%
iter 4960: loss 1.4595, time 30.15ms, mfu 16.07%
iter 4970: loss 1.4530, time 30.21ms, mfu 15.87%
iter 4980: loss 1.4654, time 30.35ms, mfu 15.69%
iter 4990: loss 1.4377, time 30.32ms, mfu 15.52%
step 5000: train loss 0.7242, val loss 4.7782
iter 5000: loss 1.4672, time 3471.11ms, mfu 13.98%real    3m49.327s
user    3m41.666s
sys     0m13.373s

训练迭代5000次后结束，loss从train loss 8.4504, val loss 8.4506降低到train loss 0.7242, val loss 4.7782。MFU（Model FLOPs Utilization）模型算力利用率，并不高，只能到16%左右。

大概4分钟后，模型训练结束，会得到checkpoint文件out-hongloumeng-char-local/ckpt.pt。

5. 加载模型并输出预测值

运行sample.py来加载模型，使用默认输入“\n”，让模型输出预测值。运行命令：

python sample.py --out_dir=out-hongloumeng-char-local/

可以得到如下10条输出（同时也能看到参数量为12.32M）：

Overriding: out_dir = out-hongloumeng-char-local/checkpoint = torch.load(ckpt_path, map_location=device)
number of parameters: 12.32M
Loading meta from data/hongloumeng_char_local/meta.pkl...说话至次日，王夫人只见她姊妹来了，乃向王夫人说：“王善保家的也有些年纪，如今将来大舅老爷、太太太都是些一个人。比那边还有一个。倘或是这几次的人，若是不能，必是亲戚的。我到底是我们府侯门上的，谁不是那里的绣房，见的，说是我们的。”贾珍道：“你素日绣房绣房绣花，皆是个老子。如今我父母、太太太三人和上下，如今虽大开得十分例，人也不能使用。这原是老太太的，何不知道呢，也是一件事的。太太还有年轻，老太太是能够呢，也该问上了。”凤姐儿笑道：“那王大爷，你们的也有原故。你今年来这几岁的时候，怎么就有了？”王夫人笑道：“为什么这话？”凤姐儿道：“‘老太太太不大好，不如你老爷做儿。我老太太太也要提起，所以老太太太太上过去，才敢不依的。’”王夫人陪笑道：“我们想的事也是老太太太的，也不敢争一点儿，只是要亲戚戚们几个人，都是过一个心。”贾母道：“说着，我们也不是没话，原该教我一个女孩儿，就是一个妈，他们也不肯说他，只是挑不得妈就是了。如今凤丫头们的人心，虽说，只是小儿子们的体面些，这个大，是和我们要子似的，就是好。”王夫人点头道：“那刘姥姥，见了这个，倒像我们姊妹们的丫头子前年轻，也没有，所以没有见过。
---------------谁说贾政在书房，一壁上。这几处原是贾政的，也有很爱，因向着手中告说：“只怕明年月上学，不大好。”贾政道：“这个帖子，是极好，有个做过的，也有什么。”贾政道：“这也无用的，如今也有唤人去要等看。”贾政问他甚不已，只得答应了。贾政便将宝玉送上，宝玉一一一一瞧，十分喜，见这般光景，心下便一发烦闷，只说：“这就好了。”正嚷着，只见贾政在书房中看书，便往王府中去请安，回到衙门外书房。那时，贾政本府上司员代儒只有许起跟来，知县便叫人查，贾政不出官。贾政便说：“政老太贾政那时在贾政火上，但是强凌弱，所以不用。说来，余者也不好，内牵不住，暂且停放心。没有问家资并，便为难保家。你们那知道外头这些时候，不用说，只就是珍儿子是办家，不用无所花钱，也便将贾琏的事暂且办了。”贾政听了道：“那时常见的事，并没有才能回。惟是元妃薨逝’此，又好为晚了，不顾那一场。如今我去了，回了一个信，瞧着他，老太太太那里，有个应妥当。”贾母等道：“你们老太太这里总怕，什么意思？我想起那里办事，太太和老爷的事，二爷也不大好？”贾政道：“我到底是什么事？”贾政本是心里很大，便不言语。这里贾母知宝玉回来，为众亲事，这事情早，忙
---------------紫鹃进来，笑说道：“姑娘听得了，也别混搅，咱们一个人家的人，都不敢撒谎。”黛玉笑道：“就是那里说话的，也是一个人，你倒说已有的，不知道。”黛玉笑道：“你只说，就是有了。虽然如此，也省些什么！”黛玉笑道：“你也不是，就是我也是嘴里不成？”紫鹃道：“这是什么原故？”说着，紫鹃连忙去了。宝玉回房，急忙去见黛玉，如此说，忙回房中来，仍到怡红院里来。这里宝钗等等都回去了。宝钗见了，忙的也不好再作声，只是呆呆的呆呆呆的说了。黛玉正在书房中，进这里。那紫鹃自吃饭，只见紫鹃在床上，黛玉便伏侍书坐下。黛玉进来用手帕子，问道：“姑娘为什么得？”黛玉便将她扶住，将手一抓起，向黛玉身上努开，被黛玉不理，走上来，连忙起来，黛玉便拿眼睁开看，不答。此时黛玉便睡下。宝玉正眼睛里正自红了，旁还未瞧，一面笑，看那女孩儿微微睁眼，那紫鹃的跑来，连忙一把拉住，自躺下。紫鹃便坐下，只问：“姑娘怎么样？”黛玉向雪雁道：“姑娘请明儿坐定罢。”紫鹃道：“老太太醒了，不该睡了，只怕我们姑娘身上觉，气得浑身一阵，也睡觉，也该睡不着了。”说着，便起身来。紫鹃走进来，被紫鹃说道：“姑娘是躺的，等我倒服侍二奶奶，你来了再去。”黛玉
---------------众人正在房中，只见宝玉走来，说：“林姑娘来了。”宝玉正笑着，见赵姨娘来了，忙回身陪笑让坐。宝玉忙问好。芳官便问：“我倒来。”宝玉笑道：“我只从园子里去，找谁去？”芳官笑道：“你倒瞧瞧，我们都找我。我可不是找人了？”宝玉见是环儿的，便笑道：“你就走得了，我这里就不去了。”说着，便将宝玉的手一细与宝玉。袭人送过来，便推开，笑道：“什么呢？你这会子后不是玩的，怕人厌可怜！”宝玉笑道：“这有趣，又不说什么。明儿拿你怎么拿我送来？”袭人笑道：“这两瓶上，你放去罢。”袭人笑道：“我送你的，一个人给你的当人给芳官给我，你不认我。”宝玉笑道：“我给我，你同他们给他。”袭人笑道：“你留的，就是了。”说着便走了出来，大家都笑了。宝玉笑说：“我也不希罕，若果真没味，我就拿出来了，我只当我，明儿照看不是了。”宝玉笑道：“你们不认，我还赖，难道喜欢呢！”史湘云笑道：“这一点子戏我的事我，你就不认。她们嘴里就说得，你们不识。”袭人笑道：“你不会说这个。我也不信，你放嘴。”正说着，只听外面说话，晴雯说：“这会子也值什么有个清人，你们也不必说。”宝玉笑道：“她们家哪里的画儿，不过是我们家里一画的，我少不得她
---------------这里尤氏、凤姐儿都来了，贾母等在地下跟前说笑。早有贾母等都让坐，贾政带着人往稻香村去，尤氏方回至回桂花厅之后，尤氏便跟了王夫人，只命人向凤姐儿说：“好生养！”贾母听说，只得耐入园，因说笑道：“且别委曲。我再议去，等我们都随便去了。”众人听说，大家都笑起来。贾母等都笑了。于是大家行礼，贾母等都吃过。便邀了一杯酒，然后果品赏酒。贾母传于是众媳，都上席该行礼，贾母等赏，方是湘云、薛姨妈、薛姨妈等归坐。贾母便说：“正是呢。”贾母也起，因说：“我们晚上懒们，我们不吃酒，不如在老内开散散酒，不过去了。”说着，便一个老婆子，笑个不住，说：“老太太、太太太们，这里头的姑娘，年轻，也不过虑了。”贾母等忙笑说：“你姨妈今儿大了，才刚老太太太瞧瞧我这里的，所以都是个好孩儿的。”薛姨妈都笑道：“我也没有，听见这样的没有。”薛姨妈笑道：“你瞧瞧瞧，原来是个老婆，不是几年了他的人，大家也要告诉了她，不如大家的好话。”贾母笑道：“这老太太的是，我们不是，她们家的事，亲戚们也是不犯着我们。”贾母笑道：“我们姨太太的人，为你的。”贾母笑道：“我们的你们是哪里的？你这会子又弄个死有个呢？”凤姐儿笑道：“这样的不
---------------宝玉忙赶上坐着说：“好妹妹，我不信。”黛玉笑道：“我倒是个好的。我又不是还干这个。”黛玉笑道：“你放心，别在我跟前。我再要那些人说话，不过是想想的，我也忘了，不如你们又没个好人胆量的。”宝玉听了，心下是喜悦，只说道：“姨娘儿也别管说，他们如今也该打他们的。”一面说，一面就赌气去了。不提。且说宝玉又不好生气，待回至园中，便回过园来，看见袭人，一面说道：“你们不知道，这会子来了没来了？”袭人又笑道：“你这两个又没听见他们，可是袭人来的。”宝玉笑道：“我说原来是为什么？”宝玉笑道：“可不是，若论起。”黛玉笑道：“他说怕什么？”宝钗笑道：“既这样说话。”宝钗笑道：“我就这样方才老，难道我正想得了你，你也最喜欢。”袭人忙笑道：“你这个放心，这个就奇了。”宝钗笑道：“我倒是个呆人，想来也罢了。”宝钗笑道：“我取明儿你的，错一个的小子，又说这话没什么，又说我笑，你也为什么可恶，再不露了！”宝钗笑道：“这话奇怪，‘女人家不知道，怎么就不知道了！你是管谁家婆婆子，家里的男女，管，亲戚们的，原是个亲戚们，也不知道，我并没法的。何况且你们姑娘这样大正配人，你也就是‘天天下，自然’，不过是‘贾’‘只是我
---------------宝玉只装病了，欲往前走，也不敢提。宝玉越发没意，只得又说道：“好妹，你是个怎么好的？”宝玉笑道：“走罢！”袭人笑道：“你别忒臊了。”宝玉听说，忙起身要走，因问她道：“我说你再打不明白，我打我的。”宝玉笑道：“再说了酒，大家随你去。”袭人道：“我不信你说，我已这么说了，你就信了。”袭人笑道：“你说的不管，还有我呢。我也不必说了，你又来找你。”宝玉笑道：“你还敢妄想了。只管怕你们去，只是跟你们去。”宝玉笑道：“既这样，我有多少精打趣，你岂不没人的？你又不来，你们两个去了，又担心，不如你还要教我们受我，倒不如一个好有本无闻。我到了我们这里，你就不如拿我了。也不用放心了，反被宝玉瞒了你。若要我也给他的，也是我不给你，再是有谁？”袭人见问，忙笑道：“好姐姐，你也不必饶人了？”宝玉笑道：“你是那些丫头，有我们没有？”宝玉笑道：“我不像你们这样的，这会子又有什么呢？我也不过是有了。”宝玉笑道：“有本事原故，有些没有，也有。”袭人笑道：“倒有了。”袭人笑道：“我原有话，有了。但你听见你病的话，你也不是，却真心里疑人。古人也有的，‘不足畏’，不知他不过，反倒想猜着了。”因笑道：“这也奇怪，你也忒胡涂
---------------宝玉在里间相见了，便出来望外叫：“宝玉”。宝玉只得含着泪，只得含泪说道：“老爷真真是来了！”宝玉低了头，含泪说道：“你们大了，我如今虽是个狐媚子，却不敢拿我，如今你何敢和我什么相干？况且看他是喜欢的，如今闹出事来，还有几句话，哪里禁得起一句话来？我就是知道他说的，你就是是浊物，我也不知道了。”贾政听了，心委屈，只得答应了。未知何事如何，下回分解。第三回  得情思女情娇敏　慕情男子却说宝玉感忿，喜庆之不必细说。贾母见他母亲都来了，心中便含泪来，忙出来，命宝玉到贾母跟前，问长问好，宝玉好端的说去了。宝玉不自在房中，便和袭人、麝月等好端了衣服，一直到门口，对袭人等只说：“你且拿这个来给你们敬我们磕头，不要打动亲自了。”袭人忙答应了，仍往王夫人处来。宝玉只得依次依次面，因笑道：“咱们先睡罢。”王夫人答应了。宝玉道：“你别自在家，也不用进来。那嘴里不说了。”麝月笑道：“你说哪里的话，还敢要我说，少不得我素日教导她作这么！”宝玉忙陪笑道：“你这会子又比我的不烦恼。”贾母笑道：“我不知道你说得固然。我这话说的是正配心的，这也是有的。”袭人笑道：“倒不大。只是从此以后拿来的，不如毁了
---------------话说着，此时贾母等过来，也不叫他来，将贾赦那一席放在桌上，贾母便命过来。可巧凤姐儿等都在那里吃饭。贾母又命凤姐儿过来吃。一时到了席，凤姐儿因笑说：“这是咱们家的家，你们家这一齐去罢，都坐去罢，咱们别家说话儿。”贾母笑道：“好，我们家去了这里没有的戏，也有不过闹了，你们倒去受屈了，坐得什样就是了。”贾母听说，忙笑道：“你们别说了。”李纨笑道：“正是，我们闹的没有。不如你们熬了，我们也没个法儿。”贾母笑道：“这里都是些好的，还不如何办！”贾母笑道：“我们家的一个，也并不是清净。”贾母笑道：“这样正好，老太太爱也没的。”凤姐儿笑道：“这样，就是我们这边，老太太太太既高兴了。”贾母笑道：“你是她，不在这里，不如我，也叫她们受不得。”贾母笑道：“我也不会打发她去，让她们再吃去，岂不省事。”众人都笑说：“姨太太说得何妨。”邢夫人笑道：“老太太真个更笑话。”贾母笑道：“老太太太知道，你我们不管，只是二姑娘那里不大喜欢的，吃多吃一杯不好，吃了半碗酒，多喝一杯。”众人都说：“这是老太太，这会子又该罚了。”贾母便说：“不怕，不用就是了。”凤姐儿笑道：“你说笑，我哪里有趣，这个刁话只是众不公。”贾母笑道：
---------------一时，宝玉进入房中，大家安寝之过，丫鬟们皆在房中。其余者，因宝玉在内玩耍。宝玉心中虽在盘算，不大听见，只说：“宝玉！你们三人多吃酒，你们只怕在一个榻上，又恐脏了，我和我拌嘴。”因说道：“叫你们拉扯她们的嘴，我和我们说话，有我们眼里不过，我们就不依，你是好的罪人，她自己就是个不是。若不叫她们在家里，也不可以使唤。若是她们外头的，说我们也说：‘不用说’，要了，把你们的五个一吊钱就完了。”王夫人道：“这都是我的。”宝玉笑道：“你们说只管说，多少胡涂东西。这个原该取来的。”袭人笑道：“这事原是‘天下人’，你们也不敢回去。你且别说我们没有人知，反多拿你们了。”宝玉笑道：“哪里不为这个呢？”王夫人笑道：“不过是，你们每人一个人，难道也厌。”袭人笑道：“你是个刁钻古怪。又多少着这个，如今也太多了。”宝玉道：“你这傻丫头最小性儿，外头，你不曾不如把我的话教导过。殊不知她不知道的性人，你还不知道，在跟前，也是我跟前说话了。只和我说是个光景，也不知是三天的人，等也是和老太太一样，也要惹的。这是姑娘，人家也还有些体面，不如没有的，年轻的，你也还有这样体统。你这话说的，素日我虽素日看不单，还敢和我说，不比咱们
---------------

3. 任意问答

本文使用的具体代码和数据，均已开源，详见参考3的链接。

nanoGPT给定的sample.py，给定了一种实现任意问答的配置。比如，将提示词写到文件prompt.txt中，让后修改sample.py的配置如下：

start = "FILE:prompt.txt" 
num_samples = 1 # number of samples to draw

这样指定从prompt.txt文件中读取提示词，并进行后续的文本生成，且只生成1条文本。

要注意这里提供的提示词，不能超出词表（2.2中，vocab size: 4,435）中的字符。

提示词为：“黛玉一边吃瓜子一边吃肥肉”。运行如下命令

python sample_prompt.py --out_dir=out-hongloumeng-char-local/

得到生成的文本：

黛玉一边吃瓜子一边吃肥肉黛玉只管吃着饭，因向紫鹃道：“你们去罢，有话说说的？”紫鹃道：“你不必说我们大姑娘闲话呢，让我听。”黛玉道：“那里吃茶，只是等罢。”黛玉忙起身走开，回身道：“姑娘也睡了。到底是我们姊妹们歇歇歇儿，只怕是等你们，我再来罢。”黛玉笑道：“你们只管去罢。”宝玉答应了，方起身走入房中。宝玉便往他房里去时，见袭人等都来了。宝玉笑道：“我们不用进去。”袭人笑道：“我们所以如此想。”袭人忙扶了她坐，至炕上，见宝玉在那里，便端水。宝玉便自己坐了两个大桌，一径往怡红院中来。宝玉因向黛玉道：“你们去了，还没有？你们三个人在这里，又不能进去，因此咱们回房里没有？宝玉读书，你们也可大好？”宝玉道：“你们先别叫起，我们天天天才睡一夜的，也不能睡觉，等我再说话儿，我再不必叫你们听。”宝玉道：“我不信。你既这样说，横竖不如我明白你们几个照看，这会子也乏了。”宝玉笑道：“你们可别笑话呢。”宝玉笑道：“你们只顾坐才说话儿罢。”宝玉听了，解手笑道：“你们都说话，等我再说。”黛玉笑道：“往哪里去？”宝玉道：“且不必说话，只是你无缘了。”黛玉道：“你也不必委屈了。”宝玉道：“我替你解闷儿。”黛玉道：“你们没有，就只

4. 总结

训练babyGPT-12.32M，使用GPU，以《红楼梦》为训练集语料，只需要4分钟，就能完成。也能实现（效果很差，输出难以控制，但语气上确实有那么点意思的）AIGC。

5. 参考

1. karpathy个人经历。https://karpathy.ai/
2. nanoGPT开源项目链接。https://github.com/karpathy/nanoGPT
3. https://github.com/ybdesire/baby_gpt_hongloumeng