TorchAir集合通信入图完整指南patch_for_hcom让AllReduce/AllGather直接跑进NPU计算图【免费下载链接】torchairTorchAir 支持用户基于PyTorch框架和torch_npu插件在昇腾NPU上使用图模式进行推理。项目地址: https://gitcode.com/Ascend/torchairTorchAir 是基于 PyTorch 框架和 torch_npu 插件、面向昇腾 NPU 的图模式推理方案。它的patch_for_hcom接口可以把 AllReduce、AllGather 等集合通信算子直接打补丁patch进 NPU 计算图避免断图让通信与计算在整图层面并行调度——这正是多卡推理性能优化的关键一步。本文面向新手带你快速理解「集合通信入图」的原理、用法与注意事项。什么是集合通信入图为什么要关心它在多卡并行推理中torch.distributed的集合通信 API如all_reduce、all_gather经常出现。如果这些算子无法进入计算图编译器就只能在通信处「断图」——整张图被切成多段失去统一的资源调度、算子融合和通信-计算重叠优化的机会。集合通信入图让通信算子作为图中的一个节点参与编译✅ 避免断图成图范围更大✅ 获得更大的资源调度与融合收益✅ 在整图层面实现通信与计算并行优化 前提模型脚本中所有算子都能以 Eager 模式正常运行。patch_for_hcom 是如何工作的原生 PyTorch 社区对集合通信算子入图的支持尚不完善尤其在 PyTorch 2.1 版本。TorchAir 的做法很巧妙打补丁torchair.patch_for_hcom()会把 PyTorch 的集合通信 API如all_reduce替换为可被 Dynamo 追踪的版本。追踪期间调用被转成 functional collective如all_reduce的 tensor 语义从而能被捕获进 FX 图。转换器接管TorchAir 的 GE 转换器为这些通信算子实现了 Ascend Converter将其映射为图中的通信节点并记录通信组process group信息。相关转换逻辑位于 hcom_allreduce.py、hcom_allgather.py、hcom_alltoall.py、hcom_broadcast.py、hcom_reducescatter.py通信组管理见 hcom_utils.py。默认集成在较新的 PyTorch 版本上npu_backend默认已集成集合通信入图能力多数场景直接torch.compile即可而dynamo_export等导出流程在 PyTorch 2.3 以下版本会自动调用patch_for_hcom见 npu_export.py保证功能对齐。最快上手三步让 AllReduce 跑进 NPU 计算图 只需在图执行之前调用一次patch_for_hcom()然后正常torch.compileimport torch import torch_npu import torchair # 第 1 步在图执行之前调用 patch torchair.patch_for_hcom() # 第 2 步创建后端并编译模型 config torchair.CompilerConfig() npu_backend torchair.get_npu_backend(compiler_configconfig) opt_model torch.compile(model, backendnpu_backend) # 第 3 步正常调用AllReduce 已作为图内节点执行 out opt_model(x, y)一个多卡 AllReduce 的完整冒烟用例可以参考 hcom_smoke_patch_for_allreduce.py初始化 hccl 进程组 → 编译含torch.distributed.all_reduce的模型 → 断言结果与期望一致。更多多卡测试可看 hcom_st.py、hcom_cache_st.py。支持入图的集合通信 API 清单根据官方文档通信入图特性说明目前支持入图的 API 如下| PyTorch 集合通信 API | 支持情况 | | -- | -- | |torch.distributed.all_gather| √ | |torch.distributed.all_gather_into_tensor| √ | |torch.distributed.all_reduce| √ | |torch.distributed.all_to_all/all_to_all_single| √ | |torch.distributed.broadcast| √ | |torch.distributed.reduce_scatter_tensor| √ | |torch.distributed.send/recv| √ | |torch_npu.distributed.all_gather_into_tensor_uneven| √ | |torch_npu.distributed.reduce_scatter_tensor_uneven| √ |新手必知的注意事项 ⚠️async_opTrue不受支持编译模式下集合通信 API 不支持异步参数传入会抛出明确断言错误见 hcom_allreduce.py。send / recv 需配套使用二者要成对出现dynamicTrue场景下不同 shape 会生成不同的 FX 图多 send/recv 共存时建议设置图遍历顺序为 StableRDFS稳定拓扑序策略。ReduceOp 约束reduce_scatter_tensor_uneven不支持ReduceOp.PRODUCT。DeepSpeed 适配补丁同样会作用于 DeepSpeed 的all_reduce/inference_all_reduce方便存量 DeepSpeed 推理服务平滑迁移到图模式。延伸阅读与相关模块API 文档patch_for_hcom.md通信图特性communication_graph.mdnpugraph_ex 中的通信算子定义ops/_hcom_allgather.py、ops/_hcom_send_recv.py一句话总结在 PyTorch 2.1 上多卡跑图模式torchair.patch_for_hcom()一行代码即可让 AllReduce/AllGather 等通信算子直接跑进 NPU 计算图配合npu_backend编译获得不断图的完整推理图。【免费下载链接】torchairTorchAir 支持用户基于PyTorch框架和torch_npu插件在昇腾NPU上使用图模式进行推理。项目地址: https://gitcode.com/Ascend/torchair创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考