跳到主要内容

MCCL over ACE

1. 概述

异步通信引擎 Async Copy Engine (ACE) 是 MTGPU 的一个独立硬件单元,用于在不同内存域之间实现高带宽、低延迟的异步数据搬运。

异步通信引擎 ACE

运行原理

在典型的多芯片(Multi-Die)GPU 结构中,每个 GPU 包含多个计算芯片,ACE 的任务是为这些芯片之间及 GPU 之间提供高效的数据通路,主要包括:

  1. Host 内存 ↔ 显存(H2D/D2H)
  2. 本地 GPU 显存 ↔ 远程 GPU 显存
  3. 芯片内部 Die 间数据搬移(D2D - Die to Die)

对称内存 (Symmetric Memory)

MCCL over ACE 使用对称内存机制,允许每个 GPU 直接访问其他 GPU 的内存地址空间。

Rank 0: [Base + 0*stride] ───┐
Rank 1: [Base + 1*stride] ───┼─→ 对称内存窗口
Rank 2: [Base + 2*stride] ───┘

关键函数:

  • mcclMemAlloc(): 分配对称内存
  • mcclCommWindowRegister(): 注册对称内存窗口

图执行 (Graph Execution)

MCCL over ACE 预先构建 MUSA Graph,包含所有通信节点(同步、拷贝、原子操作),运行时动态更新参数后执行。

优势

  • 减少 API 调用开销
  • 硬件级流水线优化
  • 更低的通信延迟

2. 支持的集合通信原语

AllReduce

AllReduce 操作在设备间对数据(例如,求和、最小值、最大值)进行归约,并将结果存储在每个 rank 的接收缓冲区中。

以 sum 为例,在 k 个 rank 之间的求和 AllReduce 操作中,每个 rank 将提供一个包含 N 个值的数组,并在包含 N 个值的数组 out 中接收相同的结果,其中 out[i] = in0[i] + in1[i] + ... + in(k-1)[i]

AllReduce 操作示意图

AllGather

AllGather 操作从 k 个 rank 中收集 N 个值到大小为 k*N 的输出缓冲区,并将该结果分发到所有 rank。输出按 rank 索引排序。因此,AllGather 操作会受到不同 rank 到设备映射的影响。

AllGather 操作示意图

ReduceScatter

ReduceScatter 操作执行与归约相同的操作,但结果在 rank 之间以等大小的块分散,每个 rank 根据其 rank 索引获得一块数据。由于 rank 决定数据布局,ReduceScatter 操作会受到不同 rank 到设备映射的影响。

ReduceScatter 操作示意图

AlltoAll

在 k 个进程组(rank)之间的 AllToAll 操作中,每个进程组会提供一个大小为 k×N 个值的输入缓冲区:其中第 j 个包含 N 个值的数据块(chunk)会发送至目标进程组 j;每个进程组会接收一个大小为 k×N 个值的输出缓冲区:其中第 i 个包含 N 个值的数据块源自源进程组 i。

ACE 版本的 mcclAllToAllv 接口支持非对称输入。

AlltoAll 操作示意图


3. API 参考

初始化与销毁

mcclAceCommInitRank

初始化 MCCL over ACE 通信器。

mcclResult_t mcclAceCommInitRank(
mcclComm_t* comm, // [out] 通信器句柄
int nranks, // [in] 进程总数
mcclUniqueId commId, // [in] 唯一 ID
int rank // [in] 当前进程 rank
);

与 SHADER 版本的区别

mcclAceCommInitRankmcclCommInitRank 的主要区别在于是否启用 ACE。

// SHADER 版本
mcclResult_t mcclCommInitRank(mcclComm_t* newcomm, int nranks,
mcclUniqueId commId, int myrank) {
...
// isAceComm 默认为 false
MCCLCHECK(mcclCommInitRankDev(...));
...
}

// ACE 版本
mcclResult_t mcclAceCommInitRank(mcclComm_t* newcomm, int nranks,
mcclUniqueId commId, int myrank) {
...
// 显式传递 isAceComm = true
MCCLCHECK(mcclCommInitRankDev(..., true));
...
}

示例

mcclUniqueId id;
if (rank == 0) {
mcclGetUniqueId(&id);
}
MPI_Bcast(&id, sizeof(id), MPI_BYTE, 0, MPI_COMM_WORLD);

mcclComm_t comm;
mcclAceCommInitRank(&comm, worldSize, id, worldRank);

mcclCommDestroy

销毁通信器。

mcclResult_t mcclCommDestroy(mcclComm_t comm);

内存管理

mcclMemAlloc

分配通信显存。

mcclResult_t mcclMemAlloc(
void** ptr,
size_t size
);

mcclMemFree

释放通信显存。

mcclResult_t mcclMemFree(void* ptr);

mcclCommWindowRegister

注册为对称内存。

mcclResult_t mcclCommWindowRegister(
mcclComm_t comm, // [in] 通信器
void* localPtr, // [in] 本地内存指针
size_t size, // [in] 内存大小
void** symmetricHandle, // [out] 对称内存句柄
int flags // [in] 标志
);

