New Scaling Algorithm and Initialization with NVIDIA Collective Communications Library 2.23

作者: Sylvain Jeaugey

背景介绍

NVIDIA集合通信库（NCCL）是为NVIDIA GPU和网络优化的多GPU及多节点通信原语实现。它是多GPU深度学习训练的核心软件组件，负责处理各种GPU间通信，无论是通过PCI、NVLink还是网络。NCCL利用先进的拓扑检测、优化的通信图和调优模型，在NVIDIA GPU平台上实现开箱即用的最佳性能。本文将讨论NCCL 2.23版本中发布的新功能和修复。

主要贡献

NVIDIA Magnum IO NCCL库旨在优化GPU间和多节点通信，这对AI和高性能计算（HPC）应用中的高效并行计算至关重要。此版本的价值在于其新功能：
* 针对ReduceScatter和AllGather的全新PAT算法：引入了基于Brucks算法的并行聚合树（PAT）算法，用于AllGather和ReduceScatter，实现了对数级扩展性。
* 加速初始化：提升了初始化性能，包括能够使用带内网络进行引导通信。
* ncclCommInitRankScalable：一个新的初始化API，用于在创建通信器时使用多个ncclUniqueId，以加速大规模场景下的初始化。
* 节点内用户缓冲区注册：利用已注册的用户缓冲区进行节点内操作，提升性能。
* 新的分析器插件API：提供API钩子以测量细粒度的NCCL性能。

方法细节

PAT对数级扩展算法（用于ReduceScatter和AllGather）

• PAT算法原理与优势：PAT算法是Bruck算法的一种变体，其特点是在大规模、小数据量场景下具有对数级的网络步骤数。随着数据量的增加，它会逐步增加网络传输次数，以保持缓冲需求最小化。该算法同时适用于AllGather和ReduceScatter。预计在使用PAT算法后，中小型消息的性能会更好，且这种提升会随着工作负载规模的扩大而愈发明显。
• 算法执行方式与适用性：该算法为每个rank执行一个移位的二项树。与递归加倍等类似算法相比，它的优势在于适用于任意数量的rank，而不仅限于2的幂次方。
• 当前限制与应用场景：目前，PAT算法仅支持每节点一个GPU。每节点单GPU的ReduceScatter和AllGather场景对于大语言模型（LLM）训练非常重要，因为在LLM训练中，流水线并行和张量并行维度通常与数据并行维度正交。张量并行维度通常与节点内的NVLink连接对齐，这意味着其他并行维度在每个节点上只会使用一个GPU。
• 未来展望：NVIDIA表示将发表一篇论文详细描述该算法的细节。

新的ncclCommInitRankScalable API

• 新API的功能与目标：此功能增加了一个新的初始化函数 ncclCommInitRankScalable，它允许在创建通信器期间利用多个唯一ID。这一新增功能避免了初始化过程中的“多对一”通信模式，从而提供了更具扩展性的初始化性能。
• 传统初始化方法的瓶颈：在创建通信器时，NCCL需要获取通信器中所有rank的地址（即引导步骤）。为此，NCCL依赖于一个所有rank都知道的唯一ID。在通信器初始化的引导步骤中，每个rank都与这个已知的唯一ID交换其地址，这会产生一个“多对一”的通信模式，在大规模场景下造成严重的瓶颈。
• ncclCommInitRankScalable的工作原理：通过ncclCommInitRankScalable，用户现在可以自由地提供多个唯一ID用于引导过程。为了获得最大增益，NCCL会将负载分散到多个唯一ID上。如果提供的唯一ID数量与通信器的规模成比例增长，就可以实现恒定的引导时间。
• 最佳实践建议：这个新API需要多个rank来创建唯一ID。为了获得最佳性能，建议将这些唯一ID尽可能均匀地分布在各个rank中。

加速的引导操作

