大型语言模型（LLMs）通过将复杂问题拆解为一系列小步骤来进行推理，擅长处理高级编程和多步骤规划等复杂任务。然而，推理模型的开发过程计算量巨大且能耗高，主要由于训练过程中的资源利用效率不高。在训练过程中，部分高性能处理器持续处理复杂查询，而其他处理器则处于空闲状态。

麻省理工学院（MIT）及其他机构的研究人员发现了一种利用这些计算空闲时间来加速推理模型训练的新方法。该方法自动训练一个较小且更快的模型来预测大型推理模型的输出，随后由大型模型进行验证，从而减少了推理模型的计算负担，加快了训练速度。

该系统的关键在于其自适应训练和部署小模型的能力，只有当部分处理器空闲时才启动。通过利用本来会被浪费的计算资源，训练速度得以提升且没有额外开销。

在多个推理大型语言模型上的测试表明，该方法在保持模型准确性的同时，将训练速度提升了一倍。这不仅降低了开发成本，也提高了能效，对于金融趋势预测、电网风险检测等应用具有重要意义。

MIT博士后、该技术论文的共同第一作者胡庆浩表示：“人们希望模型能处理更复杂的任务，但如果这是模型开发的目标，就必须优先考虑效率。我们找到了一个无损的解决方案，并开发了一个完整系统，实际应用中能显著提升速度。”

论文共同第一作者还有电气工程与计算机科学（EECS）研究生杨尚和郭俊贤，资深作者包括EECS副教授、电子研究实验室成员及NVIDIA杰出科学家韩松，以及来自NVIDIA、苏黎世联邦理工学院、MIT-IBM Watson AI实验室和马萨诸塞大学阿默斯特分校的其他研究人员。该研究成果将在ACM国际编程语言与操作系统架构支持会议上展示。

训练瓶颈

推理大型语言模型需要具备识别并纠正自身推理错误的能力，以应对复杂查询。训练这类模型通常采用强化学习（RL）技术，模型生成多个候选答案，基于最佳答案获得奖励并更新模型，重复数千次以实现学习。

研究人员发现，生成多个答案的过程（称为rollout）占据了强化学习训练执行时间的85%。相比之下，模型更新所需时间非常短。

这一瓶颈源于标准强化学习算法要求所有处理器完成响应后才能进入下一步。由于部分处理器处理的响应较长，其他处理器生成较短响应后只能等待，导致资源闲置。

胡庆浩说：“我们的目标是将这些空闲时间转化为加速效果，而不产生额外成本。”

他们尝试利用一种称为推测解码（speculative decoding）的技术加速训练。该技术训练一个较小的“草稿”模型快速预测大型模型的未来输出，大型模型随后验证草稿模型的预测，接受的结果用于训练。由于大型模型可以同时验证所有预测，而非逐条生成输出，训练过程得以加速。

自适应解决方案

然而，传统推测解码中草稿模型通常只训练一次并保持静态，这在强化学习中不可行，因为推理模型在训练过程中会更新数千次，静态草稿模型很快就会失效。

为此，研究团队开发了名为“Taming the Long Tail”（简称TLT）的灵活系统。

TLT的第一部分是自适应草稿模型训练器，利用空闲处理器的计算时间动态训练草稿模型，确保其与目标模型保持同步，且不增加额外计算资源消耗。

第二部分是自适应rollout引擎，管理推测解码过程，自动为每批输入选择最优策略。该机制根据训练负载特征（如草稿模型处理的输入数量和目标模型验证通过的输入数量）调整推测解码配置。

此外，研究人员设计了轻量级草稿模型，便于快速训练。TLT还复用了部分推理模型训练流程中的组件，进一步提升加速效果。

胡庆浩介绍：“一旦部分处理器完成短查询变为空闲，我们立即切换它们进行草稿模型训练，使用与rollout相同的数据。关键机制是我们的自适应推测解码，没有它这些加速是不可能实现的。”

在多个使用真实数据集训练的推理大型语言模型上测试TLT，训练速度提升了70%至210%，且模型准确性未受影响。

额外好处是，训练得到的小型草稿模型可作为高效部署的免费副产品。

未来，研究团队计划将TLT集成到更多训练和推理框架中，并探索更多可利用该方法加速的强化学习应用。

韩松评价道：“随着推理成为推动推理模型需求的主要工作负载，胡庆浩的TLT为解决训练计算瓶颈提供了极佳方案。我认为该方法将在高效AI计算领域发挥重要作用。”

该研究由MIT-IBM Watson AI实验室、MIT AI硬件项目、MIT亚马逊科学中心、现代汽车公司及美国国家科学基金会资助。