跳到主要内容

使用 MCCL 与 MUSA Graphs

MCCL 操作可以被 MUSA Graph 捕获。

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

MCCL 的集体、点对点和组操作都支持 MUSA Graph 捕获。这种支持需要相应的 MUSA 运行时版本。

操作启动是否被 MUSA Graph 捕获被认为是该操作的集体属性,因此必须在所有参与启动的 rank 上统一(对于集体操作来说,这是通信器中的所有 rank;对于点对点操作来说,是发送方和接收方)。包含被捕获的 MCCL 操作的 MUSA Graph(通过 musaGraphLaunch 等)的启动被认为是对捕获中相同 rank 集合的集体操作,这些 rank 中的每一个都必须使用从该集体捕获派生的 MUSA Graph。

以下示例代码显示了如何在 MUSA Graph 中捕获计算内核和 MCCL 操作:

musaGraph_t graph;
musaStreamBeginCapture(stream);
kernel_A<<< ..., MUSA stream >>>(...);
kernel_B<<< ..., MUSA stream >>>(...);
mcclAllReduce(..., stream);
kernel_C<<< ..., MUSA stream >>>(...);
musaStreamEndCapture(stream, &graph);

musaGraphExec_t instance;
musaGraphInstantiate(&instance, graph, NULL, NULL, 0);
musaGraphLaunch(instance, stream);
musaStreamSynchronize(stream);

当 MCCL 通信被捕获并且使用 CollNet 算法时,MCCL 允许通过用户缓冲区注册进一步提高性能。有关详细信息,请参见环境变量 MCCL_GRAPH_REGISTER

支持有多个未完成的 MCCL 操作,这些操作可以是 MUSA Graph 捕获的或非捕获的任何组合。我们不建议用户从同一线程对多个通信器的 MUSA Graph 调用 musaGraphLaunch() ,否则可能导致 MUSA 死锁。