• 初始化代码的整体性能优化：在2.23版本中，我们改进了初始化代码的整体性能。我们减少了一些不必要的引导集合操作，并对引导步骤进行了性能调优。
• 利用高速网络加速带外通信：现在可以使用高速网络（如IB/RoCE）进行带外通信，以加速初始化的两个线性步骤：引导（bootstrap）和allgather。此功能默认禁用，以避免使用配置错误的设备（因为ncclNet设备的使用发生在拓扑检测之前）。可以通过设置 NCCL_OOB_NET_ENABLE=1 来启用它。
• 指定网络接口：此外，用户可以使用 NCCL_OOB_NET_IFNAME 来指定应使用的网络接口。默认情况下，NCCL将使用在该网络上找到的第一个 ncclNet 设备。

节点内用户缓冲区注册

• NCCL用户缓冲区的默认行为与性能权衡：NCCL的设计允许用户无需注册和维护任何持久性缓冲区即可工作。这极大地简化了可用性，但也带来了一些性能上的权衡。如果没有直接访问权限，NCCL在传输数据时必须进行更多的控制流操作和缓冲处理。与显式注册和映射的缓冲区相比，这会消耗更多的GPU资源，并导致移动相同数据量时产生更高的开销。
• 推荐实践与2.23版本的新支持：我们建议NCCL开发者尽可能使用 ncclCommRegister 注册他们的缓冲区，以便NCCL能够利用所有可用的优化。NCCL团队一直在努力为注册的用户缓冲区增加更多的使用场景。2.23版本为NvLink和PCIe P2P传输实现了节点内用户缓冲区（UB）注册支持。
• 节点内UB注册的主要优势：节点内UB注册的主要好处是避免了节点内节点间的额外拷贝。这减轻了内存子系统的压力，提升了NCCL的通信性能，并改善了计算与通信的重叠。除了 ncclReduce 和 ncclReduceScatter（因为它们无法从中受益），所有NCCL集合操作以及基于send/recv的操作都支持此功能。
• 启用节点内UB注册的两种方式：第一种方式是通过 ncclCommRegister 显式注册缓冲区，缓冲区仅在调用相应的NCCL集合操作时才被注册。第二种方式是通过CUDA Graphs捕获NCCL操作，所有用户缓冲区将在图捕获期间自动注册。更多指南和要求，请参考NCCL文档。
• 适用系统：除了通过NVLink和PCIe进行的节点内通信外，该功能还适用于多节点NVLink（MNNVL）系统中每个NVLink域内的通信。

新的分析器插件API

• 设计动机：随着GPU集群规模的增加，性能异常变得越来越难以检测和定位。需要领域特定的监控和诊断工具，以最小的开销为正在运行的作业收集和分析遥测数据。NCCL分析器插件接口就是为了解决这些问题而设计的。该接口的设计也使其易于被PyTorch Kineto等深度学习框架的分析器采用。

• 插件加载与事件激活：新的 NCCL_PROFILER_PLUGIN 环境变量控制分析器插件的加载和初始化，其方式与其他NCCL插件相同。加载后，分析器插件可以通过设置事件激活掩码来启用NCCL事件分析。NCCL在初始化期间将此掩码暴露给分析器。事件激活掩码是一个32位整数，其中每一位代表一个NCCL分析器事件。目前，NCCL支持以下事件：

  • ncclProfileGroup (bit-0): 组事件
  • ncclProfileColl (bit-1): 集合操作事件
  • ncclProfileP2p (bit-2): 点对点事件
  • ncclProfileProxyOp (bit-3): 代理进度通道事件
  • ncclProfileProxyStep (bit-4): 代理进度步骤事件
  • ncclProfileProxyCtrl (bit-5): 代理进度内部状态事件

• 事件的层次结构：NCCL以分层形式表示事件。例如，集合操作可以组合在一起，而代理操作则协助GPU通过可用的网络通信通道进行单个数据块的点对点传输。因此，NCCL向分析器呈现相应的事件时会保留这种关系。NCCL事件层次结构图如下所示：

