在家用显卡上从零训练GPT-2，损失竟低于原版

一位开发者用JAX框架，在自己的RTX 3090上从零训练了一个GPT-2 small模型，耗时37小时15分钟，比等效PyTorch模型快了约3小时。更关键的是，这个JAX模型在测试集上的损失（3.418784）低于原始GPT-2 small的3.499677，也低于他自己PyTorch版本的3.538161。PyTorch用了混合精度（16位），JAX全程用全32位精度。

项目从2024年12月22日开始读《从零构建大语言模型》，到2026年7月5日完成JAX实现，超过一年半。训练数据分两阶段：短训练约9200万token，长训练约32.6亿token。

构建过程被拆解成一系列可逐步验证的步骤。从最简单的A-to-A（输入等于输出）开始，每添加一个标准组件，模型效果就稳定提升，训练损失持续下降。LayerNorm把损失从5.909降到5.734。单头注意力在没有残差连接时让模型变差（损失升至5.934），加上残差连接后降至5.553。位置嵌入使损失降到5.354。多头注意力相比单头改进很小（5.336 vs 5.354）。添加前馈网络后降到5.295。

直接堆叠12层而不加层归一化，训练崩溃：损失飙升到10.75，接近随机模型。在Transformer块内正确放置层归一化后，12层模型训练成功，最终长训练损失3.637。随后添加dropout，损失从3.399升至3.524，说明单epoch场景下dropout未必有益。作者还发现多头注意力的输出投影缺少偏置，修复后损失几乎不变（3.398 vs 3.399）。

尽管JAX模型在损失指标上超越了原始GPT-2，在实际指令微调测试中仍远不如OpenAI的预训练权重。JAX模型得分19.25，OpenAI small权重26.73，OpenAI medium权重41.62。损失下降并不直接等于模型更“聪明”，预训练数据的质量和规模可能比架构细节更重要。

JAX比PyTorch快的原因还不明确。作者猜测可能是全32位精度带来的数值稳定性，也可能是JAX编译优化，或者只是训练初始化运气好。多头注意力输出投影的偏置对损失几乎无影响的原因也不清楚。PyTorch启用torch.compile后是否能在速度上追上JAX，还没有答案。Dropout在更大数据集上是否仍有价值，也值得验证。

这个项目的意义在于：个人开发者或小团队不再需要昂贵的云GPU集群，就能在自己的消费级显卡上复现和实验基础语言模型。现代深度学习中的“小技巧”——如层归一化、残差连接——不是可选项，而是训练深层网络的必需品。构建语言模型像搭积木，每个组件都经过验证，降低了学习门槛。