用户缓冲区注册

用户缓冲区注册是一个功能，它允许 MCCL 直接通过用户缓冲区发送/接收/操作数据，无需额外的内部复制（零复制）。它可以加速集体操作，并大大减少资源使用（例如#通道使用）。MCCL 提供了两种注册用户缓冲区的方式；一种是MUSA 图注册，另一种是本地注册。MCCL 要求对于所有 MCCL 通信函数调用（例如，allreduce、sendrecv 等），如果任何 rank 在通信器中传递注册缓冲区给 MCCL 通信函数，同一通信器中的所有其他 rank 也必须传递其注册缓冲区；否则，混合注册和非注册缓冲区可能导致未定义行为。

MTLink Sharp 缓冲区注册

MCCL 支持 MTLink Sharp（MTLS）的用户缓冲区注册；任何支持 MTLS 算法的 MCCL 集体操作（例如，allreduce）都可以利用这一特性。

要启用基于MUSA 图的 MTLS 缓冲区注册，用户必须遵守几个要求：

• 缓冲区是通过 mcclMemAlloc() 或合格的分配器分配的（见 内存分配器）。
• MCCL 操作是在每个 rank 的 MUSA 图捕获的流上启动的。
• 在每个 rank 的集体操作中，缓冲区的头部地址偏移是相同的。

注册的缓冲区将在 MUSA 图被销毁时取消注册。这里是一个基于 MUSA 图的缓冲区注册示例：

void* sendbuff;
void* recvbuff;
size_t count = 1 << 25;
CHECK(mcclMemAlloc(&sendbuff, count * sizeof(float)));
CHECK(mcclMemAlloc(&recvbuff, count * sizeof(float)));

musaGraph_t graph;
CHECK(musaStreamBeginCapture(stream, musaStreamCaptureModeThreadLocal));
CHECK(mcclAllReduce(sendbuff, recvbuff, 1024, mcclFloat, mcclSum, comm, stream));
// 每个 rank 的 sendbuff 和 recvbuff 头部地址偏移相同
CHECK(mcclAllReduce((void*)((float*)sendbuff + 1024), (void*)((float*)recvbuff + 2048), 1024, mcclFloat, mcclSum, comm, stream));
CHECK(musaStreamEndCapture(stream, &graph));

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

CHECK(mcclMemFree(sendbuff));
CHECK(mcclMemFree(recvbuff));

另一方面，要启用基于本地的 MTLS 缓冲区注册，用户必须遵守以下要求：

• 缓冲区是通过 mcclMemAlloc() 或合格的分配器分配的（见 内存分配器）。

• 在每个 rank 调用集体操作之前，使用 mcclCommRegister() 注册缓冲区。

• 像往常一样调用 MCCL 集体操作，但同样保持缓冲区头部地址的偏移对每个 rank 都是相同的。

注册的缓冲区将在用户显式调用 mcclCommDeregister() 时取消注册。这里是一个基于本地的缓冲区注册示例：

void* sendbuff;
void* recvbuff;
size_t count = 1 << 25;
void* sendRegHandle;
void* recvRegHandle;
CHECK(mcclMemAlloc(&sendbuff, count * sizeof(float)));
CHECK(mcclMemAlloc(&recvbuff, count * sizeof(float)));

CHECK(mcclCommRegister(comm, sendbuff, count * sizeof(float), &sendRegHandle));
CHECK(mcclCommRegister(comm, recvbuff, count * sizeof(float), &recvRegHandle));

CHECK(mcclAllReduce(sendbuff, recvbuff, 1024, mcclFloat, mcclSum, comm, stream));
CHECK(mcclAllReduce((void*)((float*)sendbuff + 1024), (void*)((float*)recvbuff + 2048), 1024, mcclFloat, mcclSum, comm, stream));
CHECK(musaStreamSynchronize(stream));

CHECK(mcclCommDeregister(comm, sendRegHandle));
CHECK(mcclCommDeregister(comm, recvRegHandle));

CHECK(mcclMemFree(sendbuff));
CHECK(mcclMemFree(recvbuff));

