ARTICLE DETAIL

资讯详情

深耕郑州网站建设与运营推广的一线实战洞察。

ATB AllToAll 算子源码导读:基于 HCCL/LCCL 双后端的全交换集合通信实现与参数解析

ATB AllToAll 算子源码导读:基于 HCCL/LCCL 双后端的全交换集合通信实现与参数解析 ATB AllToAll 算子源码导读基于 HCCL/LCCL 双后端的全交换集合通信实现与参数解析【免费下载链接】ascend-transformer-boost本项目是CANN提供的是一款高效、可靠的Transformer加速库基于华为Ascend AI处理器提供Transformer定制化场景的高性能融合算子。项目地址: https://gitcode.com/cann/ascend-transformer-boostAllToAll 是 CANN ascend-transformer-boostATB加速库中承担「各卡间数据全交换」的集合通信算子每个通信域内的卡将输入张量按 rank 数量均匀切分向所有卡发送相同数据量的分片同时从所有卡接收相同数据量的数据。本文以路由文档 .agent/knowledge/routing/all_to_all.md 为骨架结合 src/ops/ops_infer/all_to_all/ 目录下的完整源码、参数头文件与测试用例讲解该算子的文件结构、推荐阅读顺序、CreateRunner双后端决策逻辑、AllToAllParam全部字段语义以及 HCCL / LCCL 两条执行路径的实现细节帮助开发者快速定位代码、理解约束并完成二次开发与调试。1. 算子元信息与定位分类infer推理算子复杂度S简单算子单一路径6 个文件Runner 类型OpsRunner、Operation基于 Runner 框架实现无独立 KernelACLNN否纯原生 Ops API 路径Op 目录src/ops/ops_infer/all_to_all/Kernel 目录无独立 Kernel通信能力由 HCCL/LCCL 集合通信库提供参数头文件include/atb/infer_op_params.h在 .agent/knowledge/README.md 的主索引中all_to_all 被归入 communication集合通信算子分类共 11 个此类算子all_to_all 属于纯 Ops 路径无 ACLNN 封装这也是通信类算子的共同特征——集合通信语义无法用单卡 Kernel 表达必须依赖 HCCL/LCCL 底层通信原语。2. 文件清单与角色划分路由文档给出该算子共 6 个文件全部位于src/ops/ops_infer/all_to_all/目录下#文件角色1all_to_all_hccl_runner.cpp源码HCCL 后端执行实现2all_to_all_hccl_runner.h头文件HcclRunner 派生类声明3all_to_all_lccl_runner.cpp源码LCCL 后端执行实现4all_to_all_lccl_runner.h头文件LcclRunner 派生类声明5all_to_all_operation.cppOperation 定义参数校验、InferShape、CreateRunner6all_to_all_operation.hOperation 定义AllToAllOperation 类声明结构上呈现清晰的「1 个 Operation 2 个 Runner 后端」模式Operation 负责统一的参数校验与形状推导实际通信执行则根据backend字段分派给 HCCL 或 LCCL 两个 Runner。3. 推荐阅读顺序路由文档建议按如下顺序阅读源码由浅入深顺序文件重点关注1all_to_all_operation.h输入输出数量、InferShape 签名、参数成员2all_to_all_operation.cppCreateRunner()决策逻辑、参数合法性校验3all_to_all_hccl_runner.cppHcclAlltoAll调用方式4all_to_all_hccl_runner.h三种构造方式rank 元组 / rankTable / 外部 comm5all_to_all_lccl_runner.cpplccl_-All2All的两种调用形态含 transpose6all_to_all_lccl_runner.h通信域初始化参数传递4. Operation 层参数校验与形状推导4.1 类结构all_to_all_operation.h 中定义的AllToAllOperation继承自OperationBase持有唯一的infer::AllToAllParam param_成员并实现五个核心接口GetInputNum()/GetOutputNum()输入输出各 1 个张量InferShapeCheckImpl()执行前的形状与尺寸约束检查InferShapeImpl()推导输出张量形状SetupCheckImpl()运行时张量描述一致性校验CreateRunner()按 backend 决策创建执行器4.2 创建入口与 backend 合法性检查all_to_all_operation.cpp 中的CreateOperationinfer::AllToAllParam模板特化是算子的统一入口依次完成backend 枚举校验backend必须是hccl或lccl否则报ERROR_INVALID_PARAMSoC 兼容性校验通过aclrtGetSocName()获取芯片型号并约束后端适用范围详见第 6 节分布式初始化检查调用OperationUtil::DistributedInitCheckinfer::AllToAllParam(opParam)完成通信域初始化一致性检查rankSize 校验LCCL 后端要求rankSize为偶数rankSize % 2 ! 0直接报错这与 LCCL 全交换算法要求成对通信有关。构造AllToAllOperation时还会根据param.transpose是否为真拼接 OpKeyAllToAllOperation或AllToAllOperationTranspose用于从AtbOperationIrCfg配置中心拉取对应的算子 IR 配置。4.3 InferShape 规则InferShapeImpl的核心逻辑all_to_all_operation.cppoutTensorDescs.at(0) inTensorDescs.at(0); if (param_.backend lccl param_.transpose) { outTensorDescs.at(0).shape.dims[0] inTensorDescs.at(0).shape.dims[0] * param_.rankSize; outTensorDescs.at(0).shape.dims[1] inTensorDescs.at(0).shape.dims[1] / param_.rankSize; }即非 transpose 模式下输出形状与输入完全一致LCCL transpose 模式下输出为[dims[0] * rankSize, dims[1] / rankSize]相当于把每个卡持有的[dims[0], dims[1]/rankSize]分片聚合后再沿第一维拼接为完整的[dims[0]*rankSize, dims[1]/rankSize]。InferShapeCheckImpl在 transpose 模式下还强制以下约束all_to_all_operation.cpp输入张量必须为2 维dimNum 2dims[1]必须能被rankSize整除单卡分片宽度wSize / rankSize不能超过90KB源码常量MAX_W_SIZE 90LL * 1024LL输入张量总大小不能超过190MB源码常量MAX_TENSOR_SIZE 190LL * 1024LL * 1024LL。SetupCheckImpl则保证运行时实际传入的张量描述与上述推导一致非 transpose 模式下输入输出 desc 必须完全相同transpose 模式下输出dims[0]必须等于输入 dims[0] * rankSize、输出 dims[1] * rankSize必须等于输入dims[1]且输入输出均为 2 维。5. CreateRunner双后端决策逻辑CreateRunnerall_to_all_operation.cpp是理解该算子执行路径分派的关键std::shared_ptrRunner AllToAllOperation::CreateRunner(Context context) const { if (param_.backend hccl) { if (param_.hcclComm nullptr) { return std::make_sharedAllToAllHcclRunner(param_, !param_.rankTableFile.empty()); } else { return std::make_sharedAllToAllHcclRunner(param_, param_.hcclComm); } } else if (param_.backend lccl) { return std::make_sharedAllToAllLcclRunner(param_, context); } return std::shared_ptrRunner(); }HCCL 后端当用户未传入hcclComm时根据rankTableFile是否为空选择通信域初始化方式rank 元组方式或 rankTable 文件方式当用户传入hcclComm时直接复用用户管理的通信域LCCL 后端始终由加速库根据rank / rankSize / commMode / commDomain自行初始化 LCCL 通信域。6. 双 Runner 后端执行细节6.1 HCCL 后端AllToAllHcclRunnerall_to_all_hccl_runner.h 中AllToAllHcclRunner派生自HcclRunner提供三个构造函数对应三种通信域来源(param)默认构造内部走 rank 元组方式(param, bool useRankTableFile)useRankTableFile true时以param.rankTableFile初始化通信域单机/多机 rankTable 场景(param, HcclComm hcclComm)复用外部传入的通信域句柄。ExecuteImplall_to_all_hccl_runner.cpp的执行逻辑非常直接——封装了 HCCL 的HcclAlltoAll原语HcclResult ret HcclAlltoAll( runnerVariantPack.inTensors[0].deviceData, Utils::GetTensorNumel(runnerVariantPack.inTensors[0]) / param_.rankSize, // 每个 rank 发送的元素数 GetHcclDtype(runnerVariantPack.inTensors[0].desc.dtype), runnerVariantPack.outTensors[0].deviceData, Utils::GetTensorNumel(runnerVariantPack.inTensors[0]) / param_.rankSize, // 每个 rank 接收的元素数 GetHcclDtype(runnerVariantPack.outTensors[0].desc.dtype), hcclComm_.get(), GetExecuteStream(runnerVariantPack.context));关键点每个 rank 发送/接收的元素数 张量总元素数 / rankSize即输入被均匀切分为 rankSize 份这正是 AllToAll「相同数据量全交换」语义的体现。执行前会校验hcclComm_与设备侧输入输出地址非空。6.2 LCCL 后端AllToAllLcclRunnerall_to_all_lccl_runner.cpp 的ExecuteImpl根据transpose区分两种调用形态非 transpose调用 5 参数版本的lccl_-All2All数据指针、元素数、dtype、streamtranspose先计算burstlen width / rankSize其中width为输入dims[1]再调用 8 参数版本的All2All额外传入burstlen与width由 LCCL 在通信过程中直接完成数据重排从而让输出呈现「第一维 ×rankSize、第二维 ÷rankSize」的转置效果。此外 LCCL 执行路径对LCAL_ERROR_PARA_CHECK_FAIL错误码有专门提示当返回该错误时说明环境变量LCCL_PARALLEL必须设置为0或false才能保证多通信域并发的正确性。6.3 后端与硬件平台的匹配约束从 all_to_all_operation.cpp 的创建入口可以看出后端选择与硬件平台严格绑定后端支持的平台额外限制hcclAtlas 800I A2/A3 或 Atlas 900 A3 SuperpodIs910B()或 SoC 名含Ascend910_93不支持 transposeparam.transpose为真时报错lccl非 transpose 模式仅支持 Atlas 800T A3 或 Atlas 900 A3 SuperpodSoC 名含Ascend910_93transpose 模式不受此限rankSize必须为偶数infer_op_params.h中AllToAllParam的注释也印证backend仅 Atlas 800 A3 推理产品支持配置为lccl。7. AllToAllParam 参数全解include/atb/infer_op_params.h 定义了AllToAllParam各字段语义与取值约束如下字段类型默认值含义与约束rankint0当前卡所属通信编号rankSizeint0通信的卡数量LCCL 后端要求为偶数rankRootint0主通信编号backendstringhccl通信计算类型仅允许hccl/lcclLCCL 仅 Atlas 800 A3 推理产品支持hcclCommHcclCommnullptrHCCL 通信域指针默认为空时由加速库创建用户可传入自管理的通信域指针复用commModeCommModeCOMM_MULTI_PROCESS通信模式枚举HCCL 多线程场景只支持外部传入通信域方式rankTableFilestring空集群信息配置文件路径适用单机及多机通信场景当前仅支持 hccl 后端单机配置了 rankTable 时以其初始化通信域commDomainstring空通信域名多通信域时使用LCCL 多进程模式需设置为 0-65535 的数字transposeboolfalse通信结果是否对输入转置仅 backend 为lccl时生效rsv[62]uint8_t0预留参数7.1 LCCL 多通信域环境变量配合结合参数注释使用 LCCL 后端进行多通信域并发时需注意环境变量配合多线程模式不支持确定性计算LCCL_DETERMINISTIC需设置为0或false多进程/多线程多通信域并发LCCL_PARALLEL需设置为1或true使用结束后必须恢复为0或false否则会导致基础场景性能下降。7.2 一个典型的最小参数配置参考 test_all_to_all_operation.py 中实际使用的参数构造方式HCCL 后端最小配置为{rank: 0, rankSize: 4, rankRoot: 0, backend: hccl}测试中每个进程rank通过torch.classes.OperationTorch.OperationTorch(AllToAllOperation)创建算子、set_param注入参数后执行execute([inTensor])即对应 ATB 的 Operation 标准调用流程。8. 执行语义与测试验证8.1 数据流语义AllToAll 的语义可以概括为输入切分为 rankSize 份后全交换。以 world_size 4 为例每个 rank 的输入[3, 4, 6]共 72 个元素会被切成 4 份、每份 18 个元素rank i 把自己输入的第 j 份发给 rank j同时接收来自所有 rank 的第 i 份并按 rank 顺序拼接得到与输入同形状的输出。HCCL 后端通过HcclAlltoAll的sendCount recvCount numel / rankSize参数直接实现这一点。8.2 测试用例与 Golden 生成tests/apitest/opstest/python/operations/all_to_all/test_all_to_all_operation.py 提供了完整的单机多卡验证逻辑运行前提导出HCCL_WHITELIST_DISABLE1且通过ATB_HOME_PATH指定 ATB 安装路径需先source set_env.sh使用 lccl 后端时需unset HCCL_MTE_ENABLE并将lcal.o复制到当前目录环境约束world_size 4且ASCEND_RT_VISIBLE_DEVICES可见设备数不得少于 4覆盖 dtypeint8 / int16 / int32 / int64 / float32 / float16 / bfloat16共 7 种Golden 逻辑test_all_to_all_operation.py把每个 rank 的输入拉平后按[i*numel/ws : (i1)*numel/ws]切分rank i 依次取各 rank 输入的第 i 段拼接成 golden 输出再与算子实际输出做torch.allclosertolatol0.001比对。这段逻辑正是 AllToAll 语义的直接翻译可作为理解算子的最小参考实现。此外仓库还提供了 LCCL 与多机场景的测试test_all_to_all_operation_lccl.py、test_all_to_all_operation_multi_server.py以及 tests/apitest/opstest/csv/all_to_all.csv 对应的 opstest 用例清单可用于回归验证两种后端在不同集群规模下的行为。9. 相关算子家族在 communication 分类下all_to_all 还有两个语义相近的兄弟算子可对照阅读all_to_allv可变长度版各卡间交换的数据量不再固定均分支持按 rank 指定不同长度目录为 src/ops/ops_infer/all_to_allv/all_to_allvv2增强版在 all_to_allv 基础上进一步增强目录为 src/ops/ops_infer/all_to_allvv2/。三者对应路由文件分别为 .agent/knowledge/routing/all_to_allv.md 与 .agent/knowledge/routing/all_to_allvv2.md。从文件数量看all_to_all 为 6 个文件双 Runner而 all_to_allv / all_to_allvv2 各为 4 个文件后两者在源码结构上更精简。10. 快速导航完整知识条目.agent/knowledge/ops/communication/all_to_all/index.md知识库主索引含 82 个算子总览与定位模式.agent/knowledge/README.md算子源码目录src/ops/ops_infer/all_to_all/参数定义include/atb/infer_op_params.h测试用例tests/apitest/opstest/python/operations/all_to_all/小结AllToAll 算子是理解 ATB 集合通信算子架构的一个典型样例Operation 层统一承担参数校验、形状推导与 Runner 分派HCCL/LCCL 两个 Runner 分别封装HcclAlltoAll与lccl_-All2All完成实际数据交换并通过transpose、backend、rankTableFile、commDomain等参数控制不同平台与部署形态下的行为。开发者可沿着「路由文件 → Operation → 双 Runner → 参数头文件 → 测试用例」的路径快速掌握该算子的全部实现细节并据此推广到 all_to_allv、all_to_allvv2 等其他通信算子。【免费下载链接】ascend-transformer-boost本项目是CANN提供的是一款高效、可靠的Transformer加速库基于华为Ascend AI处理器提供Transformer定制化场景的高性能融合算子。项目地址: https://gitcode.com/cann/ascend-transformer-boost创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考
返回列表