Ascend for PyTorch框架HCCL对称内存，提升通信带宽30%

在大模型训练与推理中，集合通信贯穿多个关键环节：专家并行依赖All-to-All完成token分发和回传，分片参数的聚合使用AllGather，梯度归约与分片使用ReduceScatter，张量并行则经常调用AllReduce汇总局部计算结果。通信性能不仅取决于链路带宽，还受到内存拷贝、地址映射和同步开销的影响。Ascend for PyTorch（下称TorchNPU）新增基于MemPool实现HCCL对称内存方案。开发者通过HCCL专用分配器创建MemPool，再将池内内存注册到通信域，使HCCL能够直接使用业务缓冲区执行集合通信，减少中间缓冲区的拷贝开销。该方案采用与PyTorch NCCL后端相近的接口，开发者只需修改一行代码便可以将代码迁移到昇腾NPU上。在Atlas 800T A3上测试，HCCL对称内存让集合通信算子性能提升高达30%。

HCCL对称内存核心原理

HCCL对称内存通过一致的虚拟地址布局建立跨rank映射，使集合通信能够直接访问已注册的业务内存，减少中间拷贝。TorchNPU通过MemPool统一管理通信缓冲区，自动维护内存段的注册关系，便于缓冲区复用并减少手动管理开销。这一机制可在沿用现有PyTorch内存管理机制的基础上优化通信路径、提升通信带宽，经测试，使能HCCL对称内存后，特定集合通信算子通信带宽最多提升约30%。

建立跨rank的对称内存访问关系，减少通信中间拷贝

参与集合通信的每个rank都拥有自己的本地物理内存，HCCL对称内存将各rank的物理内存映射到布局一致的虚拟地址窗口，使通信库可以根据目标rank和窗口内偏移访问已经注册的内存，减少中间拷贝。其基本内存布局如下：

可以从三个方面理解这套布局：

为每个rank申请虚拟内存，假设每个rank对应的虚拟内存大小为heap_size，通信域中的rank数量为rank_size，则通信域中总虚拟内存大小为heap_size*rank_size。

每个rank的虚拟地址布局都是相同的。

不同rank的物理地址映射到每个rank对应位置的虚拟地址，实现对其他rank内存的访问。

TorchNPU负责将MemPool中的内存段接入HCCL注册流程，并跟踪其生命周期；HCCL负责底层对称窗口的注册与地址映射。满足相关约束后，集合通信可以直接使用注册的业务内存，无需经过中间缓冲区(HCCL buffer)，从而减少内存拷贝开销。

以MemPool为单位管理通信内存，简化分配与注册

以MemPool为单位管理通信内存，可减少逐段维护注册关系的负担，并复用已有的PyTorch内存管理流程，降低上层框架的接入与维护成本。TorchNPU将这套能力组织为三个接口：

通过以上接口，开发者可以将指定范围内的张量分配路由到这个MemPool。三个接口的基本用法如下所示：

#使用HCCL专用分配器创建内存池pool = torch_npu.npu.MemPool(backend.mem_allocator)#注册内存池backend.register_mem_pool(pool, symm=True)#通信完成后注销内存池backend.deregister_mem_pool(pool)

自动维护内存段的注册关系，降低手动管理开销

HCCL对称内存特性的会自动维护内存段的注册关系，可减少动态分配与释放时的手动操作，降低漏注册、漏注销的风险，便于应用复用通信缓冲区。注册MemPool时，当前实现会遍历内存池的分配器快照，将该MemPool已有内存段逐个注册到 HCCL通信域。

注册完成后，分配器的回调函数负责处理新增和释放的内存段：新增内存段自动注册，释放内存段自动注销。因此，应用可以在初始化阶段完成分配和注册，后续反复复用通信缓冲区，避免在每次集合通信前重新注册内存池。

明确使用路径，避免接口混用

本文讨论的是HCCL对称内存的MemPool注册路径。它通过 backend.mem_allocator和register_mem_pool()使用HCCL对称窗口。与 NPUSHMEM的empty()和rendezvous()方式使用不同的分配器。两者都涉及对称内存，但分配、注册和调用接口不同, 内存访问关系不会自动互通。按各自流程配套使用接口，可避免内存来源不匹配或遗漏注册，减少接入错误和排查成本。两者的区别如下图所示：

TorchNPU上对称内存的两套API实现

HCCL对称内存使用指南及收益

确认运行条件

HCCL对称内存特性对硬件平台和软件版本都存在约束，使用者在测试该特性时需要确认以下条件：