flags 参数支持

/* Window Registration flags */
#define MCCL_WIN_DEFAULT 0x00 // 普通注册,一般不用
#define MCCL_WIN_COLL_SYMMETRIC 0x01 // 对称内存注册(ACE 集合操作必需)

集合操作

mcclAllReduce

mcclResult_t mcclAllReduce(
const void* sendbuff, // [in] 发送缓冲区
void* recvbuff, // [out] 接收缓冲区
size_t count,
mcclDataType_t datatype,
mcclRedOp_t op,
mcclComm_t comm,
musaStream_t stream
);

两种模式

  1. Out-of-place (sendbuff != recvbuff):

    • sendbuff: 对称内存
    • recvbuff: 普通内存
    • 每个 rank 收集所有进程的数据,pull mode
  2. In-place (sendbuff == recvbuff):

    • recvbuff: 对称内存,总大小 = count * nranks * sizeof(datatype)
    • 每个 rank 的数据已在对应位置

支持的归约操作

操作支持的数据类型
mcclSum (求和)mcclInt32, mcclUint64, mcclFloat16, mcclBfloat16, mcclFloat32, mcclFloat64
mcclMin (最小值)mcclInt32, mcclInt64, mcclUint32, mcclUint64
mcclMax (最大值)mcclInt32, mcclInt64, mcclUint32, mcclUint64

重要说明

  • 浮点类型(mcclFloat16/32/64, mcclBfloat16)仅支持 mcclSum 操作
  • 整数类型支持所有三种操作,但 mcclSum 不支持 mcclInt64

mcclAllGather

mcclResult_t mcclAllGather(
const void* sendbuff, // [in] 发送缓冲区
void* recvbuff, // [out] 接收缓冲区
size_t count,
mcclDataType_t datatype,
mcclComm_t comm,
musaStream_t stream
);

mcclReduceScatter

mcclResult_t mcclReduceScatter(
const void* sendbuff, // [in] 发送缓冲区
void* recvbuff, // [out] 接收缓冲区
size_t count, // [in] 每个进程接收的元素数量
mcclDataType_t datatype,
mcclRedOp_t op, // [in] 归约操作
mcclComm_t comm,
musaStream_t stream
);

内存要求

  • sendbuff: 对称内存,大小 = count * nranks * sizeof(datatype)
  • recvbuff: 普通内存,大小 = count * sizeof(datatype)

mcclAllToAllv

mcclAllToAllv 接口支持非对称,recvcounts[i] 直接支持非均匀。

mcclResult_t mcclAllToAllv(
const void* sendbuff, // [in] 发送缓冲区 (对称内存)
const size_t sendcounts[], // [in] 发送计数数组
const size_t sdispls[], // [in] 发送偏移数组
void* recvbuff, // [out] 接收缓冲区
const size_t recvcounts[], // [in] 接收计数数组
const size_t rdispls[], // [in] 接收偏移数组
mcclDataType_t datatype,
mcclComm_t comm,
musaStream_t stream
);

4. 示例代码

AllReduce 示例

AllReduce 是一个用于将所有进程的数据进行归约操作(如求和)的集合操作,结果存回所有 rank。

关键要点

  • 必须使用 mcclMemAlloc 分配通信显存
  • 必须注册为对称内存:调用 mcclCommWindowRegister
#include <mpi.h>
#include <mccl.h>
#include <musa_runtime.h>
#include <vector>
#include <cassert>
#include <cstdio>

