Skip to main content

使用 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 操作支持纳入 MUSA Graph 捕获。当前版本中,graph capture 期间的用户缓冲区注册保持禁用状态,因此 MCCL_GRAPH_REGISTER 不作为该场景下的性能调优开关。

支持有多个未完成的 MCCL 操作,这些操作可以是 MUSA Graph 捕获的或非捕获的任何组合。需要注意的是,MCCL 内部用来实现这一点的机制在多个通信器的 MUSA Graph 从同一线程 musaGraphLaunch() 时,已被观察到会导致 MUSA 死锁。要禁用此机制,请参见环境变量 MCCL_GRAPH_MIXING_SUPPORT