NVSHMEM 团队概念：Team 的定义、创建与销毁

聊到 NVSHMEM 的集合通信，绕不开一个核心概念——Team。说白了，Team 就是一组 PE 的集合。你想想看，如果每次做集合操作都要把所有 PE 拉进来，那多浪费？有些场景下，我们只想让一部分 PE 协作，比如做子区域的归约。这时候，Team 就派上用场了。

什么是 Team？

Team 本质上是一个逻辑分组。它把全局的 PE 空间划分成若干个子集。每个 Team 内部有自己独立的排名（rank），从 0 开始编号。我习惯把 Team 理解成「通信域」——在这个域里，PE 之间可以高效地做广播、归约、收集等操作。

举个例子，假设你有 8 个 GPU，编号 0 到 7。你可以创建一个 Team 包含 GPU 0、2、4、6。在这个 Team 里，GPU 0 的 rank 是 0，GPU 2 的 rank 是 1，以此类推。这样做广播时，只有这 4 个 GPU 参与，其他 GPU 不受影响。

核心要点： Team 是 NVSHMEM 集合通信的基本单位。所有集合操作（如 shmem_barrier、shmem_sum_reduce）都基于 Team 执行。

Team 的创建

创建 Team 有两种方式：从全局拆分和从已有 Team 拆分。我个人更推荐后者，因为更灵活。

先看一个最简单的例子——从全局 Team（NVSHMEM_TEAM_WORLD）拆分出一个子 Team：

// 创建一个包含偶数 PE 的 Team
int pe_stride = 2;          // 步长
int pe_start = 0;           // 起始 PE
int size = npes / 2;        // Team 大小

nvshmem_team_t even_team;
nvshmem_team_split_strided(NVSHMEM_TEAM_WORLD, 
                           0,          // 当前 PE 在父 Team 中的 rank
                           npes,       // 父 Team 大小
                           pe_start, 
                           pe_stride, 
                           size, 
                           NULL,       // 配置参数，通常传 NULL
                           &even_team);

嗯，这里要注意：nvshmem_team_split_strided 是创建 Team 最常用的 API。它通过步长和起始位置来选取 PE。我在项目中遇到过一种情况——需要把 16 个 GPU 分成 4 组，每组 4 个连续 PE。用这个 API 配合不同的 stride 和 size 就能轻松搞定。

还有一种更通用的方式——通过 PE 列表创建：

// 通过 PE 列表创建 Team
int members[] = {0, 3, 5, 7};
nvshmem_team_t custom_team;
nvshmem_team_split_2d(NVSHMEM_TEAM_WORLD, 
                      0, 
                      members, 
                      4, 
                      NULL, 
                      &custom_team);

这种方式适合那些「不规律」的分组。比如你想把计算节点 0 的 GPU 0 和计算节点 1 的 GPU 3 放在一个 Team 里——用列表方式最直接。

Team 的销毁

创建了 Team，用完就得销毁。不然会泄漏资源。销毁很简单：

nvshmem_team_destroy(even_team);
nvshmem_team_destroy(custom_team);

我曾经犯过一个错误——在销毁 Team 之后，还用它做集合通信。结果程序直接崩溃。后来我养成了一个习惯：销毁后立即把 team 句柄置为 NVSHMEM_TEAM_INVALID，这样再用的时候能及时发现。

注意： 不要销毁 NVSHMEM_TEAM_WORLD 和 NVSHMEM_TEAM_SHARED。这两个是预定义的 Team，由 NVSHMEM 管理生命周期。你销毁它们会导致未定义行为。

Team 在集合通信中的角色

Team 在集合通信中扮演什么角色？说白了，它就是通信的边界。

所有的集合操作——shmem_barrier、shmem_sum_reduce、shmem_broadcast、shmem_collect——都需要指定一个 Team。Team 决定了哪些 PE 参与这次操作。

举个例子，假设你有两个 Team：Team A（PE 0-3）和 Team B（PE 4-7）。当你调用：

// 只在 Team A 内做同步
nvshmem_barrier_all_in_team(team_a);

只有 PE 0-3 会等待彼此。PE 4-7 完全不受影响。它们可以继续做自己的计算，或者在自己的 Team 里做别的集合操作。

这种「隔离性」在实际项目中非常有用。我做过一个多物理场模拟的程序——流体部分用 PE 0-7，结构部分用 PE 8-15。两个 Team 各自做归约和广播，互不干扰。如果只用全局 Team，每次同步都要等所有人，性能会差很多。

Team 的层级结构

NVSHMEM 支持 Team 的层级结构。你可以从 Team A 再拆出 Team A1、Team A2。这种层级关系在分层并行中特别有用。

我画了一张图，帮你理解 Team 的层级关系：

NVSHMEM Team 层级结构 NVSHMEM_TEAM_WORLD (PE 0-7) Team A PE 0,2,4,6 Team B PE 1,3,5,7 Team C PE 0,1,2,3 Team A1 PE 0,4 Team A2 PE 2,6 Team B1 PE 1,5 Team B2 PE 3,7 图例 全局 Team 子 Team 再拆分

从图中你能看到：全局 Team 可以拆成多个子 Team，子 Team 还能继续拆。这种层级结构让通信模式变得非常灵活。比如在分层并行中，你可以先做节点内的归约（用节点内 Team），再做节点间的归约（用跨节点 Team）。

Team 的查询与属性

创建了 Team 之后，经常需要查询它的属性。常用的 API 有：

API 功能 返回值
nvshmem_team_n_pes 获取 Team 中的 PE 数量 int
nvshmem_team_my_pe 获取当前 PE 在 Team 中的 rank int
nvshmem_team_translate_pe 将一个 Team 中的 rank 映射到另一个 Team int
nvshmem_team_get_config 获取 Team 的配置信息 int

其中 nvshmem_team_translate_pe 特别有用。我在做跨 Team 通信时经常用它——比如 Team A 的 rank 2 对应全局的 PE 4，我想知道它在 Team B 中的 rank 是多少。这个 API 直接搞定。

小技巧： 如果你不确定当前 PE 是否属于某个 Team，可以用 nvshmem_team_my_pe 检查返回值。如果返回 -1，说明当前 PE 不在这个 Team 中。

预定义的 Team

NVSHMEM 提供了两个预定义的 Team：

  • NVSHMEM_TEAM_WORLD：包含所有 PE。这是默认的通信域。
  • NVSHMEM_TEAM_SHARED：包含共享内存空间内的 PE。通常用于节点内的快速通信。

这两个 Team 不需要创建，也不能销毁。它们由 NVSHMEM 在初始化时自动创建。我建议你尽量使用自定义 Team，而不是直接操作 NVSHMEM_TEAM_WORLD。原因很简单——自定义 Team 能更好地控制通信范围，减少不必要的同步开销。

好了，关于 Team 的基本概念就聊到这里。记住：Team 是 NVSHMEM 集合通信的基石。理解 Team 的创建、销毁和层级关系，你就能灵活地组织通信模式，写出高效的并行程序。

公众号:蓝海资料掘金营，微信deep3321

← 上一章 📖 返回目录 下一章 →