新方法大幅提升大语言模型训练效率

推理大语言模型通过将复杂问题分解为一系列较小步骤来解决难题。这些强大的模型在高级编程和多步规划等具有挑战性的任务上表现出色。

但是，开发推理模型需要巨大的计算量和能耗，这是由于训练过程中的低效率造成的。当少数高性能处理器持续处理复杂查询时，群组中的其他处理器却处于闲置状态。

来自MIT等机构的研究人员找到了一种利用这些计算空闲时间来有效加速推理模型训练的方法。

他们的新方法能够自动训练一个更小、更快的模型来预测较大推理大语言模型的输出，然后由较大模型进行验证。这减少了推理模型必须完成的工作量，从而加速了训练过程。

该系统的关键在于能够自适应地训练和部署较小的模型，使其仅在某些处理器空闲时才介入工作。通过利用原本会被浪费的计算资源，它在不产生额外开销的情况下加速了训练。

在多个推理大语言模型上进行测试时，该方法在保持准确性的同时将训练速度提高了一倍。这可以降低成本并提高开发高级大语言模型的能源效率，这些模型可应用于金融趋势预测或电网风险检测等领域。

MIT博士后、该技术论文的共同第一作者胡庆豪表示："人们希望模型能够处理更复杂的任务。但如果这是模型开发的目标，那么我们需要优先考虑效率。我们找到了这个问题的无损解决方案，然后开发了一个能够在实践中实现相当显著加速的全栈系统。"

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

训练瓶颈

开发人员希望推理大语言模型能够识别并纠正其批判性思维过程中的错误。这种能力使它们能够处理会困扰标准大语言模型的复杂查询。

为了教授这项技能，开发人员使用一种称为强化学习的技术来训练推理大语言模型。模型为查询生成多个潜在答案，为最佳候选答案获得奖励，并根据最佳答案进行更新。这些步骤重复数千次，模型在此过程中学习。

但研究人员发现，生成多个答案的过程（称为rollout）可能消耗强化学习训练所需执行时间的多达85%。

胡庆豪说："相比之下，更新模型——这才是实际的'训练'部分——消耗的时间很少。"

这个瓶颈出现在标准强化学习算法中，因为训练组中的所有处理器必须完成其响应后才能进入下一步。由于某些处理器可能正在处理很长的响应，而生成较短响应的其他处理器则等待它们完成。

胡庆豪补充说："我们的目标是将这种空闲时间转化为加速，而不产生任何浪费的成本。"

他们试图使用一种称为推测解码的现有技术来加速过程。推测解码涉及训练一个称为草案模型的较小模型来快速猜测较大模型的未来输出。

较大模型验证草案模型的猜测，其接受的响应用于训练。

因为较大模型可以一次验证草案模型的所有猜测，而不是按顺序生成每个输出，所以它加速了过程。

自适应解决方案

但在推测解码中，草案模型通常只训练一次并保持静态。这使得该技术对强化学习不可行，因为推理模型在训练期间会更新数千次。

静态草案模型在几步之后很快就会变得过时和无用。

为了克服这个问题，研究人员创建了一个名为"驯服长尾"（TLT）的灵活系统。

TLT的第一部分是自适应草案训练器，它利用空闲处理器的空闲时间来动态训练草案模型，使其与目标模型保持良好对齐，而不使用额外的计算资源。

第二个组件是自适应rollout引擎，管理推测解码以自动为每批新输入选择最优策略。该机制根据训练工作负载特征改变推测解码配置，如草案模型处理的输入数量和验证期间目标模型接受的输入数量。

此外，研究人员将草案模型设计得轻量化，以便可以快速训练。TLT重用推理模型训练过程的一些组件来训练草案模型，从而获得额外的加速收益。

胡庆豪说："一旦某些处理器完成了短查询并变得空闲，我们立即切换它们使用相同的数据进行草案模型训练，这些数据也用于rollout过程。关键机制是我们的自适应推测解码——没有它这些收益是不可能的。"

他们在使用真实世界数据集训练的多个推理大语言模型上测试了TLT。该系统在保持每个模型准确性的同时，将训练速度提高了70%到210%。

作为额外的好处，小型草案模型可以作为免费副产品用于高效部署。

未来，研究人员希望将TLT集成到更多类型的训练和推理框架中，并找到可以使用这种方法加速的新强化学习应用。

韩松说："随着推理继续成为驱动推理需求的主要工作负载，庆豪的TLT是应对训练这些推理模型计算瓶颈的出色工作。我认为这种方法在高效AI计算的背景下将非常有用。"

这项工作由MIT-IBM Watson人工智能实验室、MIT AI硬件项目、MIT亚马逊科学中心、现代汽车公司和美国国家科学基金会资助。

Q&A

Q1：TLT方法是什么？它如何提升大语言模型训练效率？

A：TLT（驯服长尾）是MIT研究人员开发的新方法，它通过训练一个较小的草案模型来预测大语言模型的输出，然后由大模型验证。关键是利用空闲处理器的计算时间来自适应训练草案模型，将原本浪费的计算资源转化为训练加速，在不增加额外开销的情况下将训练速度提高70%到210%。

Q2：为什么推理大语言模型训练会出现效率瓶颈？

A：瓶颈主要出现在强化学习的rollout过程中，这个过程需要生成多个潜在答案，可能消耗高达85%的执行时间。由于所有处理器必须等待最慢的处理器完成长响应后才能进入下一步，导致处理短响应的处理器长时间闲置，造成计算资源浪费。

Q3：TLT方法与传统推测解码有什么不同？

A：传统推测解码中的草案模型是静态的，只训练一次，在强化学习中会快速变得过时无用。TLT的创新在于自适应特性：它包含自适应草案训练器，能在训练过程中动态更新草案模型；还有自适应rollout引擎，能根据工作负载特征自动选择最优策略配置。