从零训练大语言模型的学习心得

从零训练大语言模型的研究虽然有趣，但同时也令人沮丧，因为影响因素众多，很难分辨哪个因素让模型变得稍好或稍差——而这正是整个深度学习领域的探索过程！如果你阅读任何预训练大语言模型的架构代码，大约只有200-300行，但若想真正理解每一行，300页的书都不够。

在任何模型训练中，主要贡献因素包括：

1. 模型架构
2. 数据
3. 优化器（控制学习过程）
4. 学习算法（目前必须是梯度下降法）

学习“训练大语言模型”最快的方法反而是最慢的方式：从零开始构建、反复调试、诊断其成功与失败的原因。你不可能通过阅读和观看视频，第二天就能训练出SOTA级别的大语言模型。但在调试过程中，你可以快速失败并快速学习。

我训练了一个约1.6亿参数的模型，结合了多种架构，如MTP（多令牌预测）、HRM（层次推理模型）以及类似GPT的因果语言模型。它的表现甚至优于nanoGPT，并超越了CORE准确率。

在进行了数百次（具体记不清了，大概300次）实验后，许多都没有成功，而效果最好的那些实验，我则在此基础上继续改进。我阅读了许多架构论文，如文本扩散模型、HRM、MOE等，老实说我都尝试构建了约1.5亿参数的版本，但有些架构（如扩散模型）需要长时间训练且收敛缓慢，效果不佳。另一个目标是快速训练，使用较少的tokens（大约20-30分钟），类似于nanochat的速通模式——虽然有些小模型使用了400-500亿tokens训练，但我只想在20-30亿tokens的规模上快速完成训练。

学习心得

第一点： 我学到了，当我想实验一个想法时，如何将其简化并快速测试。我之前犯的错误是试图复现完整的流程。如果我从像HRM这样的论文中汲取灵感，我会找到其核心机制——即“循环回路”，它能提供潜在推理能力，并问自己：“如何提取这个核心想法并将其融入我的架构进行测试，而不是对架构做大的改动？”我理解了HRM为何有效，并受其启发在自己的架构中构建了一个简单的循环并加以测试。结果成功了，它为我带来了一些额外的CORE分数，而没有增加过多参数或进行重大的架构变更。

第二点： 显而易见，像1.6亿参数这样的小型语言模型在狭窄任务上可以表现出色，但在多任务上则力不从心。我尝试在数学数据集上进行监督微调，模型学会了gsm8k这样的数学问题，但其他任务的表现也没有变差。这是因为它们没有足够的神经元来捕捉这些信息。我们知道，大语言模型本质上是一种压缩算法，将所有知识压缩到其潜在空间中，但这个潜在空间必须足够大才能容纳这些信息。这就是为什么缩放定律指出，如果你扩展Transformer模型（增加参数），它最终会吸收所有知识并成为专家。当然，这种压缩也与架构有关：更好的架构能更有效地压缩知识，而较差的则不能。因此，我在架构空间进行了大量探索，并整合了不同的灵感，使模型变得更好。

第三点： 强化学习只有在模型对特定主题拥有超过一般水平的知识时才有效，否则它不会从中受益。我原本打算在监督微调之后对模型进行后续训练，目标是让模型学会gsm8k数据集，但不幸的是，监督微调后的模型数学能力很差，几乎没有学到什么。它在某些数学问题上获得了一些奖励，但无法解决大部分问题。如果一直训练下去，模型最终可能会在更高温度下摸索出答案，但这不是我们想要的。因此，我们需要先对模型进行预训练或更充分的数学数据监督微调，让它具备平均水平的数学知识和推理能力，然后强化学习才能将模型提升为专家（锦上添花）。简言之，强化学习无法让一个对某领域几乎一无所知的人成为专家，但可以让一个超过平均水平的人成为专家。

直觉上，强化学习对错误给予0奖励，对正确给予1奖励，你也可以添加中间步骤（如步骤和格式），但通常这种奖励非常稀疏。因此，如果模型没有内部推理能力和解决问题的知识，它只能依赖试错，碰运气获得奖励，这会导致模型在整个token空间中搜索那些能带来奖励的token组合，因而不可靠。

第四点： 我在开头定义的四个基本要素——架构、数据、优化器和学习算法——并不是模型训练得更好的唯一因素，基础设施也同样重要。无论是早期实验还是生产级训练，都需要良好的基础设施来训练和扩展模型。我在训练时浪费了大量GPU分钟，因为基础设施代码不够优化、效率低下，导致训练时间过长、等待时间、冷启动以及分布式训练中的梯度同步问题。在大型语言模型领域进行系统研究时，系统知识非常重要。

确保充分利用GPU潜力很重要。我在做强化学习实验时，发现训练只用了38%的GPU，因为在强化学习中，你需要收集轨迹，然后对这些轨迹进行反向传播，而轨迹是按每个样本顺序收集的，这对GPU来说并不是最优的基础设施。因此，我必须将所有生成任务堆叠到一个批次中，在这个批次中，所有生成都进行下一个token的预测，但256个样本同时预测下一个token，这能快速生成用于训练的人工数据。

第五点也是最后一点： 记录实验，保持逻辑而非情绪化，根据基准做决策。起初，每当我有想法并好奇想去尝试时，我就直接做不同的实验，但不做记录，只是情绪化地说“好吧，试试这个，看看能不能行，再调大这个参数也许能更好”。我浪费了大量时间这样操作，这并不是我不知道如何进行规范的研究，而是我过于好奇，跳过了遵循正确流程的步骤。但很快我意识到了问题，创建了一个排行榜，开始对实验进行排名和记录。

今天的分享就到这里，感谢阅读！代码在这里：https://github.com/harshbhatt7585/deep-learning-papers-implementation/blob/main/tinyGroot