新方法可能提升 LLM 训练效率
Source: MIT News - AI
推理大型语言模型（LLMs）
推理 LLM 旨在通过将复杂问题拆解为一系列更小的步骤来求解。这些强大的模型在高级编程和多步骤规划等具有挑战性的任务上表现出色。
然而，开发推理模型需要巨大的计算和能源消耗，因为训练过程效率低下。虽然少数高功率处理器持续处理复杂查询，许多其他处理器却处于空闲状态。
来自 MIT 及其他机构的研究人员找到了一种利用这些计算空闲时间 高效加速推理模型训练 的方法。
工作原理
- 更小、更快的模型（draft‑er） – 自动训练以预测大型推理 LLM 的输出。
- 验证 – 大型模型检查 draft‑er 的预测。
- 工作负载降低 – 推理模型的工作量减少，从而加快训练速度。
该小模型以 自适应 方式进行训练和部署，仅在部分处理器空闲时启动。通过利用本来会被浪费的资源，训练速度提升且无需额外开销。
在多个推理 LLM 上测试时，该方法将训练速度提升了一倍，同时保持了准确性。这有望降低成本并提升能源效率，适用于金融趋势预测或电网风险检测等高级 LLM 应用。
“人们希望模型能够处理更复杂的任务。但如果这正是模型开发的目标，那么我们必须把效率放在首位。我们找到了一个无损的解决方案，并开发了一个完整的系统，能够在实际中实现相当显著的加速，” 胡庆浩 说，他是 MIT 的博士后，也是该技术论文的共同第一作者。
胡庆浩在论文中与共同第一作者 杨尚（EECS 研究生）、郭俊贤（EECS 研究生）、资深作者 韩松（EECS 副教授、电子研究实验室成员、NVIDIA 杰出科学家），以及来自 NVIDIA、苏黎世联邦理工学院、MIT‑IBM Watson AI 实验室和马萨诸塞大学阿默斯特分校的合作者共同完成。该研究将在 ACM 国际编程语言与操作系统体系结构支持会议 上进行报告。
训练瓶颈
开发者希望推理型 LLM 能够 识别并纠正其批判性思维过程中的错误，从而处理那些会让普通 LLM 卡住的查询。
为教授这一技能，开发者使用 强化学习（RL）：
- 模型为一个查询生成多个潜在答案。
- 它为最佳候选答案获得奖励。
- 模型根据该最佳答案进行更新。
这些步骤会重复数千次，模型在此过程中学习。
问题所在
- Rollout（生成多个答案）可能会消耗 高达 85 % 的 RL 训练所需执行时间。
- 更新 模型——实际的“训练”部分——相对耗时很少。
“更新模型——即实际的‘训练’部分——相比之下耗时极少，”胡说。
在标准 RL 中，所有处理器必须在继续之前完成各自的响应。如果有些处理器在处理较长的响应，其他已经完成的处理器就会空闲等待。
“我们的目标是把这段空闲时间转化为加速，而不产生任何浪费成本，”胡补充道。
投机解码
研究人员转向 投机解码，其做法包括：
- 训练一个 更小的“草稿模型”（drafter），快速猜测大模型未来的输出。
- 让大模型 验证 这些猜测。
- 使用被接受的猜测进行训练。
由于大模型可以一次性验证大量猜测，这一过程得以加速。
自适应解决方案：“驯服长尾” (TLT)
传统的投机解码使用静态草稿模型，一次训练后保持不变。这在强化学习（RL）中行不通，因为推理模型会被更新数千次；静态草稿模型很快就会变得陈旧。
TLT 组件
-
自适应草稿模型训练器
- 利用空闲的处理器时间即时训练草稿模型，使其与目标模型保持一致，而无需额外的计算资源。
-
自适应 rollout 引擎
- 管理投机解码，自动为每一批新输入选择最优策略。
- 根据工作负载特征（例如草稿模型处理的输入数量与目标模型接受的输入数量）调整配置。
草稿模型被刻意设计为轻量级，以实现快速训练。TLT 还复用推理模型训练流水线的组件，从而获得额外加速。
“一旦某些处理器完成了短查询并变为空闲，我们就立刻让它们使用与 rollout 过程相同的数据进行草稿模型训练。关键机制是我们的自适应投机解码——没有它这些收益是不可能实现的，”胡说。
结果
- 在多个推理 LLM 和真实世界数据集上进行测试。
- 训练速度提升 70 % 至 210 %，同时保持模型准确性。
- 小型草稿模型还可以重新用于高效部署，作为免费的副产品。
展望
TLT 展示了利用闲置计算资源可以在不牺牲性能的情况下显著加速推理模型的训练。这种方法有望降低成本并提升未来大语言模型开发的能源效率。
集成与未来方向
研究人员旨将 TLT 集成到更广泛的训练和推理框架中，并探索能够通过该方法加速的新的强化学习应用。
“随着推理持续成为驱动推理需求的主要工作负载，青浩的 TLT 是应对这些推理模型训练计算瓶颈的出色工作。我认为该方法在高效 AI 计算的背景下将非常有帮助，” — 宋晗
资助来源
- MIT‑IBM Watson AI 实验室
- MIT AI 硬件计划
- MIT 亚马逊科学中心
- 现代汽车公司
- 美国国家科学基金会