HCCL 集成 PyTorch 分布式后端：Ascend NPU 上 AllReduce 实战指南

HCCL 集成 PyTorch 分布式后端Ascend NPU 上 AllReduce 实战指南【免费下载链接】hccl集合通信库Huawei Collective Communication Library简称HCCL是基于昇腾AI处理器的高性能集合通信库为计算集群提供高性能、高可靠的通信方案项目地址: https://gitcode.com/cann/hccl本文基于 HCCL 仓库中的 PyTorch AllReduce 示例examples/03_ai_framework/01_pytorch讲解如何在单机多卡昇腾 NPU 环境下通过torch.distributed指定hccl后端完成集合通信。读完本文你可以完整理解该示例从设备探测、多进程拉起、通信域初始化到dist.all_reduce()下发 AllReduce 算子的全流程并掌握HCCL_OP_EXPANSION_MODE等关键配置项的使用方法从而将 HCCL 作为标准分布式后端接入自己的 PyTorch 训练代码。一、背景HCCL 如何接入 PyTorch根据 主流框架集成文档HCCLHuawei Collective Communication Library是面向昇腾 AI 处理器的高性能集合通信库。AI 框架主要有单算子模式、图模式Ascend IR和图捕获模式aclgraph三种编程执行形态HCCL 均提供对应工作方式。针对 PyTorch 与 MindSporeHCCL 的调用已集成到 TorchNPU 插件代码中开发者只需在框架 API 中指定使用 HCCL 作为分布式后端直接使用框架原生通信 API如torch.distributed即可实现分布式能力无需手写 HCCL C 接口调用。这一点在示例代码 hccl_pytorch_allreduce_test.py 中体现得非常直接——全部通信逻辑只有三行框架代码而通信域初始化、算子编排、结果校验等底层工作均由torch_npu插件桥接到 HCCL 完成。二、样例能力说明本样例展示如何使用 PyTorch 接口执行 AllReduce 操作覆盖以下四个功能点步骤使用的接口说明设备检测torch_npu.npu.device_count()查询当前可用的 NPU 数量作为进程数与 world_size多进程拉起torch.multiprocessing.spawn()按卡数启动 N 个进程每个进程负责一张 NPU通信域初始化torch.distributed.init_process_group()指定backendhccl初始化进程组通信域集合通信torch.distributed.all_reduce()执行 AllReduce 求和归约三、环境准备3.1 环境要求该样例支持以下产品组网要求为单机 N 卡N ≥ 2Ascend 950PR / Ascend 950DTAtlas A3 训练系列产品 / Atlas A3 推理系列产品Atlas A2 训练系列产品Atlas 训练系列产品Atlas 推理系列产品Python 侧依赖两个包torchPyTorch 深度学习框架提供torch.distributed与torch.multiprocessing等分布式能力torch_npuPyTorch 的昇腾 NPU 适配插件HCCL 通过该插件接入 PyTorch 的分布式后端。3.2 配置环境变量运行前需加载 CANN 环境变量以 root 用户默认安装路径为例# 设置 CANN 环境变量 source /usr/local/Ascend/cann/set_env.sh四、示例代码逐行解析完整代码见 hccl_pytorch_allreduce_test.py核心结构如下4.1 主进程设备探测与多进程拉起def main(): print(Executing AllReduce collective operation via HCCL backend) ip 127.0.0.1 port 50001 print(Listening on %s:%d % (ip, port)) rank_size torch_npu.npu.device_count() # 查询可用 NPU 数量 print(Available NPU count: %d % rank_size) # 启动多进程每个进程绑定一张卡 mp.spawn(run_hccl, args(rank_size, ip, port), nprocsrank_size, joinTrue)要点说明torch_npu.npu.device_count()返回本机可用 NPU 数同时用作nprocs进程数和world_size通信域大小天然保证“一进程一卡”mp.spawn(..., joinTrue)为阻塞调用主进程会等待所有子进程执行完毕后再返回这保证了进程退出顺序清晰便于调试127.0.0.1:50001是单机场景下torch.distributed的 rendezvous 地址。扩展到多机训练时可将ip替换为各节点可达的 master 地址。4.2 子进程设备绑定、通信域初始化与 AllReducedef run_hccl(rank: int, world_size: int, master_ip: str, master_port: int): # 指定当前进程使用的 NPU 设备 torch_npu.npu.set_device(rank) # 初始化进程组后端使用 HCCL init_method ftcp://{master_ip}:{master_port} dist.init_process_group( backendhccl, rankrank, world_sizeworld_size, init_methodinit_method ) # 构造输入数据1行8列值为0~7 torch_tensor torch.arange(world_size, dtypetorch.float32, devicenpu) print([Rank %d] Input: %s % (rank, torch_tensor)) try: # 调用 HCCL 接口下发 AllReduce 集合通信算子 dist.all_reduce(torch_tensor, opdist.ReduceOp.SUM) except Exception as e: print([Rank %d] Error occurred: %s % (rank, e)) else: print([Rank %d] Output: %s % (rank, torch_tensor))从框架 API 到 HCCL 接口的映射关系如下可与 HCCL 原生 C 示例 examples/02_collectives/01_allreduce/main.cc 对照理解框架侧调用底层对应的 HCCL 能力说明torch_npu.npu.set_device(rank)aclrtSetDevice将当前进程绑定到指定 NPUdist.init_process_group(backendhccl, ...)HcclGetRootInfoHcclCommInitRootInfo生成 RootInfo 并初始化通信域HcclComm由插件代业务完成dist.all_reduce(tensor, opSUM)HcclAllReduce(sendBuf, recvBuf, count, FP32, HCCL_REDUCE_SUM, comm, stream)在任务流上异步下发 AllReduce 算子并在框架侧同步等待完成mp.spawn/ 进程退出HcclCommDestroyaclrtFree通信域与 Device 内存由插件在进程生命周期结束时释放其中HcclAllReduce的接口签名、缓冲区对齐要求及“注册对称内存后 sendBuf 可能被就地修改”等使用约束可参见 HcclAllReduce 接口文档。代码中值得注意的实现细节输入张量直接创建在 NPU 上torch.arange(world_size, dtypetorch.float32, devicenpu)避免了 Host→Device 拷贝每个 rank 的输入为[0., 1., 2., ..., 7.]8 个 float32 元素异常隔离try/except包裹all_reduce调用单个 rank 通信失败时打印错误而不影响其他 rank 的日志输出方便定位问题world_size 与数据长度一致这里torch.arange(world_size, ...)使张量长度等于卡数如 8 卡即 8 个元素与结果示例中“0~7 初始化、求和后为 0~56”完全对应。五、运行样例与配置展开模式5.1 运行python hccl_pytorch_allreduce_test.py5.2 可选配置通信算子展开模式HCCL_OP_EXPANSION_MODE可通过HCCL_OP_EXPANSION_MODE环境变量配置通信算子的展开模式即通信算子在哪个执行单元上展开执行不同产品支持的取值不同。以 HCCL_OP_EXPANSION_MODE 环境变量文档 为准各产品典型取值如下产品默认值支持的主要取值Ascend 950PR / 950DTAICPU_TSAI_CPU后续版本废弃由AICPU_TS替代、AICPU_TS、AICPU_CacheDisable、AIV、CCU_MS、CCU_SCHEDAtlas A3 训练/推理系列AI_CPUAI_CPU、AICPU_CacheDisable、AIVAtlas A2 训练/推理系列HOSTHOST、HOST_TS、AI_CPU仅 AllGather/AlltoAll 系算子、AIVAtlas 300I Duo 推理卡HOSTHOST、AI_CPU仅 AllReduce单机单通信域配置示例# 设置通信算子的展开模式为AI CPU通信引擎 export HCCL_OP_EXPANSION_MODEAI_CPU使用注意摘自环境变量文档的使用约束若通过 HCCL C 接口以HcclCommConfig的hcclOpExpansionMode参数按通信域粒度配置了展开模式则通信域粒度配置优先于环境变量图模式Ascend IR或图捕获aclgraph场景下通信算法采用 AI CPU 模式时单卡上的并发图数量不能超过 6 个否则可能因 AI CPU 核被占满而导致通信阻塞AIV模式仅支持对称组网、推理特性不支持多通信域并行场景且部分算子在数据量超过一定阈值时系统会自动切换为AI_CPU模式具体算子与数据类型支持范围请以该环境变量文档为准AICPU_CacheDisable用于关闭 AI CPU cache同一通信算子二次执行时复用首次展开结果的特性适合通信数据量频繁变化的服务场景以降低显存开销。六、运行结果与验证每个 rank 的数据初始化为 0~7经过 AllReduce求和操作后每个 rank 的结果是所有 rank 对应位置数据的和以 8 卡为例即 8 份数据相加第 i 个元素结果为 8i[Rank 0] Input: tensor([0., 1., 2., 3., 4., 5., 6., 7. ], devicenpu:0) [Rank 1] Input: tensor([0., 1., 2., 3., 4., 5., 6., 7. ], devicenpu:1) [Rank 2] Input: tensor([0., 1., 2., 3., 4., 5., 6., 7. ], devicenpu:2) [Rank 3] Input: tensor([0., 1., 2., 3., 4., 5., 6., 7. ], devicenpu:3) [Rank 4] Input: tensor([0., 1., 2., 3., 4., 5., 6., 7. ], devicenpu:4) [Rank 5] Input: tensor([0., 1., 2., 3., 4., 5., 6., 7. ], devicenpu:5) [Rank 6] Input: tensor([0., 1., 2., 3., 4., 5., 6., 7. ], devicenpu:6) [Rank 7] Input: tensor([0., 1., 2., 3., 4., 5., 6., 7. ], devicenpu:7) [Rank 0] Output: tensor([0., 8., 16., 24., 32., 40., 48., 56. ], devicenpu:0) [Rank 1] Output: tensor([0., 8., 16., 24., 32., 40., 48., 56. ], devicenpu:1) [Rank 2] Output: tensor([0., 8., 16., 24., 32., 40., 48., 56. ], devicenpu:2) [Rank 3] Output: tensor([0., 8., 16., 24., 32., 40., 48., 56. ], devicenpu:3) [Rank 4] Output: tensor([0., 8., 16., 24., 32., 40., 48., 56. ], devicenpu:4) [Rank 5] Output: tensor([0., 8., 16., 24., 32., 40., 48., 56. ], devicenpu:5) [Rank 6] Output: tensor([0., 8., 16., 24., 32., 40., 48., 56. ], devicenpu:6) [Rank 7] Output: tensor([0., 8., 16., 24., 32., 40., 48., 56. ], devicenpu:7)验证方法若卡数为 N则每个 rank 的输出张量第 i 个元素应恒等于N × i0 ≤ i N且所有 rank 输出一致。若某个 rank 打印Error occurred可结合 CANN 日志定位 HCCL 通信域初始化或算子执行阶段的具体报错。七、延伸阅读与扩展方向C 接口版本对照不依赖框架、直接调用 HCCL C API 的 AllReduce 单机多卡示例见 examples/02_collectives/01_allreducemain.cc中展示了aclrtMalloc申请 Device 内存、HcclCommInitRootInfo初始化通信域、HcclAllReduce下发算子、aclrtSynchronizeStream同步等待的完整生命周期可与本文的 PyTorch 版本逐行对照TensorFlow 集成HCCL 通过 TF Adapter 对接 TensorFlow 的调用示例见 examples/03_ai_framework/02_tensorflow更多集合通信算子HCCL 接口文档中还覆盖 Broadcast、AllGather、ReduceScatter、AlltoAll、Scatter、Send/Recv 等算子见 通信算子接口参考环境变量全集HCCL 各环境变量如缓冲、超时、确定性计算等说明见 HCCL 环境变量文档。将本样例作为最小可运行基线即可在自有 PyTorch 分布式训练中通过init_process_group(backendhccl)直接获得 HCCL 的集合通信能力。【免费下载链接】hccl集合通信库Huawei Collective Communication Library简称HCCL是基于昇腾AI处理器的高性能集合通信库为计算集群提供高性能、高可靠的通信方案项目地址: https://gitcode.com/cann/hccl创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考