Tevatron Meets Megatron: Expert-Parallel LLM Reranker Training on an Academic Budget
本文介绍了 Tevatron 3.0，这是一个集成 Megatron-Core 以实现大规模 MoE 重排序器高效专家并行训练的学术预算友好型框架，证明了 30B 参数量的 MoE 模型可以在激活更少参数并实现更高推理吞吐量的同时，达到 8B 稠密模型的质量水平。
原始论文采用 CC BY 4.0 许可（http://creativecommons.org/licenses/by/4.0/）。 这是对下方论文的AI生成解释。它不是由作者撰写或认可的。如需技术准确性，请参阅原始论文。 阅读完整免责声明
想象一下，你正试图在一大堆混乱的干草堆中寻找一根特定的针。在计算机科学领域，这被称为“搜索”。首先，计算机投下一张大网，抓取一千根可能的针（这是“检索”阶段）。但这张网很乱；它抓住了许多稻草和一些错误的针。为了解决这个问题，第二台更聪明的计算机介入了，它仔细观察每一个候选者，并决定哪一个是那根真正的针。这个第二、更细致的步骤被称为“重排序”（reranking）。
长期以来，这些聪明的计算机就像是小型、高效的计算器。但最近，研究人员意识到，如果让他们变得巨大——赋予他们数十亿个“神经元”或“专家”——他们会变得极其擅长寻找正确的答案。然而，这有一个代价：这些巨大的大脑如此沉重，以至于需要像仓库一样规模的超级计算机来训练它们。大多数大学实验室和小型研究小组负担不起一个仓库。他们被困在较小的、较弱的计算机上，这些计算机根本无法承载这些巨大的模型。这篇论文正是针对这个问题的：我们如何在没有超级计算机的情况下，用极低的预算训练这些庞大、超智能的搜索引擎？
这篇论文的作者们是一个来自犹他大学、滑铁卢大学和卡内基梅隆大学等大学的团队，他们构建了一个名为 Tevatron 3.0 的新工具。把以前训练这些模型的方式想象成试图让一个人一次搬运一架巨大的、沉重的钢琴走狭窄的楼梯。这很慢，而且通常钢琴太大了，根本无法通过，所以你不得不放弃。他们以前使用的工具（称为 PyTorch FSDP）就像那段楼梯；他们试图把钢琴拆成碎片，但碎片仍然太重，无法高效地搬运，而且他们无法处理一种特定类型的钢琴（称为“混合专家模型”或 MoE），这种钢琴拥有许多只有在需要时才会运作的不同内部零件。
团队的解决方案是用一辆带有特殊装卸平台的货运卡车来取代楼梯。他们将一个强大的引擎“Megatron”集成到了他们的工具包中。这个新引擎不仅能运送钢琴，还能将钢琴分解成一个个微小、易于管理的箱子，以便同时装载到不同的卡车（计算机）上。这个新引擎最神奇的地方在于它处理“专家并行”（Expert Parallelism）的能力。想象一下一个由 128 位不同专家组成的团队（比如一位厨师、一位机械师和一位诗人）正在处理一个问题。旧有的工具试图让每一位专家都参与处理每一个问题，这是一种浪费时间和空间的行为。新的 Megatron 引擎足够聪明，它会说：“好的，对于这个特定的问题，只需要厨师和机械师站出来，”而其他人则留在后方。这使得团队能够训练一个拥有 300 亿参数（一个巨大的大脑）的庞大模型，而这在以前是不可能实现的预算水平下完成的。
论文表明，这种新方法的效果与旧的、昂贵的方法一样好。事实上，当他们进行测试时，新系统训练一个标准的 80 亿参数模型的速度比旧方法快了约 22%。但真正的魔力发生在那个巨大的 300 亿参数模型身上。旧有的工具根本无法运行它；计算机内存会崩溃。然而，新工具却成功训练了它。
更令人惊讶的是，作者发现这个巨大的 300 亿参数模型——它在回答每个问题时只“唤醒”约 30 亿个参数——其表现与较小的 80 亿参数模型一样出色。这就像是拥有一个藏有 3000 万本书的图书馆，但只需要打开其中的 300 万本就能找到答案，而且得到的结果与阅读完较小的 800 万本书的图书馆所得到的结果一样。不仅质量相同，由于它在处理每个问题时做的工作更少，它的速度也更快。当他们测试每秒可以回答多少个问题时，在使用高速服务器的情况下，这个新的巨型模型比较小的模型快了 1.43 倍。
研究人员还测试了不同的教学方式。他们尝试通过直接展示正确答案（对比学习）以及通过让它模仿一个“老师”模型（蒸馏）来进行教学。他们发现，两者都没有明显的胜出者；这取决于具体的问题类型。他们还展示了你可以使用“低秩”方法（LoRA）来训练模型，这就像是只更新书中的笔记而不是重写整本书，而且它仍然能做得几乎和重写整本书一样好，从而节省了大量的内存。
简而言之，这篇论文并不仅仅是在说“我们做了一个更大的模型”。它证明了你可以在标准的学术预算下，构建这些庞大且高效的搜索引擎。他们搭建了一座桥梁，让小型实验室也能接触到以前仅限于顶级科技巨头拥有的强大技术。他们发布了所有的代码和训练好的模型，以便任何人都可以尝试。结果是一个训练成本更低、使用速度更快、且同样智能的系统，其性能足以媲美最昂贵的替代方案。
您所在领域的论文太多了？
获取与您研究关键词匹配的最新论文每日摘要——附技术摘要，使用您的语言。