CANN SHMEM Team通信域深度解析：如何用team_split灵活切分多卡分组？

CANN SHMEM Team通信域深度解析：如何用team_split灵活切分多卡分组？

【免费下载链接】shmemCANN SHMEM 是面向昇腾平台的多机多卡内存通信库，基于OpenSHMEM 标准协议，实现跨设备的高效内存访问与数据同步。项目地址: https://gitcode.com/cann/shmem

CANN SHMEM 是面向昇腾平台的多机多卡内存通信库，基于 OpenSHMEM 标准协议实现跨设备的高效内存访问与数据同步。在实际训练中，算子往往只需要部分 rank 参与计算（例如张量并行组内做 AllReduce、组间做数据交换），这时就要用到 SHMEM 的Team 通信域机制：通过team_split系列接口，你可以从全局通信域中灵活切分出任意子分组，并基于子分组做同步与远程内存访问。本文将带你快速吃透 Team 的三大核心概念：team_split_strided、team_split_2d以及成员查询/地址转换接口。

在 SHMEM 的整体架构中，"通信域管理"是 Host APIs 与 Device APIs 共同提供的核心能力，与初始化、内存管理、同步机制并列，构成多卡编程的骨架。

一、什么是Team通信域？如何理解mype与全局mype的区别？

Team 是 SHMEM 中的通信域概念，相关接口通过team_id访问。初始化完成后，SHMEM 会自动创建一个默认全局通信域：

#define ACLSHMEM_TEAM_WORLD 0

每个 team 内部记录了：当前进程在该 team 内的 rank（mype）、起始 rank、rank 步长、rank 数量等信息。

⚠️新手最容易踩的坑：team内部存储的mype/size是team 内部视角，而全局 state 里的mype/npes是全局视角。

举例：4 个 rank 初始化 SHMEM，前 2 卡和后 2 卡各组成一个 team：

经验法则：team 级别的mype用作算子内部资源数组的索引；全局mype用作全局共享内存地址信息的索引。

初始化时，SHMEM 会创建大小为ACLSHMEM_MAX_TEAMS的全局 team 池，并把全局 team（start=0, stride=1, size=npes）填入其中，同时分配 team 级同步所需的sync_pool、sync_counter等资源。之后team_split切分出的每个子 team 都会占用池中的一个槽位。

二、team_split_strided一键切分：掌握start、stride、size三个参数

SHMEM 提供了专门的集合接口进行子 Team 切分：

int aclshmem_team_split_strided(aclshmem_team_t parent_team, int pe_start, // 起始pe int pe_stride, // 步长 int pe_size, // 新team中pe个数 aclshmem_team_t *new_team); // 出参：新team_id

以初始化好 8 个 rank 的场景为例，从全局域中 idx 为 1 的 pe 开始、步长 2、切分 3 个 pe：

aclshmem_team_t new_team; aclshmem_team_split_strided(ACLSHMEM_TEAM_WORLD, 1, 2, 3, &new_team);

切分后new_team包含全局 pe 1、3、5，它们在 team 内部的 mype 分别是 0、1、2。

三个参数的直观理解：

• pe_start：新 team 的第一个成员（父 team 视角的 pe 编号）
• pe_stride：成员之间的间隔，stride=2就是"隔一个选一个"
• pe_size：选到多少个成员后停止

由于这是集合操作，所有 rank 都要调用该接口；当前 pe 不属于新 team 时，返回值为ACLSHMEM_TEAM_INVALID，后续查询接口对它返回 -1，注意判断。

切分得到的 team 信息会存放在设备的元数据空间中，host 侧 state 变化后会自动同步到设备，因此设备侧算子可以直接使用 team_id完成通信域内操作，无需额外传递。

三、team_split_2d二维切分：并行场景的终极利器

对于张量并行（TP）这类"行维组内通信 + 列维跨组通信"的 2D 网格场景，SHMEM 提供了基于二维笛卡尔空间的切分接口：

int aclshmem_team_split_2d(aclshmem_team_t parent_team, int x_range, // 第一维元素数量 aclshmem_team_t *x_team, // 出参：x轴team aclshmem_team_t *y_team); // 出参：y轴team

一次调用同时拿到 x 轴 team 和 y 轴 team。在示例 examples/tp_allreduce_udma/main.cpp 中，正是用它从全局域切分出 TP 组内/跨 TP 组两个 team，完成 2D 并行下的 AllReduce：

status = aclshmem_team_split_2d(ACLSHMEM_TEAM_WORLD, TP_SIZE, &tpTeam, &crossTpTeam);

类似的实战用法也可以参考 examples/rdma_sync_barrier_demo/main.cpp 中用team_split_strided构造子组做 barrier 同步的完整流程。

四、配套接口速查：成员查询与地址翻译

切分完成只是第一步，日常开发中更高频的是这几个"查询 + 翻译"接口：

其中aclshmem_team_translate_pe特别值得记住：当你手里只有"子 team 里的编号"，却需要在全局域（或另一个 team）中定位同一张卡时，它一步完成翻译。设备侧还有aclshmem_team_pe_mapping可以把 team 内 pe 直接映射到全局 pe，方便算子内核中使用。

以上接口的完整声明可以查看 include/host/team/shmem_host_team.h 与 include/device/team/shmem_device_team.h，host 侧实现位于 src/host/team/shmem_team.cpp。

五、动手实践：从示例与Python接口快速上手

• C++ 示例：docs/example/api_demo.md 提供了 host 侧切分子通信域的完整样例（含取值判断与aclshmem_team_destroy资源释放），可直接照搬到你的工程中。
• 原理文档：docs/principles/team.md 对 team 池、mype 视角差异与切分细节有更深入的说明。
• Python 绑定：如果你走 PyTorch 集成路线，src/host/python_wrapper/pyshmem.cpp 暴露了team_split_strided、team_split_2d等绑定，接口参数说明见 docs/api/pythonAPI.md。

三条实用建议：

1. 🧩 切分前先想清楚通信拓扑：TP 组等"行/列"结构用team_split_2d，奇偶分组、跨机子集等结构用team_split_strided。
2. 📏 注意pe_start + (pe_size-1)*pe_stride不能超过父 team 大小，越界会导致切分失败。
3. 🧹 用完记得aclshmem_team_destroy：team 池容量有限（ACLSHMEM_MAX_TEAMS），长期不释放会耗尽可用槽位。

六、小结

• Team = 通信域：全局域ACLSHMEM_TEAM_WORLD是起点，一切子分组都从它（或其他 team）切分而来。
• 两种切分：team_split_strided（步长式，灵活）+team_split_2d（网格式，面向 2D 并行）。
• 两套视角：team 内 mype 管"局部索引"，全局 mype 管"地址索引"，用translate_pe/pe_mapping打通。
• 同步接口同样吃 team_id：aclshmem_barrier(aclshmem_team_t tid)等接口直接以 team 为范围，实现组内细粒度同步。

掌握 team_split 之后，你就可以像"切蛋糕"一样把任意规模的多机多卡集群切成最贴合算子结构的通信分组，让每一组 rank 只做自己该做的通信，这正是 SHMEM 在多卡并行编程中灵活性的来源。

【免费下载链接】shmemCANN SHMEM 是面向昇腾平台的多机多卡内存通信库，基于OpenSHMEM 标准协议，实现跨设备的高效内存访问与数据同步。项目地址: https://gitcode.com/cann/shmem