TorchNPU版本不低于26.2.0、PyTorch版本大于v2.9.0以及CANN版本不低于9.2.0。

使用昇腾A3训练或推理系列产品，并满足超节点内HCCS通信及对称组网要求。

仅使用受支持的集合通信算子：AllReduce、AllGather、ReduceScatter或AllToAll。

各rank按一致的顺序分配和注册对应内存段，保证注册段大小一致。

准备测试环境

以下步骤以Atlas 800T A3单机8卡为例，展示了从配置脚本到成功运行HCCL测试脚本的整体流程：

a) 下载CANN-9.2.0和TorchNPU 26.2.0的安装包

b) 执行以下命令安装CANN和TorchNPU，这里以aarch64架构为例:

mkdir ./cann-9.2.0Ascend-cann-toolkit_9.2.0_linux-aarch64.run --install –install-path=`realpath ./cann-9.2.0`Ascend-cann-A3-ops_9.2.0_linux-aarch64.run –install –install-path=`realpath ./cann-9.2.0`pip install torch_npu-2.13.0-cp312-cp312-manylinux_2_28_aarch64.whl

c) 使能CANN的环境变量配置

source ./cann-9.2.0/ascend-toolkit/set_env.sh

d) 基于配置好的运行环境，运行Atlas 800T A3单机8卡的测试脚本，该脚本基于TorchNPU的HCCL对称内存特性，如果成功运行便说明HCCL对称内存特性使能成功：

import osimport torchimport torch.distributed as distimport torch_npulocal_rank = int(os.environ["LOCAL_RANK"])torch.npu.set_device(local_rank)device = torch.device("npu", local_rank)dist.init_process_group(backend="hccl")rank, world_size = dist.get_rank(), dist.get_world_size()# 1. 预热，确保 HCCL 通信域已建立dist.all_reduce(torch.ones(8, device=device))torch.npu.synchronize(device)# 2. 创建通信专用内存池并注册backend = dist.group.WORLD._get_backend(device)pool = torch_npu.npu.MemPool(backend.mem_allocator)backend.register_mem_pool(pool, symm=True)with torch_npu.npu.use_mem_pool(pool, device=device):    x = torch.full((1024,), rank + 1.0, device=device)# 3. 在池内张量上执行集合通信dist.all_reduce(x) # 默认即 ReduceOp.SUMtorch.npu.synchronize(device)expected = world_size * (world_size + 1) / 2torch.testing.assert_close(x.cpu(), torch.full((1024,), expected))if rank == 0:     print(f"AllReduce 校验通过，结果为 {expected}")# 4. 收尾dist.barrier()torch.npu.synchronize(device)backend.deregister_mem_pool(pool)del x, pooldist.destroy_process_group()

在满足上述条件的单机8卡环境中启动：

torchrun --standalone --nproc_per_node=8 hccl_symm_pool_demo.py

执行成功后可以看到如下的打印：

AllReduce 校验通过，结果为 36.0

技术收益

测试Atlas 800T A3单机8卡上不同的通信数据量下AllGather，AllReduce，AlltoAll，ReduceScatter这四个集合通信算子相比于原生集合通信算子的性能收益。HCCL对称内存特性在通信数据量较大的情况下效果较好。其中，AllReduce、AllToAll和ReduceScatter算子的性能收益最大，在数据量大于等于1GB时，性能提升约25%-38%，AllGather的性能提升幅度较小，可以提升约10%的带宽。当集合通信数据量较小时（小于等于8MB），所有集合通信算子提升幅度均较小，其中AllGather、AllReduce和ReduceScatter带宽提升约为5%，AllToAll的带宽提升不明显。

结语

HCCL对称内存将通信优化的一部分工作前移到内存分配和注册阶段。应用仍然调用普通的集合通信接口，通过专用内存池和对称窗口，实现对集合通信操作的加速。

对于大模型训练与推理，这套机制提供了一条减少中间拷贝、改善通信延迟的路径。其最终价值，应由实际业务中的正确性、稳定性与端到端性能共同验证。

后续Ascend for PyTorch会进一步完善支持对称内存后端NPUSHMEM，围绕对称内存、单边读写和细粒度同步提供完整的编程能力，使开发者能够根据业务需求灵活组织跨rank数据交换，并探索通信和计算的细粒度融合，为训练和推理场景提供更多的性能优化空间。若您对Ascend for PyTorch和对称内存感兴趣，可以关注：

Ascend for PyTorch：

https://gitcode.com/Ascend/pytorch

SHMEM对称内存: 

https://gitcode.com/cann/shmem

HCCL通信基础库：

https://gitcode.com/cann/hcomm