使用 NCCL 与 CUDA Graphs

从 NCCL 2.9 开始，NCCL 操作可以被 CUDA Graphs 捕获。

CUDA Graphs 提供了一种将工作流程定义为图形而不是单个操作的方法。 它们可以通过单个 CPU 操作启动多个 GPU 操作来减少开销。 有关 CUDA Graphs 的更多详细信息，请参阅 CUDA 编程指南。

NCCL 的集合操作、P2P 操作和组操作都支持 CUDA Graph 捕获。 此支持需要最低 CUDA 版本 11.3。

操作启动是否被图形捕获被认为是该操作的集合属性，因此必须在参与启动的所有 rank 上保持一致（对于集合操作，这是通信器中的所有 rank，对于点对点操作，这是发送者和接收者）。 包含捕获的 NCCL 操作的图形的启动（通过 cudaGraphLaunch 等）被认为是针对捕获中存在的同一组 rank 的集合操作，并且这些 rank 中的每一个都必须使用从该集合捕获派生的图形。

以下示例代码展示了如何在 CUDA Graph 中捕获计算内核和 NCCL 操作

cudaGraph_t graph;
cudaStreamBeginCapture(stream);
kernel_A<<< ..., stream >>>(...);
kernel_B<<< ..., stream >>>(...);
ncclAllreduce(..., stream);
kernel_C<<< ..., stream >>>(...);
cudaStreamEndCapture(stream, &graph);

cudaGraphExec_t instance;
cudaGraphInstantiate(&instance, graph, NULL, NULL, 0);
cudaGraphLaunch(instance, stream);
cudaStreamSynchronize(stream);

从 NCCL 2.11 开始，当 NCCL 通信被捕获并且使用 CollNet 算法时，NCCL 允许通过用户缓冲区注册进一步提高性能。 有关详细信息，请参阅环境变量 NCCL_GRAPH_REGISTER。

支持具有多个未完成的 NCCL 操作，这些操作可以是图形捕获或非捕获的任意组合。 需要注意的是，已经观察到 NCCL 内部用于完成此操作的机制在从同一线程 cudaGraphLaunch() 多个通信器的图形时会导致 CUDA 死锁。 要禁用此机制，请参阅环境变量 NCCL_GRAPH_MIXING_SUPPORT。