集合通信入图

功能简介

集合通信入图能够避免断图，并拥有更大的成图范围，从而获得更大的资源调度与融合收益，同时在整图层面实现通信与计算并行优化。

原生PyTorch社区对集合通信算子入图的支持度尚不完善，功能正在持续增强中。针对这一现状，TorchAir提供的NPU图编译后端实现了集合通信算子的Ascend Converter，调用torch.compile时，默认已支持集合通信算子入图。入图方法参见使用方法。

目前支持入图的集合通信API如下表所示，请根据实际业务需求按需调用。注意，集合通信算子入图的前提是PyTorch脚本中所有算子均能正常以Eager模式运行。

使用约束

本功能要求PyTorch 2.6.0及以上版本。

使用方法

无需修改PyTorch脚本，直接调用torch.compile，NPU图编译后端npu_backend默认集成了集合通信算子入图能力。

1
2
3
4
5
6
7
8

import torch
import torch_npu
import torchair
config = torchair.CompilerConfig()
npu_backend = torchair.get_npu_backend(compiler_config=config)
.........
# 多卡模型调用compile，后端提供集合通信入图能力
model = torch.compile(model, backend=npu_backend)