int main(int argc, char* argv[]) {

// ========== 1. 初始化 MPI ==========
CHECK_MPI(MPI_Init(&argc, &argv));

int worldRank, worldSize;
CHECK_MPI(MPI_Comm_rank(MPI_COMM_WORLD, &worldRank));
CHECK_MPI(MPI_Comm_size(MPI_COMM_WORLD, &worldSize));

// ========== 2. 设置 GPU 设备 ==========
int ngpus = 0;
CHECK_MUSA(musaGetDeviceCount(&ngpus));
int dev = worldRank % ngpus; // 轮询分配 GPU
CHECK_MUSA(musaSetDevice(dev));

// 创建 MUSA stream
musaStream_t stream;
CHECK_MUSA(musaStreamCreate(&stream));

// ========== 3. 初始化 MCCL 通信器 ==========
mcclUniqueId id;
if (worldRank == 0) {
CHECK_MCCL(mcclGetUniqueId(&id));
}
CHECK_MPI(MPI_Bcast(&id, sizeof(id), MPI_BYTE, 0, MPI_COMM_WORLD));

mcclComm_t comm;
CHECK_MCCL(mcclAceCommInitRank(&comm, worldSize, id, worldRank));

// ========== 4. 执行集合操作 ==========
// ========== 4.1 分配和初始化内存 ==========
size_t count = 16 * 1024 * 1024; // 16M 个元素
std::vector<uint64_t> data(count, 1); // 初始化为 1

void* buff = nullptr;
void* handle = nullptr;

// 分配对称内存
CHECK_MCCL(mcclMemAlloc(&buff, count * sizeof(uint64_t)));

// 拷贝数据到 GPU
CHECK_MUSA(musaMemcpy(buff, data.data(),
count * sizeof(uint64_t),
musaMemcpyHostToDevice));

// 注册对称内存窗口
CHECK_MCCL(mcclCommWindowRegister(comm, buff,
count * sizeof(uint64_t),
&handle,
MCCL_WIN_COLL_SYMMETRIC));

// ========== 4.2 执行 AllReduce ==========
// 多次调用测试
for (int i = 0; i < 4; i++) {
CHECK_MCCL(mcclAllReduce(buff, buff, // in-place!
count,
mcclUint64,
mcclSum, // 求和操作
comm,
stream));
}
CHECK_MUSA(musaStreamSynchronize(stream));

// ========== 4.3 验证结果 ==========
uint64_t* host_data = (uint64_t*)malloc(count * sizeof(uint64_t));
CHECK_MUSA(musaMemcpy(host_data, buff,
count * sizeof(uint64_t),
musaMemcpyDeviceToHost));

// 每次 AllReduce: 1*worldSize -> worldSize
// 4 次后:worldSize^4
uint64_t expected = 1;
for (int i = 0; i < 4; i++) {
expected *= worldSize;
}

for (size_t i = 0; i < count; i++) {
assert(host_data[i] == expected);
}

free(host_data);
printf("Rank %d: AllReduce PASSED!\n", worldRank);

// ========== 4.4 清理 ==========
CHECK_MCCL(mcclMemFree(buff));

// ========== 5. 清理资源 ==========
CHECK_MUSA(musaStreamDestroy(stream));
CHECK_MCCL(mcclCommDestroy(comm));
CHECK_MPI(MPI_Finalize());

return 0;
}

5. 错误处理

MCCL over ACE 主要返回以下错误码:

错误代码错误消息说明
mcclSuccess (0)"no error"操作成功
mcclInvalidUsage (5)"invalid usage"API 使用方式不正确

获取错误信息:

mcclResult_t result = mcclAllReduce(...);
if (result != mcclSuccess) {
fprintf(stderr, "错误:%s\n", mcclGetErrorString(result));
}

mcclInvalidUsage 错误场景

1. 要求 in-place 操作

错误代码

mcclAllReduce(sendbuff, recvbuff, count, mcclUint64, mcclSum, comm, stream);
// sendbuff != recvbuff

源码位置src/ce_coll.cc:518

if (sendbuff != recvbuff || isAtomicOpSupported(op) == false) {
return mcclInvalidUsage;
}

解决方案

// 必须 in-place
mcclAllReduce(buff, buff, count, mcclUint64, mcclSum, comm, stream);

2. 缓冲区未注册为对称内存

错误代码

void* buff = nullptr;
musaMalloc(&buff, size); // 普通内存
mcclAllReduce(buff, buff, count, mcclUint64, mcclSum, comm, stream);

源码位置src/ce_coll.cc:524-526

MCCLCHECK(mcclRegFindSymmetric(comm, sendbuff, ...));
if (symPtr == nullptr) {
return mcclInvalidUsage; // 未找到对称内存
}

解决方案

void* buff = nullptr;
void* handle = nullptr;

// 分配对称内存
mcclMemAlloc(&buff, size);

// 注册为对称内存
mcclCommWindowRegister(comm, buff, size, &handle, MCCL_WIN_COLL_SYMMETRIC);

// 正常使用
mcclAllReduce(buff, buff, count, mcclUint64, mcclSum, comm, stream);

3. AllReduce: 不支持的归约操作

错误代码

mcclAllReduce(buff, buff, count, mcclFloat32, mcclProd, comm, stream);
// mcclProd 不支持

源码位置src/ce_coll.cc:518 + src/ce_coll.cc:288-364 (isAtomicOpSupported)

支持的操作和数据类型请参考 mcclAllReduce

解决方案

// 使用支持的操作
mcclAllReduce(buff, buff, count, mcclFloat32, mcclSum, comm, stream);

4. Batch 操作:批量数超过限制

错误代码

int groupnum = 5; // 超过限制
mcclBatchAllToAllv(sendbuffs, sendcounts, sdispls, recvbuffs,
recvcounts, rdispls, datatypes, groupnum, comm, stream);

源码位置src/ce_coll.cc:1147-1149, src/ce_coll.cc:1052-1054

if (groupnum > NUM_MAX_CE_BATCH_SIZE) { // NUM_MAX_CE_BATCH_SIZE = 4
return mcclInvalidUsage;
}

解决方案

int groupnum = 4; // 不超过 NUM_MAX_CE_BATCH_SIZE
mcclBatchAllToAllv(sendbuffs, sendcounts, sdispls, recvbuffs,
recvcounts, rdispls, datatypes, groupnum, comm, stream);

相关文档