nanoGPT训练实战：从零跑通微型GPT模型全流程

4月7日我留出了一整天时间专门跑了一次nanoGPT的训练。说实在的训练这种微型GPT模型看起来好像没什么技术含量但真正动手跑一遍你会发现里面的门道比想象中多得多。nanoGPT是一个极简的GPT实现代码量很小完全基于PyTorch训练流程透明很适合拿来验证自己对自回归语言模型的理解。这次训练对象是莎士比亚作品集目标就是让模型学会模仿莎翁的文风然后自动生成看起来有模有样的戏剧台词。如果你正准备入手大模型训练或者想自己跑一个GPT模型的完整流程这篇文章应该能帮你少走不少弯路。我会从环境准备、数据预处理、参数配置、训练监控、常见问题到结果分析一步一步复盘我这次的完整过程。1. 项目概述与环境准备1.1 为什么选择nanoGPT而不是完整的GPT很多人一提起训练GPT第一反应是参数量巨大、显存不够、训练时间按周算。但nanoGPT这套开源代码给了另一条路径用最少的功能实现把一个真正的GPT模型从头训起来。它保留了Transformer decoder结构、自回归语言建模目标、训练与采样功能代码单文件就能看懂非常适合我这种习惯把每个环节都抠清楚的人。用我这次跑的莎士比亚数据集做例子整个训练集大概1MB左右字符集也就几十种模型用几百万个参数就能有不错的效果。相比动辄几GB的通用语料训练省去了数据清洗、分布式并行这些外围工作我可以把精力集中在理解模型本身的行为上。1.2 软硬件环境与依赖版本训练GPT还是要一张能用的显卡显存越大越省心。我这次用的是一块入门级显卡显存12GB全程没有出过显存溢出的问题。如果你只有CPU小数据集也能跑就是速度慢到有点劝退所以只要能借到朋友一台带N卡的机器体验会完全不同。依赖方面我推荐用一个干净的虚拟环境避免跟其他项目冲突。关键库版本大概是Python 3.10以上PyTorch 2.0以上NumPytiktoken用来做tokenizer简单的tqdm可选用于进度显示安装命令我不多说了pip直接装。重点提醒一下PyTorch的CPU版和GPU版一定要区分开否则后面训练一跑起来就发现计算图全在CPU上等半天没动静。1.3 训练目标和硬件规划这次训练我的目标是复现一个最基础的nanoGPT训练流程不追求生成质量多高而是把整套链路跑通。硬件规划上因为数据量小我把训练规模控制在几百万参数级别默认batch size在GPU上跑很轻松整个训练过程不到20分钟就完成了。如果你准备在单卡上跑更大的模型建议提前用watch -n 1 nvidia-smi观察显存占用情况并预留一个能让中间checkpoint落盘的目录。我这次只保存了最好的一个模型文件没有做周期性保存因为数据规模太小没必要。2. 数据准备与预处理2.1 数据集选取与下载我选的是莎士比亚数据集这是nanoGPT项目自带的一个入门数据也是训练文本生成模型最经典的语料之一。原因很直观数据量小、语言风格鲜明、字符集基本固定训练出来的模型能被直观“看出来”是不是学会了莎翁的语气。下载这个数据集最方便的方式是从公开的模型库里拿原始文本或者直接用项目里自带的下载脚本。我建议不要选那种需要额外注册才能下载的渠道直接把txt文件下载到本地然后放到data/shakespeare目录下。如果是别的语料记得先统一编码最好全部转成UTF-8避免后面tokenizer遇到乱码。2.2 Tokenizer的选择与实现细节nanoGPT有两种处理文本的方式一种是用tiktoken这种BPE级别的tokenizer另一种是字符级tokenizer。字符级最简单就是把文本按字符拆开建立字符到整数的映射表。我没有额外引入复杂的tokenizer因为莎士比亚数据只有1MB字符表只有不到100个字符用字符级训练足够。自己写字符级tokenizer最关键的一点是字符表必须从训练集和验证集里一起建立不能只从训练集建立否则验证集里出现训练集没有的字符推理时会直接报错。我这次把整个文本拆成训练集和验证集之后再去构建字典保证两边都能覆盖。2.3 训练集与验证集的划分数据划分我用的是最简单的办法把全部文本按顺序切成长度相等的两段前90%做训练后10%做验证。注意这里不能随机打乱再切因为文本本身是有顺序的随机打乱会破坏句子间的连贯性模型就学不到上下文关系了。切分完之后我把两段文本分别编码成整数数组再用PyTorch的Tensor封装好。nanoGPT代码里有个DataLoader的逻辑每次随机从训练集里抽一段固定长度的上下文再预测下一个token。所以本质上它并不是把一个长文本直接输进去而是不断随机截取片段来增强学习。3. 模型配置与训练参数3.1 模型结构中的关键参数我用的模型结构是标准的GPT decoder里面几个核心参数值如下n_layer4Transformer块的深度层数越多模型能捕捉到的语义层次越复杂n_head4多头注意力头数每个头负责关注不同位置的关系n_embd128embedding维度表示每个token被映射到多少维的向量空间block_size64上下文窗口长度也就是模型最多能看到前面多少个token这个配置大概是几百万参数量对于莎士比亚这种小数据集已经足够。如果你增加n_layer和n_embd生成质量通常会更好但训练时间和显存占用也会同步上升需要自己权衡。3.2 训练超参数的选择与调优思路训练参数主要看这几个batch_size12每个batch里同时训练12个样本max_iters1000总共迭代1000次learning_rate3e-4初始学习率。nanoGPT代码里面有学习率warmup和余弦退火的设置这里我直接沿用默认值。选学习率是门学问太大会导致loss震荡太小则训练半天不见loss下降。我习惯先跑一个比较小的实验比如200次迭代观察loss曲线再决定要不要调低或者调高学习率。这次用3e-4跑下来loss走势很稳就没有再动。3.3 权重初始化和检查点逻辑权重初始化部分nanoGPT用了常规的正态分布初始化模型本身对初始化有一定容忍度。我唯一调整的是把损失函数里的ignore_index设置成-1这样在计算loss时能屏蔽掉padding token。检查点保存我建议除了保存最终模型也保留一个最近迭代的记录以防训练中途崩溃。nanoGPT代码里提供了torch.save的逻辑保存的内容包含了模型参数、优化器状态和当前迭代数。恢复训练时把三者一起加载回来是最稳妥的。4. 训练过程与监控4.1 训练启动与日志输出我用的命令是python train.py直接跑nanoGPT的主训练脚本。启动时终端会每隔一段时间打印iter、loss、time这些信息。刚开始的几次迭代loss在4.0左右几百次迭代后能降到1.5左右这个过程我反而觉得是整个训练里最有成就感的时刻。日志格式里有一个val loss这是在验证集上算出来的loss用来判断模型有没有过拟合。当训练集loss下降但验证集loss不再下降甚至上升时就是过拟合的信号。我这组数据本身较小最后验证集loss在1.4左右和训练集loss差距不算太大说明模型还有继续训练的空间。4.2 loss曲线分析与过拟合判断我将每次打印的loss记录到本地然后用一个小脚本画出来。训练loss整体下降明显验证loss初期下降后期开始平缓这是一个比较健康的曲线。如果你发现验证loss突然掉头往上通常的做法是提前打断训练减小学习率或者增加一点dropout。这次的模型正则化项我用的是dropout0.0因为数据量很小dropout反而可能干扰收敛速度。如果你用大语料建议把dropout调到0.1-0.2防止过拟合。4.3 训练时间与资源消耗整个训练用时大概十几分钟这个时长非常适合快速迭代。GPU利用率我并没有刻意追求100%因为数据量小很多时间可能浪费在数据加载上。如果后续要放大到更大数据集CPU数据预处理的耗时就会凸显出来可以考虑用num_workers参数加速或者提前把数据全部加载到内存里。5. 常见问题与调试心得5.1 显存不足的应急处理我这次虽然没有遇到显存爆掉但以前在其他项目里踩过坑。如果你在训练时突然报CUDA out of memory第一反应不要急着换显卡而是把batch_size调小一半比如从12调到6然后再试。如果显存还不行就把block_size调小比如从64调到32牺牲的是上下文长度但至少训练能跑下去。5.2 训练不收敛与loss震荡的排查有时候你会看到loss一直在某个区间来回跳不往下走。这种情况我通常会先看学习率是不是太高尝试降到1e-4。再看数据有没有问题比如tokenizer里出现了大量重复字符导致模型学不到有效结构。另一个容易忽略的点是学习率warmup在训练初期先用一个很小的学习率然后逐步升到目标值这能避免一开始参数抖动太大。5.3 生成结果太差的原因分析训练完成后我用sample.py生成台词但前几次结果完全不像英文更像随机字母串。排查下来发现是因为我用了字符级tokenizer输出的确实是字符序列看起来自然没有词级那么“通顺”。不过对于小数据集字符级模型已经很正常只要结构上符合语言模式比如有标点、空格、常见词缀就算成功了。如果你连基本单词都拼不出来大概率是训练的模型还没有收敛或者采样时温度设置太低导致输出太机械。我一般是设置temperature0.8在随机性和稳定性之间取一个平衡。温度越低输出越保守温度越高输出越疯狂。0.8左右是个不错的起点。5.4 训练流程中的时间成本优化很多人觉得训练大语言模型时间成本很高但在nanoGPT这种规模下完全不用担心。我建议新手把精力放在多跑几组不同参数上比如变换n_layer和n_head观察它们对loss的影响。每组实验也就多花一顿饭的功夫却能让你对模型容量有直观感受。参数项我这次的值调试建议n_layer4从2开始逐层增加观察loss变化n_head4保持与n_embd比例协调n_embd128太小学不到好特征太大训练慢block_size64可与数据集平均句子长度对齐batch_size12显存不足先调这个learning_rate3e-4波动大就降到1e-4dropout0.0大模型建议0.1-0.26. 结果展示与经验复盘6.1 生成结果的有效性判断我保存了最终模型然后用一段随机起始文本去采样输出的文本虽然离莎翁原作还有距离但已经能看出明显的英文结构片段。比如会生成类似“Now is the winter of our discontent”的句子结构只是内容会变得奇奇怪怪。这种结果说明模型已经学到了词汇搭配和长短句节奏只是还没学到深层的语义逻辑。判断文本生成质量时我建议不要只看是不是“完美”而是看有没有特色。从无到有让模型具备某种风格特征这就是一个里程碑。如果你看到输出文本里反复出现同一个词说明模型可能学得过度拟合了高频词汇可以增加温度或者降低学习率再训练一次。6.2 我第一次踩坑的复盘这次训练整体很顺但开头其实也翻过车。第一次跑的时候我直接把整个莎士比亚文本加载进内存由于数据预处理管道没走对导致训练时读取单个文件每次都重复做tokenization训练速度慢得离谱。后来仔细读了代码才发现应该把编码后的数据保存成二进制格式接下来训练时直接load。另一个小坑是把block_size设得太短只有32生成出来的句子特别容易逻辑断裂上下文接不回来。后来改成64效果提升明显。对于短文本语料窗口大小反而比模型层数更影响结果。6.3 仔细分析训练记录的后续扩展思路跑通一次nanoGPT训练之后后续可以扩展的方向其实很多。比如把莎士比亚换成某一领域的业务文档尝试让模型学会写报告摘要。或者把字符级tokenizer换成BPE tokenizer看效果提升多少。又比如换一个更大的上下文窗口测试模型对长距离依赖的记忆能力。我个人的习惯是每一次训练结束之后都会保存好训练日志、模型检查点和生成样本然后统一放进一个实验目录里方便对照不同参数的效果。这种做法不花时间但复盘的时候能省很多事。最后分享一个小技巧训练时别只盯着验证集loss有时候看模型的采样结果更有意义。loss指标下降只能说明模型学到了数据分布里的统计规律但生成结果才能直观告诉你模型是不是真的“说话有模有样”。这次训练我最大的心得就是小模型也有小模型的乐趣过程清楚可调可控特别适合用来建立对GPT家族的直觉。