Skip to main content

用户缓冲区注册

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

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));

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

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

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 以启用它。

2.4 对称窗口辅助接口

MCCL 在对称窗口能力上提供了两个主机侧辅助接口:

它们都依赖对称窗口注册与通信器能力探测。

mcclWinGetUserPtr

该接口返回窗口最初绑定的用户指针,常用于:

  • 验证窗口与业务缓冲区的关联关系
  • 在调试或 fallback 路径中恢复应用侧视角的地址

注意:

  • 在支持对称窗口辅助能力的平台上,win 必须有效
  • 在不支持该能力的平台上,win == NULL 的特定探测路径可能返回成功并把输出置为 NULL

参考示例:examples/05_symmetric_memory/03_win_get_user_ptr

mcclGetLsaDevicePointer

该接口用于从主机侧获取对称窗口在指定 lsaRank 视角下的设备地址,适合:

  • 主机侧调试 LSA 可见地址
  • 构造需要显式设备地址的辅助逻辑

参考示例:examples/05_symmetric_memory/02_get_lsa_device_pointer

内存分配器

为了方便,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 分配器相同的要求。