想象一下，你正在尝试解决一个庞大而复杂的拼图。在人工智能的世界里，这个拼图就是“大语言模型”（LLM）——一种超级聪明的计算机大脑，能够撰写故事、回答问题并解决难题。随着这些大脑变得更大、更聪明，它们变得过于庞大，以至于单台计算机无法承载。
为了解决这个问题，科学家们将拼图碎片分散到许多协同工作的计算机（称为“节点”）上。本文旨在探讨如何使这些计算机高效地相互通信，从而避免浪费时间等待。
以下是研究人员发现并构建内容的简明分解：
1. 问题：“传话游戏”瓶颈
当计算机协同工作时，它们必须不断共享信息。
- 设置：想象一个由 32 人（GPU）组成的团队试图解决这个拼图。他们被分成若干组（节点）。在组内，他们可以瞬间互相喊话（就像使用一种名为 NVLink 的超快速内部对讲机）。但是，要与其他组交谈，他们必须使用较慢的长途电话线（节点之间的网络）。
- 问题：研究人员发现，当计算机尝试解决拼图的“解码”部分（即逐字生成内容）时，它们需要彼此发送非常小的消息。
- 类比：想象你正在参加接力赛。如果你必须跑很长一段距离才能将一张小纸条传递给下一位选手，而接收纸条的人接收速度很慢，那么你大部分时间都在等待。论文发现，标准的“电话线”软件（称为 NCCL）在处理不同建筑物（节点）之间这些微小且频繁的消息时表现极差。这就像当你需要瞬间传递一个秘密握手时，却试图通过缓慢的邮政服务发送一张明信片。
2. 比较：两种团队组织方式
研究人员测试了两种主要的团队组织方式：
- 张量并行（TP）：所有人同时处理拼图的同一部分，但他们必须不断与其他所有人核对，以确保意见一致。这对于大块工作非常有效，但会被大量的核对（通信）所拖累。
- 混合并行（HP）：他们将拼图分成大块，并将不同的块分配给不同的人。这减少了对核对的需求，但对于任务的“逐字”生成部分来说，效率不够高。
发现：对于任务的“逐字”部分（这也是大部分时间发生的情况），TP通常更好，但前提是团队能够足够快地相互交谈。标准的交谈方式太慢了，导致团队停滞不前。
3. 解决方案：NVRAR（“快速通道”）
为了解决通信缓慢的问题，研究人员构建了一个名为NVRAR的新工具。
- 工作原理：他们不使用标准的缓慢邮政服务，而是利用一种名为 NVSHMEM 的技术，构建了一条定制的“快速通道”。
- 类比：将旧的方式想象成发送一封信，这封信必须经过盖章、分拣，并由卡车投递。而 NVRAR 就像拥有一架专用无人机，直接从一个人飞向另一个人，瞬间放下便条并取回回复。
- “递归倍增”技巧：他们将通信组织得像一场“传话游戏”，每个人在每一步中都将交谈的人数翻倍。不是每个人都逐一与所有人交谈，而是他们两两配对、合并、再次配对、再次合并。这对于大型群体来说要快得多。
4. 结果：加速团队
当他们将这个新的“快速通道”（NVRAR）接入系统时：
- 更快的交谈：对于这些 AI 任务中使用的小消息，新系统比标准系统快了1.9 到 3.6 倍。
- 更好的拼图解决：当他们使用这个新系统运行庞大的"Llama 3.1 405B"模型（一个巨大的 AI 大脑）时，生成答案所需的时间显著下降。在某些情况下，团队完成的速度比之前快了1.72 倍。
- 现实世界测试：他们在现实世界的流量（模拟数千名用户提问）上进行了测试，发现该系统可以在不减速的情况下处理每秒更多的请求。
总结
这篇论文本质上是在阐明：当一支庞大的计算机团队试图解决 AI 拼图时，最大的延迟并非来自思考，而是来自交谈。不同建筑物之间标准的交谈方式对于所需的那些微小且频繁的消息来说太慢了。作者构建了一个定制的、高速的通信系统（NVRAR），它像一条快速通道，使团队能够瞬间协调，从而更快地解决拼图。
技术摘要：理解与改进多节点大语言模型推理中的通信性能
问题陈述
随着大语言模型（LLM）规模的增大，其内存占用往往超出单块 GPU 的容量，从而需要在多个节点上进行分布式推理。虽然张量并行（TP）和流水线并行（PP）等模型并行策略在单节点推理中已得到充分研究，但其在多节点环境下的性能仍未得到充分探索。多节点推理因节点间延迟远高于节点内连接（如 NVLink）而引入了显著挑战。
作者指出，虽然 TP 通常对内存受限、以解码为主的负载更优，但由于频繁的 All-Reduce 操作，其通信开销巨大。相反，结合 TP 和 PP 的混合并行（HP）虽降低了通信开销，但由于小矩阵乘法的特性，在以解码为主的场景中未能有效减少计算时间。研究识别出的一个关键瓶颈是：标准 NCCL All-Reduce 实现针对 LLM 解码阶段典型的小消息尺寸（128 KB 至 2 MB）表现次优，尤其是在跨节点扩展时。
方法论
作者利用现有最先进的推理引擎与自定义研究原型相结合，进行了系统的性能研究：
实验设置：
- 硬件： 实验在 Perlmutter 系统（NVIDIA A100 GPU，HPE Slingshot-11 互连）和 Vista 系统（NVIDIA GH200 GPU，InfiniBand 互连）上进行。
- 模型： Llama 3.1 70B 和 405B（Instruct），采用 bf16 精度。
- 工作负载： 测试了两种主要场景：“预填充为主”（大提示长度）和“解码为主”（长生成序列）。
- 引擎： 研究使用了 vLLM 和 SGLang，以及由作者开发的自定义开源推理引擎 YALIS，以支持多节点 HPC 环境下的受控实验。YALIS 具备统一的模型定义层、基于 Torch Compile 的执行机制以及基于 AxoNN 的张量并行功能。
分析方法：
- 作者进行了强扩展实验，测量了 GPU 数量从单节点增加到 32 节点配置时的端到端批处理延迟。
- 他们将执行时间分解为计算（矩阵乘法、其他计算）、通信和空闲时间，以识别瓶颈。
- 使用合成 GEMM 基准测试来隔离微批处理（PP）与张量拆分（TP）对矩阵乘法内核的影响。
** proposed 解决方案（NVRAR）：**
- 为解决通信瓶颈，作者开发了 NVRAR，这是一种基于 NVSHMEM 递归倍增的层次化 All-Reduce 算法。
- 设计： NVRAR 分为三个阶段：
- 节点内 Reduce-Scatter： 使用 NCCL 在节点内归约数据。
- 节点间递归倍增： 使用 NVSHMEM 非阻塞 RMA 原语（put_nbi），采用递归倍增拓扑在节点间执行 All-Reduce。
- 节点内 All-Gather： 使用 NCCL 将归约后的数据收集回节点内的所有 GPU。
- 优化措施：
- 分块非阻塞通信： 通过分块处理数据，实现计算与通信的重叠。
- 融合载荷： 将数据与同步标志合并为单个 8 字节载荷，避免在 Slingshot/InfiniBand 上使用昂贵的显式信号原语。
- 序列号同步： 使用唯一序列号进行全局同步，以避免全局屏障（fence/quiet）带来的延迟开销，同时确保正确性。
主要贡献
- 系统性性能研究： 本文详细剖析了多节点环境下 TP 与 HP 的扩展行为，指出 TP 通常对以解码为主的负载更优，但受限于通信开销；而 HP 在这些场景中难以减少计算时间。
- YALIS 引擎： 开发了 YALIS，这是一款面向研究的推理引擎，专为多节点 HPC 环境中的可仪器化和受控实验而设计。
- NVRAR 算法： 提出并实现了 NVRAR，这是一种针对 LLM 解码阶段典型小消息场景优化的自定义 All-Reduce 算法。它利用 NVSHMEM 在 Slingshot 和 InfiniBand 网络上实现了比标准 NCCL 实现更低的延迟。
- 集成与评估： 将 NVRAR 集成到 YALIS 和 vLLM 中，展示了其在大型模型（Llama 3.1 405B）和混合专家（MoE）模型（Qwen3-235B）上的端到端性能提升。
结果
- 扩展行为： TP 和 HP 在多节点设置下均表现出较差的强扩展性。HP 在计算受限（预填充为主）的场景中优于 TP，而 TP 在内存受限（解码为主）的场景中优于 HP。
- 通信瓶颈： 对于小消息（128 KB–2 MB），NCCL All-Reduce 在跨节点扩展时表现不佳，在这些特定场景下，其性能往往不如基于 MPI 的解决方案。
- NVRAR 性能：
- 微基准测试： 在 HPE Slingshot 和 InfiniBand 互连上，对于 128 KB 至 2 MB 的消息尺寸，NVRAR 的延迟比 NCCL 低 1.9 倍至 3.6 倍。
- 端到端推理： 集成到 YALIS 和 vLLM 后，NVRAR 在使用张量并行处理 Llama 3.1 405B 模型的以解码为主的负载时，将端到端批处理延迟降低了高达 1.72 倍。
- MoE 模型： NVRAR 还通过加速并行策略中的 TP 组件，提高了 MoE 模型（如 Qwen3-235B）的吞吐量，与基于 NCCL 的配置相比，吞吐量最高提升了 1.14 倍。
意义
本文指出，随着 LLM 规模不断扩大且推理迁移至多节点集群，通信效率成为决定性能的首要因素。作者证明，标准通信库（NCCL）并未针对 LLM 解码阶段的特定消息尺寸和延迟约束进行优化。通过引入 NVRAR，他们提供了一种针对性解决方案，显著降低了通信开销，使得大型模型的张量并行扩展更加高效。这项工作强调了下一代分布式推理系统需要自定义通信原语，并验证了在特定 HPC 互连环境中，基于 NVSHMEM 的层次化方法优于标准集合通信库的实用性。
每周获取最佳 machine learning 论文。
受到斯坦福、剑桥和法国科学院研究人员的信赖。
请查收邮箱确认订阅。
出了点问题，再试一次？
无垃圾邮件，随时退订。