——7.5 小时实现 1.23 亿参数分布式预训练文章目录准备工作注册 Kaggle 账号已有可跳过注册流程开启 TPU Runtime运行TPU 环境补丁必须最先执行踩坑修复一参数共享问题踩坑修复二BF16 混合精度训练突发数值异常问题QK-Norm 是什么踩坑修复三OOM不足之处训练你自己的 GPT总结参考资料你是不是也想过自己亲手训练一个GPT网上搜一圈教程热血沸腾地点进去结果发现要么是拿一本小说训个玩具模型生成三句话就开始胡言乱语要么标题写着对标 GPT-2看了半天发现需要 8 张 A100——直接劝退。大模型训练好像天生就是富哥们的游戏。好消息是我帮你们解决了这个问题。这个项目就是用 Kaggle 免费的 TPU v5e-88 个加速核心从零实现 GPT-2 Small 级别的分布式预训练。参数量 1.23 亿12 层 Transformer训练数据是 FineWeb-Edu 10B——真正的互联网海量数据不是那种拿一本小说训着玩的 demo。全程免费7.5 小时出结果你换上自己的语料也能训一个属于你自己的 GPT。项目地址https://github.com/suxuefenxi/nanoGPT-with-TPU如果觉得有用求一个 ⭐这篇博客侧重什么我不会花大篇幅讲 Transformer、多头注意力这些——网上已经有很多优秀的资料了我强推 Karpathy 的视频链接在文末。我想聊的是如何用 TPU 去训 GPT。目前大部分 GPT 开源项目都用 GPU而 TPU 那边的代码又以 TF/JAX 为主。PyTorch TPU 这个组合训 GPT 的项目我找了一圈确实没见着。所以我想把自己的折腾经历写下来给那些跟我一样——习惯 PyTorch、又想白嫖 TPU 算力的人省点踩坑的时间。准备工作注册 Kaggle 账号已有可跳过动手之前你需要一个 Kaggle 账号。注册流程打开 kaggle.com点击 “Register”可以用 Google / GitHub / 邮箱注册推荐用 Google 账号一步到位注册完后去 Settings → Account → Phone Verification 绑定手机号现在 Kaggle 要求绑定手机号才能用加速卡GPU / TPU国内手机号大概率收不到验证码需要国外手机号验证通过后手机号就绑定了以后不用再管它小贴士如果你已经有 Kaggle 账号但没绑手机登录后去 Settings 看看有没有红色警告就知道了。没有警告直接跳到下一步。开启 TPU RuntimeKaggle Notebook 默认用的是 CPU需要手动切换到 TPU打开 Kaggle Notebook 编辑器点击右上角的 “Settings” 面板在 “Accelerator” 下拉菜单中选择“TPU v5e-8”或者右边栏点开 Session options 也可以选择 TPU把 “Internet” 开关打开下载数据集需要联网搞定现在你拥有了 8 个 TPU 核心算力大概相当于……嗯大概相当于你白嫖了 8 张显卡先偷着乐吧。运行跑起来很简单先运行fineweb.py下载处理语料建议在 Kaggle 新建 notebook 运行因为语料很大再导入 notebookbuild-nanogpt.ipynb以及刚刚搞定的语料可以做成 Kaggle 的 Dataset 导入记得改文件路径运行整个 notebook 就可以了。下面讲一下用 TPU 去训 GPT 的踩坑经历。TPU 环境补丁必须最先执行import os os.environ[PJRT_DEVICE] TPU os.environ[OMP_NUM_THREADS] 1 if TPU_PROCESS_ADDRESSES in os.environ: del os.environ[TPU_PROCESS_ADDRESSES] if CLOUD_TPU_TASK_ID in os.environ: del os.environ[CLOUD_TPU_TASK_ID] # TPU memory optimizations os.environ[XLA_TENSOR_ALLOCATOR_MAXSIZE] 100000000 print(✅ TPU 环境补丁已设置)这个环境补丁很重要如果你用 PyTorch并且是用 Kaggle 的 TPU 来训练不管你是用 XLA 还是 Accelerate 库建议加上这个补丁否则会报错。这也算是本项目的一个小亮点吧解决了 Kaggle 上 TPU 训练恼人的环境问题可以将这段代码 copy 到其他地方用。踩坑修复一参数共享问题GPT 需要嵌入层 wte 跟输出分类头 lm_head 的参数共享一般各个项目都是用指针指向同一块内存实现的比如 lm_head.weight wte.weight。这在 GPU 上没问题但我用 Accelerate 库在 TPU 上训练就出问题了。一开始我并没有直接定位到这个 bug而是先察觉到一个异常现象每次跑几千步看看效果后面从 checkpoint 接着训练发现 loss 一开始会飙升到 1114然后很快又下降回上次训练的结果比如56。后来我觉得不对劲了因为我每次训练前后都会做 eval按理 loss 应该一样才对怎么会退化成随机初始化的 loss 值呢后面我艰难地排查了好一阵子。最后下载模型权重一通扫描发现lm_head 的参数方差是 0.02跟随机初始化时一模一样而 wte 的方差已经变了——说明 wte 已经训练好了但 lm_head的参数就像没训练过一样。如果参数共享生效这两个参数应该完全一致才对。我的第一反应是只保留一份参数 wte加载后再次判断指针是否一致发现问题依旧没有解决。折腾了好一阵子给代码加各种判断、打补丁都没用。最后我发现问题出在 accelerator.prepare() 会破坏 weight tying。眼看代码越改越长我发现其实可以用函数式 API——F.linear(x, wte.weight) 替代显式声明 lm_head这样就彻底绕开了参数共享的问题。defforward(self,idx,targetsNone):xself.transformer.wte(idx)# ... transformer blocks ...xself.ln_f(x.float()).to(x.dtype)# 直接用 F.linear手动传入 wte.weight 作为转置权重# 这样无论 prepare() 怎么处理参数都能保证是同一份权重logitsF.linear(x,self.transformer.wte.weight)returnlogits,loss这个问题 HuggingFace 的官方文档也有提到https://huggingface.co/docs/accelerate/en/basic_tutorials/tpu原文是这样说的A common language model design is to tie the weights of the embedding and softmax layers. However, moving the model to a TPU (either yourself or passing it to the prepare() method) breaksthe weight tying and you’ll need to retie the weights.但官方给的解决办法仅限于 HuggingFace 上的模型我们是自己实现的 GPT所以我建议用函数式 API, 一劳永逸。踩坑修复二BF16 混合精度训练突发数值异常问题如果不加以优化模型要二十多个小时才能训完 1 epoch19073 steps所以我必须做点什么。为了提升训练效率我开启了 bf16 混合精度。bf16 只占两个字节相比 fp32 存储减半、速度翻倍。但问题来了TPU 上开 bf16 训练到六七千步时loss 和 grad norm 突然变成 NaN。对此我排查了很久。我注意到了一个细节代码中 grad norm 只要大于 1 就会触发梯度裁剪而日志显示 grad norm 一直在 1.8 左右徘徊——说明每一步都在裁剪训练非常不稳定。我试过把 softmax、交叉熵、LayerNorm 等数值敏感的地方显式转成 fp32 计算但问题依旧。这个问题困扰我很久因为我看过其他用 GPU 训 GPT 的项目开 bf16 混合精度都没问题为什么偏偏我用 TPU 就不行这个时期项目陷入了低谷Kaggle 每周的免费 TPU 时长也烧完了。我就先做点别的逛逛 GitHub无意中发现了 modded-nanogpt它的 README 提到了很多 GPT 训练优化的技巧其中提到了QK-Norm。我尝试加上 QK-Norm并配合在 softmax、交叉熵、LayerNorm 等关键路径转 fp32 计算后来看到 PyTorch 关于torch.amp的文档也提到混合精度下这些地方要注意。另外我在 attention 计算全程也转成了 fp32——算是防御性编程吧。后面我试着去掉 attention 的 fp32 转换发现也没问题速度快不了多少但 loss 高了 0.1 左右所以最终代码保留了对 attention 全部转到 fp32 计算。修改之后grad norm 稳定在 0.3 左右相比之前的 1.8 好太多了。QK-Norm 是什么其实很简单做 attention 计算时不是有q k / sqrt(d)吗在这步之前先让 Q 和 K 分别过一层 LayerNorm相当于多了一步归一化。一开始我也不明白 QK-Norm 为什么有效后面查论文发现了一篇关键文献——《Scaling Vision Transformers to 22 Billion Parameters》[2302.05442]。这篇论文讲的是将 ViT 扩展到约80 亿参数时训练在几千步后就会发散divergence。根本原因是注意力 logits 中出现极大的值导致注意力权重变成近似 one-hot分布注意力熵趋近于零attention entropy collapse这种现象在深层、大参数的 Transformer 中尤为严重因为 Q 和 K 的范数在训练过程中会无控制地增长。论文发现在注意力点积前对 Q 和 K 做 LayerNorm能约束注意力 logits 的数值范围防止注意力熵崩溃和训练发散。我读完之后发现我遇到的问题跟论文里描述的非常像——虽然它是视觉大模型参数量比我的 GPT 大得多但 QK-Norm 确实奏效。而且自这篇论文之后QK-Norm 已被广泛采用成了大模型训练的常见技巧之一。踩坑修复三OOM因为看了 Karpathy 的视频我一开始就知道要用梯度累积gradient accumulation。简单来说梯度累积就是把一个大 batch 拆成多个小 batch 分步算每算完一个小 batch梯度先攒着不更新参数等攒够了累加值再统一更新一次参数。这样显存只装得下小 batch但效果等同于大 batch。遵循 GPT-2 论文的参数设置全局 batch 为 524288 tokens/step序列长度 1024换算成序列数就是 524288 ÷ 1024 512 条。因为是 8 卡并行每张卡每步需要处理 512 ÷ 8 64 条序列。那这 64 条怎么塞进去呢设 micro batch size m每张卡每次前向只算 m 条攒够 64 ÷ m 次再更新一次参数。但问题来了TPU v5e 单卡显存只有 16GBmicro batch size 只能开到 4~8勉强能跑但非常慢。因为 TPU 算力其实很足显存太小导致数据填不满计算单元算力利用率很低memory-bound。所以我引入了**梯度检查点gradient checkpointing**来降低显存峰值把 micro batch size 提到了 32梯度累积只要 2 次就能凑满。梯度检查点的核心思想前向传播时不保存中间激活反向传播时重新计算。不用梯度检查点 前向保存 layer0 激活 → 保存 layer1 激活 → ... → 保存 layer11 激活 反向用 layer11 激活算梯度 → 用 layer10 激活算梯度 → ... 显存12 层激活同时存在 用梯度检查点 前向计算 layer0 → 丢弃激活 → 计算 layer1 → 丢弃激活 → ... → 只保存最终输出 反向重新计算 layer11 激活 → 算梯度 → 释放 → 重新计算 layer10 激活 → ... 显存同时只有 1 层激活存在代价是前向传播要多算一次反向时重新计算计算量约增加 33%但显存从 O(N) 降到了 O(1)N 是层数典型的用时间换空间。听起来这是牺牲速度换显存对吧如果只看单步计算量确实慢了。但实际效果恰恰相反——速度反而变快了。为什么呢因为梯度检查点省下的显存让我可以把 micro batch size 从 4~8 提升到 32梯度累积从 16 次降到了 2 次。虽然每步多了检查点的开销但整体计算效率大幅提升——这才是关键TPU 不再饿着肚子干活了算力被充分填满速度自然就上来了。经过上述的一系列优化我的 GPT 最终吞吐量超过 40 万 tokens/s。对极致性能有追求的读者也可以尝试看看能不能进一步压榨算力。不足之处我后面参考modded-nanogpt仓库 README 里介绍的一些技巧在原有基础上改了模型结构比如使用 RoPE、ReLU² 作为激活函数而不是 GELU用残差投影零初始化导致最终跟 GPT-2 论文不完全一致。而且我了解到一些更现代的做法比如用 RMSNorm 替换 LayerNorm、用 SwiGLU 代替两层 MLP。说实话我不能保证这些改动对这个项目一定有提升做的探索实验也不多有兴趣的读者可以自己尝试魔改看看。代码的采样推理做得比较粗糙只实现了 top-k其实可以增加 temperature、top-p 等参数。另外生成效果算不上很惊艳这是我采样得到的结果[prompt] Im a computer science student, sample 0: Im a computer science student, working in the areas of computer technology and computer science, so Ive come up with a lot of these concepts before a while sample 1: Im a computer science student, and the only people in my class who dont want to be in college are students. 你不要乱说啊 This is because my teacher has sample 2: Im a computer science student, who has studied computer science for 25 years. It is pretty easy to find the answer - the question, How do I 好家伙学了25年是吧 sample 3: Im a computer science student, and you can watch me through the program.我这里采样长度比较短你可以调大一些。如果你尝试用不同的推理参数调一下还是有机会生成不错的结果的比如这个[prompt] The future of AI is The future of AI is still in its infancy. The world of AI is a vast, complex and expensive field that requires a tremendous amount of understanding, technology, and expertise.作为对比这是 GPT-2 官方预训练权重生成的文本Im a computer science student, so I dont have much to say about the data, but Ill tell you what: The data I can tell you about is very limited. The data that I can tell you about is very much limited. I dont know how much its going to change over the next few years. I think I can say that in the next few years, I think Ill be able to tell you whats going to be the difference between the future and the...训练你自己的 GPT看到这里不知你是否已经跃跃欲试了。你可以尝试用更多的语料训练更长时间——比如我用的是 FineWeb 语料库的 edu 子集你完全可以尝试用整个语料来训练注意预估训练时间。当然如果你想快速训练或者想训练中文的这里有一个项目供你参考我用几万首中文歌词训练了一个小一点的 GPT 模型大概训练半小时我觉得效果还是可以的我爱你/但我的心不会停/为何爱何苦哭/如风吹醒在这窗里/为何爱何苦心痛/还教我心痛恨你/偏偏爱何苦流泪/偏偏爱何苦心碎/为何泪水难分徘徊/为何爱何苦心痛/难道爱情已经伤人/难道爱何苦心碎/为何恨何苦 风吹过的夏天/所有回忆还是想念/你的泪水不能变得模糊/谁留下所有美丽/所有眼泪叫人悲伤/为你慢慢燃烧/回忆留在黑夜/幸福的梦如此强烈/谁留下所有温柔/你的忧郁里/那黎明有遥远的人/有多少思念的人/那黎明有遥远可以看到很有情歌那味了这种反复中又有变化的感觉让我很惊艳。要注意先对文本进行清洗不然就是rubbish in, rubbish out。另外用已有的 tokenizer 还是自己用 BPE 训练一个也是需要考虑的点。像我这个项目就是自己训练了一个专属的分词器。项目地址https://github.com/suxuefenxi/gpt-cn-songs其实我有个想法几万首歌词还是偏少能不能用现代诗扩充一下语料我感觉两者在文字风格上差不多。不过时间原因没去做感兴趣的读者可以尝试一下。如果你在 Kaggle 上使用 TPU 发现要排队这是正常的。根据我的经验上午训练可以不用等下午跟晚上可能要排队。如果你发现前面一直很多人在等——你我不禁会想是不是这篇博客被很多人看到了。总结Kaggle 免费提供的 TPU v5e-8 是一块被严重低估的算力——不是因为它不强而是因为 PyTorch TPU 的资料太少让很多人望而却步。大部分 TPU 教程都基于 TF/JAX而社区习惯用 PyTorch这中间存在一个鸿沟。这个项目就是想把这个缺口补上让你用最熟悉的 PyTorch白嫖最香的 TPU 算力训练自己的 GPT。另外我还发现现有的大模型对 PyTorch TPU 这套组合的理解也不太靠谱。举个例子你让它写一个用 PyTorch 在 TPU 上训练 CIFAR-10 的代码跑起来全是坑——Kaggle 环境没配好、用的还是旧 API。说白了大模型在这块喂的资料太少了。所以我把这个项目的实战经验整理成了一份SKILL.md放到了 GitHub 上你可以直接用它来增强大模型的能力也可以自己根据我的项目代码再蒸馏一份。这样以后让大模型帮你写 Kaggle TPU 项目就不会翻车了白嫖也白嫖得更得心应手。最后我想说大模型训练不应该只是少数人的玩具。一张免费的 Kaggle TPU 一份开源代码你也可以拥有自己的 GPT。希望这篇博客能帮你迈出这一步。如果觉得有用别忘了去 GitHub 点个 ⭐参考资料Andrej Karpathy - Let’s reproduce GPT-2 (124M)— B 站直接搜就有。我就是看完这个视频才决定自己动手实现GPT的讲解极其细致强烈推荐。视频配套代码https://github.com/karpathy/build-nanogpt更工程化的版本https://github.com/karpathy/nanoGPTmodded-nanogpt(KellerJordan) — README 很值得读介绍了很多性能加速的 tricks。https://github.com/KellerJordan/modded-nanogptHugging Face Accelerate 官方文档— 本项目的分布式训练核心库。PyTorch 也可以直接通过torch_xla使用 TPU但 Accelerate 封装更友好对初学者门槛更低。https://huggingface.co/docs/accelerate/index