集合通信入图

功能简介

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

原生PyTorch社区对集合通信算子入图的支持度尚不完善，功能正在持续增强中。针对这一情况，npugraph_ex实现了集合通信算子的Ascend Converter，在调用torch.compile时默认已支持集合通信算子入图，具体入图方法参见使用方法。

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

表 1 集合通信API入图支持情况

|PyTorch集合通信API|支持情况| |--|--|--| |torch.distributed.all_gather|√| |torch.distributed.all_gather_into_tensor|√| |torch.distributed.all_reduce|√| |torch.distributed.all_to_all|√| |torch.distributed.all_to_all_single|√| |torch.distributed.broadcast|√| |torch.distributed.reduce_scatter_tensor|√| |torch_npu.distributed.all_gather_into_tensor_uneven|√| |torch_npu.distributed.reduce_scatter_tensor_uneven|√| |torch.distributed.send|√| |torch.distributed.recv|√|

使用约束

本功能支持的产品型号参见使用说明。

使用方法

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

import torch
import torch_npu

# 多卡模型调用compile，后端提供集合通信入图能力
model = torch.compile(model, backend="npugraph_ex", fullgraph=True, dynamic=False)