对于基于本地的注册，用户可以在程序开始时注册缓冲区一次，并多次重用缓冲区以利用注册的好处。

为了节省内存，也可以分配一大块缓冲区并只注册一次。sendbuff 和 recvbuff 可以通过大块进一步分配，以实现零复制 MCCL 操作，只要 sendbuff 和 recvbuff 满足偏移要求。以下示例显示了一个用例：

void* buffer;
void* handle;
void* sendbuff;
void* recvbuff;
size_t size = 1 << 29;

CHECK(mcclMemAlloc(&buffer, size));
CHECK(mcclCommRegister(comm, buffer, size, &handle));

// 分配缓冲区块给 sendbuff 和 recvbuff
sendbuff = buffer;
recvbuff = (void*)((uint8_t*)buffer + (1 << 20));

CHECK(mcclAllReduce(sendbuff, recvbuff, 1024, mcclFloat, mcclSum, comm, stream));
CHECK(musaStreamSynchronize(stream));

CHECK(mcclCommDeregister(comm, handle));

CHECK(mcclMemFree(sendbuff));

IB Sharp 缓冲区注册

MCCL 支持 IB Sharp 缓冲区注册，任何支持 IB Sharp 算法的 MCCL 集体操作都可以从该特性中受益，例如 allreduce、reducescatter 和 allgather。目前，MCCL 仅支持每个节点包含 1 个 rank 的通信器的 IB Sharp 缓冲区注册，注册可以减少 MCCL SM 使用量至 1。

通过 MUSA 图启用 IB Sharp 缓冲区注册：

• 使用任何 MUSA 分配器（例如，musaMalloc/mcclMemAlloc）分配发送和接收缓冲区

• 使用 MUSA 图启动 MCCL 集体操作

通过本地注册启用 IB Sharp 缓冲区注册：

• 使用任何 MUSA 分配器（例如，musaMalloc/mcclMemAlloc）分配发送和接收缓冲区

• 在通信器中的每个 rank 使用 mcclCommRegister 注册发送和接收缓冲区

• 启动 MCCL 集体操作

通用缓冲区注册

MCCL 支持节点内缓冲区注册，它针对所有点对点节点内通信（例如，Allgather Ring）并带来较小的内存压力，更好的通信和计算重叠性能。通过在开始时使用 mcclCommRegister 注册缓冲区或应用 MUSA 图，可以为 MCCL 集体操作和 sendrecv 启用节点内缓冲区注册。

用户缓冲区可以通过 VMM API（即 muMem*）、任何基于 VMM 的分配器（内存分配器）或 mcclMemAlloc 分配。通过传统 MUSA API（例如，musaMalloc）分配的缓冲区也可以用于注册。然而，由于在执行过程中可能发生挂起，在失败或者 abort 过程中可能出现段错误。因此，不推荐使用传统缓冲区进行注册；目前，默认情况下禁用了传统缓冲区注册，用户可以设置 MCCL_LEGACY_MUSA_REGISTER=1 以启用它。

内存分配器

为了方便，MCCL 提供了 mcclMemAlloc 函数，帮助用户通过 VMM API 分配缓冲区，之后这些缓冲区可以用于 MCCL 注册。该函数是专门为 MCCL 设计的，因此不建议在应用程序中到处使用由 mcclMemAlloc 分配的缓冲区。

对于高级用户，如果您想为 MTLS UB 创建自己的内存分配器，分配器分配的缓冲区需要满足以下要求：

• 在支持的 GPU 上，分配缓冲区时需要使用共享标志 MU_MEM_HANDLE_TYPE_POSIX_FILE_DESCRIPTOR，同时也要使用 MU_MEM_HANDLE_TYPE_FABRIC。
• 缓冲区物理内存大小是 MUMEM 推荐粒度的倍数（即 muMemGetAllocationGranularity(..., MU_MEM_ALLOC_GRANULARITY_RECOMMENDED)）
• 缓冲区虚拟头部地址至少与 MUMEM 推荐粒度对齐，大小是 MUMEM 推荐粒度的倍数。

对于使用 VMM API 的通用缓冲区注册，分配器需要满足与 MTLS UB 分配器相同的要求。