ReduceScatter【免费下载链接】asc-devkit本项目是CANN 推出的昇腾AI处理器专用的算子程序开发语言原生支持C和C标准规范主要由类库和语言扩展层构成提供多层级API满足多维场景算子开发诉求。项目地址: https://gitcode.com/cann/asc-devkit产品支持情况Ascend 950PR/Ascend 950DT支持Atlas A3 训练系列产品/Atlas A3 推理系列产品支持Atlas A2 训练系列产品/Atlas A2 推理系列产品支持Atlas 200I/500 A2 推理产品不支持Atlas 推理系列产品AI Core不支持Atlas 推理系列产品Vector Core不支持Atlas 训练系列产品不支持功能说明集合通信算子ReduceScatter的任务下发接口返回该任务的标识handleId给用户。ReduceScatter的功能为将所有rank的输入相加或其他归约操作后再把结果按照rank编号均匀分散到各个rank的输出buffer每个进程拿到其他进程1/ranksize份的数据进行归约操作。函数原型template bool commit false __aicore__ inline HcclHandle ReduceScatter(GM_ADDR sendBuf, GM_ADDR recvBuf, uint64_t recvCount, HcclDataType dataType, HcclReduceOp op, uint64_t strideCount, uint8_t repeat 1)参数说明表1模板参数说明参数名输入/输出描述commit输入bool类型。参数取值如下true在调用Prepare接口时Commit同步通知服务端可以执行该通信任务。false在调用Prepare接口时不通知服务端执行该通信任务。表2接口参数说明参数名输入/输出描述sendBuf输入源数据buffer地址。recvBuf输出目的数据buffer地址集合通信结果输出到此buffer中。recvCount输入参与ReduceScatter操作的recvBuf的数据个数sendBuf的数据个数等于recvCount * rank size。dataType输入ReduceScatter操作的数据类型目前支持float、half、int8_t、int16_t、int32_t、bfloat16_t数据类型即支持取值为HCCL_DATA_TYPE_FP32、HCCL_DATA_TYPE_FP16、HCCL_DATA_TYPE_INT8、HCCL_DATA_TYPE_INT16、HCCL_DATA_TYPE_INT32、HCCL_DATA_TYPE_BFP16。HcclDataType数据类型的介绍请参考表1。op输入ReduceScatter的操作类型目前支持sum、max、min操作类型即支持取值为HCCL_REDUCE_SUM、HCCL_REDUCE_MAX、HCCL_REDUCE_MIN。HcclReduceOp数据类型的介绍请参考表2。strideCount输入当将一张卡上sendBuf中的数据scatter到多张卡的recvBuf时需要用strideCount参数表示sendBuf上相邻数据块间的起始地址的偏移量。strideCount0表示从当前卡发送数据给其它卡时相邻数据块保持地址连续。本卡发送数据到卡rank[i]且本卡数据块在sendBuf中的偏移为i*recvCount。非多轮切分场景下推荐用户设置该参数为0。strideCount0表示从当前卡发送数据给其它卡时相邻数据块在sendBuf中起始地址的偏移数据量为strideCount。本卡发送数据到卡rank[i]且本卡数据块在SendBuf中的偏移为i*strideCount。注意上述的偏移数据量为数据个数单位为sizeof(dataType)。repeat输入一次下发的ReduceScatter通信任务个数。repeat取值≥1默认值为1。当repeat1时每个ReduceScatter任务的sendBuf和recvBuf地址由服务端自动算出计算公式如下sendBuf[i] sendBuf recvCount * sizeof(datatype) * i, i∈[0, repeat)recvBuf[i] recvBuf recvCount * sizeof(datatype) * i, i∈[0, repeat)注意当设置repeat1时须与strideCount参数配合使用规划通信数据地址。图1ReduceScatter通信示例以上图为例假设4张卡的场景每份数据被切分为3块TileCnt为3每张卡上的0-0、0-1、0-2数据最终reducescatter到卡rank0的recvBuf上其余的每块1-y、2-y、3-y数据类似最终分别reducescatter到卡rank1、rank2和rank3的recvBuf上。因此对一张卡上的数据需要调用3次ReduceScatter接口完成每份数据的3块切分数据的通信。对于每一份数据本接口中参数recvCount为TileLenstrideCount为TileLen*TileCnt即数据块0-0和1-0间隔的数据个数。由于本例为内存连续场景因此也可以只调用1次ReduceScatter接口并将repeat参数设置为3。返回值说明返回该任务的标识handleIdhandleId大于等于0。调用失败时返回 -1。约束说明调用本接口前确保已调用过InitV2和SetCcTilingV2接口。若HCCL对象的config模板参数未指定下发通信任务的核该接口只能在AIC核或者AIV核两者之一上调用。若HCCL对象的config模板参数中指定了下发通信任务的核则该接口可以在AIC核和AIV核上同时调用接口内部会根据指定的核的类型只在AIC核、AIV核二者之一下发该通信任务。对于Atlas A2 训练系列产品/Atlas A2 推理系列产品一个通信域内所有Prepare接口的总调用次数不能超过63。对于Atlas A3 训练系列产品/Atlas A3 推理系列产品一个通信域内所有Prepare接口和InterHcclGroupSync接口的总调用次数不能超过63。对于Ascend 950PR/Ascend 950DT一个通信域内所有Prepare接口的总调用次数不能超过63。对于Ascend 950PR/Ascend 950DT通信服务端为CCU时单次最大通信数据量不能超过256M。调用示例非多轮切分场景如下图所示4张卡上均有300 * 41200个float16数据每张卡从xGM内存中获取到本卡数据对各卡数据完成reduce sum计算后的结果数据进行scatter处理最终每张卡都得到300个reduce sum后的float16数据。图2非多轮切分场景下4卡ReduceScatter通信extern C __global__ __aicore__ void reduce_scatter_custom( GM_ADDR xGM, GM_ADDR yGM, GM_ADDR workspaceGM, GM_ADDR tilingGM) { auto sendBuf xGM; // xGM为ReduceScatter的输入GM地址 auto recvBuf yGM; // yGM为ReduceScatter的输出GM地址 uint64_t recvCount 300; // 每张卡的通信结果数据个数 uint64_t strideCount 0; // 非切分场景strideCount可设置为0 HcclReduceOp reduceOp HcclReduceOp::HCCL_REDUCE_SUM; REGISTER_TILING_DEFAULT(ReduceScatterCustomTilingData); // ReduceScatterCustomTilingData为对应算子头文件定义的结构体 GET_TILING_DATA_WITH_STRUCT(ReduceScatterCustomTilingData, tilingData, tilingGM); Hccl hccl; GM_ADDR contextGM AscendC::GetHcclContext0(); // AscendC自定义算子kernel中通过此方式获取HCCL context if (AscendC::g_coreType AIV) { // 指定AIV核通信 hccl.InitV2(contextGM, tilingData); auto ret hccl.SetCcTilingV2(offsetof(ReduceScatterCustomTilingData, reduceScatterCcTiling)); if (ret ! HCCL_SUCCESS) { return; } HcclHandle handleId1 hccl.ReduceScattertrue( sendBuf, recvBuf, recvCount, HcclDataType::HCCL_DATA_TYPE_FP16, reduceOp, strideCount); hccl.Wait(handleId1); AscendC::SyncAlltrue(); // 全AIV核同步防止0核执行过快提前调用hccl.Finalize()接口导致其他核Wait卡死 hccl.Finalize(); } }多轮切分场景开启多轮切分等效处理上述非多轮切分示例的通信。如下图所示每张卡的每份300个float16数据被切分为2个首块1个尾块。每个首块的数据量tileLen为128个float16数据尾块的数据量tailLen为44个float16数据。在算子内部实现时需要对切分后的数据分3轮进行ReduceScatter通信任务将等效上述非多轮切分的通信结果。图3各卡数据切分示意图具体实现为第1轮通信每个rank上的0-0\1-0\2-0\3-0数据块进行ReduceScatter处理。第2轮通信每个rank上0-1\1-1\2-1\3-1数据块进行ReduceScatter处理。第3轮通信每个rank上0-2\1-2\2-2\3-2数据块进行ReduceScatter处理。每一轮通信的输入数据中各卡上相邻数据块的起始地址间隔的数据个数为strideCount以第一轮通信结果为例rank0的0-0数据块和1-0数据块或者1-0数据块和2-0数据块两个相邻数据块起始地址间隔的数据量strideCount 2*tileLen1*tailLen300。图4第一轮4卡ReduceScatter示意图extern C __global__ __aicore__ void reduce_scatter_custom( GM_ADDR xGM, GM_ADDR yGM, GM_ADDR workspaceGM, GM_ADDR tilingGM) { constexpr uint32_t tileNum 2U; // 首块数量 constexpr uint64_t tileLen 128U; // 首块数据个数 constexpr uint32_t tailNum 1U; // 尾块数量 constexpr uint64_t tailLen 44U; // 尾块数据个数 auto sendBuf xGM; // xGM为ReduceScatter的输入GM地址 auto recvBuf yGM; // yGM为ReduceScatter的输出GM地址 HcclReduceOp reduceOp HcclReduceOp::HCCL_REDUCE_SUM; uint64_t strideCount tileLen * tileNum tailLen * tailNum; REGISTER_TILING_DEFAULT(ReduceScatterCustomTilingData); // ReduceScatterCustomTilingData为对应算子头文件定义的结构体 GET_TILING_DATA_WITH_STRUCT(ReduceScatterCustomTilingData, tilingData, tilingGM); Hccl hccl; GM_ADDR contextGM AscendC::GetHcclContext0(); // AscendC自定义算子kernel中通过此方式获取HCCL context if (AscendC::g_coreType AIV) { // 指定AIV核通信 hccl.InitV2(contextGM, tilingData); auto ret hccl.SetCcTilingV2(offsetof(ReduceScatterCustomTilingData, reduceScatterCcTiling)); if (ret ! HCCL_SUCCESS) { return; } // 2个首块处理 constexpr uint32_t tileRepeat tileNum; // 除了sendBuf和recvBuf入参不同处理2个首块的其余参数相同。故使用repeat2第2个首块ReduceScatter任务的sendBuf、recvBuf将由API内部自行更新 HcclHandle handleId1 hccl.ReduceScattertrue( sendBuf, recvBuf, tileLen, HcclDataType::HCCL_DATA_TYPE_FP16, reduceOp, strideCount, tileRepeat); // 1个尾块处理 constexpr uint32_t kSizeOfFloat16 2U; sendBuf tileLen * tileNum * kSizeOfFloat16; recvBuf tileLen * tileNum * kSizeOfFloat16; constexpr uint32_t tailRepeat tailNum; HcclHandle handleId2 hccl.ReduceScattertrue( sendBuf, recvBuf, tailLen, HcclDataType::HCCL_DATA_TYPE_FP16, reduceOp, strideCount, tailRepeat); for (uint8_t i 0; i tileRepeat; i) { hccl.Wait(handleId1); } hccl.Wait(handleId2); AscendC::SyncAlltrue(); // 全AIV核同步防止0核执行过快提前调用hccl.Finalize()接口导致其他核Wait卡死 hccl.Finalize(); } }【免费下载链接】asc-devkit本项目是CANN 推出的昇腾AI处理器专用的算子程序开发语言原生支持C和C标准规范主要由类库和语言扩展层构成提供多层级API满足多维场景算子开发诉求。项目地址: https://gitcode.com/cann/asc-devkit创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考