使用 PyTorch 执行 AllReduce 操作【免费下载链接】hccl集合通信库Huawei Collective Communication Library简称HCCL是基于昇腾AI处理器的高性能集合通信库为计算集群提供高性能、高可靠的通信方案项目地址: https://gitcode.com/cann/hccl样例介绍本样例展示如何使用 PyTorch 接口执行 AllReduce 操作包含以下功能点设备检测通过torch_npu.npu.device_count()接口查询可用设备数量。通过torch.multiprocessing.spawn()接口拉起多进程。在每个进程中通过torch.distributed.init_process_group()接口初始化通信域。在每个进程中通过torch.distributed.all_reduce()接口执行 AllReduce 操作。环境准备环境要求本样例支持以下产品Ascend 950PR / Ascend 950DTAtlas A3 训练系列产品 / Atlas A3 推理系列产品Atlas A2 训练系列产品Atlas 训练系列产品 / Atlas 推理系列产品配置环境变量# 设置 CANN 环境变量以 root 用户默认安装路径为例 source /usr/local/Ascend/cann/set_env.sh执行样例python hccl_pytorch_allreduce_test.py注意可通过设置HCCL_OP_EXPANSION_MODE环境变量配置通信算子的展开模式不同产品型号支持的范围可参考环境变量列表中该环境变量的使用方法。# 设置通信算子的展开模式为AI CPU通信引擎 export HCCL_OP_EXPANSION_MODEAI_CPU结果示例每个 rank 的数据初始化为 0~7经过 AllReduce 操作后每个 rank 的结果是所有 rank 对应位置数据的和8 个 rank 的数据相加。[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)【免费下载链接】hccl集合通信库Huawei Collective Communication Library简称HCCL是基于昇腾AI处理器的高性能集合通信库为计算集群提供高性能、高可靠的通信方案项目地址: https://gitcode.com/cann/hccl创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考