ncclProfileGroup
+- ncclProfileColl
|  +- ncclProfileProxyOp
|     +- ncclProfileProxyStep
+- ncclProfileP2p
   +- ncclProfileProxyOp
      +- ncclProfileProxyStep

# ncclProfileProxyCtrl

• 层次结构的优势：这种层次化表示使分析器插件能够以更有意义、更易于理解的形式向用户呈现事件。
• 示例插件：NCCL在 ext-profiler/example 目录中提供了一个示例分析器插件，可用作开发第三方分析器插件的模板。

• API回调函数定义：分析器插件接口总共定义了以下五个函数回调：

  • ncclResult_t (*init)( void** context, int* eActivationMask);
  • ncclResult_t (*startEvent)( void* context, void** eHandle, ncclProfilerEventDescr_t* eDescr);
  • ncclResult_t (*stopEvent)( void* eHandle);
  • ncclResult_t (*recordEventState)( void* eHandle, ncclProfilerEventState_t eState, NcclProfilerEventStateArgs_t* eStateArgs);
  • ncclResult_t (*finalize)(void* context);

• init 函数：分析器的 init 函数接收一个事件激活掩码指针，并向NCCL返回一个不透明的上下文对象。该上下文为分析器实例之间提供了隔离，而事件激活掩码则由分析器用来通知NCCL应分析哪些事件；例如，设置 *eActivationMask = ncclProfileColl | ncclProfileProxyOp。

• startEvent 函数：分析器的 startEvent 函数接收一个分析器上下文和一个事件描述符。分析器使用描述符信息来分配一个新的事件对象并对其进行初始化。之后，分析器返回一个不透明的句柄，NCCL可以用它来对事件执行进一步的操作，例如记录状态更新。
• stopEvent 函数：分析器的 stopEvent 函数接收一个事件句柄，并将该事件标记为完成。此后，该事件句柄将不能再被使用（分析器可能会在内部回收相应的对象以用于未来的事件）。
• recordEventState 函数：分析器的 recordEventState 函数接收一个事件句柄、一个事件状态以及（可选的）一个事件状态参数对象。此函数使分析器能够更新在NCCL中可能经历不同状态转换的事件。一个例子是代理事件，其中代理在传输数据时需要与GPU和网络进行协调，从而在此过程中从一个状态转换到另一个状态。
• finalize 函数：分析器的 finalize 函数接收分析器上下文，并释放与其相关的所有资源。

补充细节与修复

• 异步图分配：使 cudaMalloc 和 cudaMemcpy 在图分配期间的调用变为异步，从而显著加快图捕获速度。
• 使用致命的IB异步事件停止网络操作：有助于在NCCL内捕获链路断开错误和其他致命的异步事件。
• 在AMD CPU上设置P2P级别为PXB：当每节点使用超过两个GPU时适用。
• 改进init日志以报告实际的NCCL函数：告知用户NCCL正在执行的是 ncclCommInitRank 还是 ncclCommSplit。
• 添加 NCCL_CONF_FILE 变量。
• 将默认IB超时从18增加到20。
• 为NVIDIA peermem添加新的检查：与最新的Linux内核兼容。
• 修复旧的性能回归问题：该问题在混合大小操作时出现。
• 修复当NUMA ID等于-1时的崩溃问题。
• 修复当 NCCL_CROSS_NIC 设置为1时的树图搜索问题。

实验设置

文章未提供详细的实验环境配置信息，如数据集、模型、硬件（GPU、CPU、网络）或软件（操作系统、代码库）的具体细节。

实验结果

文章主要对新功能进行了定性描述，并未提供定量的实验结果、图表或性能对比数据。文中提到的一些预期性能提升（如PAT算法对中小型消息的性能改善）是基于算法设计原理的推断。

结论

NVIDIA NCCL 2.23版本为优化AI和HPC应用中至关重要的GPU间和多节点通信引入了多项新功能和改进。关键增强功能包括：新的PAT算法、大规模场景下的加速初始化、节点内用户缓冲区注册以及新的分析器插件API。