使用 CUDA 图形与 NCCL

从 NCCL 2.9 开始，NCCL操作可以被 CUDA 图形所捕获。

CUDA 图形提供一种以图形而不是单个操作的方式来定义工作流程。

通过单个CPU操作启动多个GPU操作，可以减少开销。关于CUDA图形的更多细节可以在CUDA编程指南

NCCL的集体、点对点和组操作都支持CUDA图形捕获。这个支持需要至少CUDA 11.3 版本。

下面的示例代码展示了如何在CUDA图形中捕获计算内核和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。