由硬件推导集合通信

前面几节给出了跨芯片通信的全部基本操作和它们的开销。本节不介绍某个“标准”的集合通信算法，而是从这些开销出发，推导 all-reduce（每颗芯片的数组求和，结果分发回每颗芯片）应当怎样做，再用实验检验推导。目标不是背下一种算法，而是学会用硬件参数比较不同的设计。

推导用到的硬件事实

来自第 2、4–6 节与研究报告 31（K 为 KiB 数）：

设每颗芯片的数组为 M。三种设计：

一次交换：每颗芯片把整份数据同时发给其余三颗，再把收到的三份相加。只付一次固定开销（对角线的 2908 个周期），但每颗芯片发出 3M，其中发往对角线的那份还要经过中间芯片转发。数一数每条链路上的数据：四颗芯片排成环 0→1→3→2，每条链路有两个方向，共 8 个“有向链路”。每颗芯片发给两个邻居的数据各占一个有向链路，共 8 份 M，正好每个有向链路一份；四条对角线流每条走两跳，共 8 跳，按先 x 后 y 的路由，正好也是每个有向链路一跳。所以每个有向链路要传 2M，预计约 2908 + 2 × 24.77M 个周期（M 以 KiB 计）。

环形：四颗芯片按物理相邻排成环（第 5 节的 0→1→3→2），把数组分成 4 块。先 reduce-scatter：每一步每颗芯片把一块发给右边、收到左边的一块并累加，3 步后每颗芯片持有一块完整的和；再 all-gather：每一步把完整的一块传给右边，3 步后每颗芯片都有全部 4 块。共 6 步，每步只传 M/4，只用相邻链路，但要付 6 次固定开销。预计约 6 × (1941 + 24.77 × M/4) 个周期。环形会改写自己的数组，所以下文的实现每次开始前都从 HBM 重新读入，再加 483.6 + 1.105M 个周期。

双向环形：环形只用了每颗芯片向右的那条链路。把数组分成两半，一半沿环向右、一半沿环向左，同时进行。每步仍是 6 步，每步传 M/8，两条链路都用上。预计约 6 × (1941 + 24.77 × M/8) 个周期。

可以预见：M 小时固定开销占主导，一次交换最好；M 大时带宽占主导，双向环形最好：它的每个有向链路只传 6 × M/8 = 0.75M，而一次交换是 2M，单向环形是 1.5M。

写法

本小节实验源码、输出。

三种设计都是一个 Pallas kernel，在四颗芯片的 TensorCore 0 上运行，mesh 按物理环排列。它们共用的部分是一个四方汇合：

def barrier() -> None:
    ready = pltpu.get_barrier_semaphore()
    for rank in range(CHIPS):
        pl.semaphore_signal(ready, 1, device_id={'device': rank, 'tc': 0})
    pl.semaphore_wait(ready, CHIPS)

一次交换同时发出三个 remote DMA，写进对方 inbox 的不同格子，然后相加。offsets 默认为 (1, 2, 3)，即依次发给 mesh 中的 me + 1、me + 2、me + 3：

copies = [pltpu.make_async_remote_copy(acc, inbox.at[k], send_sems.at[k], recv_sems.at[k], device_id={'device': (me + offset) % CHIPS, 'tc': 0}, device_id_type=pl.DeviceIdType.MESH) for k, offset in enumerate(offsets)]
for copy in copies:
    copy.start()
for copy in copies:
    copy.wait_send()
    copy.wait_recv()
total = acc[...]
for k in range(len(offsets)):
    total = total + inbox[k]
inbox[0] = total

环形的每一步是“等下游许可 → 发送 → 等发出和收到 → 累加 → 给上游许可”：

if not first:
    pl.semaphore_wait(credits.at[d], 1)          # 下游已用完上一次收到的数据
transfer = pltpu.make_async_remote_copy(source, destination, ..., device_id={'device': (me + sign) % CHIPS, 'tc': 0}, ...)
transfer.start()
transfer.wait_send()
transfer.wait_recv()
part(index)[...] = part(index)[...] + inbox[d]   # reduce-scatter 阶段累加
pl.semaphore_signal(credits.at[d], 1, device_id={'device': (me - sign) % CHIPS, 'tc': 0})

这里的许可（credit）是一个普通信号量，它代替了每一步的四方汇合：只有相邻的两颗芯片需要互相等待，不必等所有芯片。第 4 节说过，remote DMA 直接写对方的 buffer，必须先确认对方已经用完上一次的数据，许可就是这个确认。sign 为 1 时向右、为 −1 时向左，双向环形同时执行两个方向。

all-gather 阶段把完整的块直接写进下游 acc 的同一位置，不经过 inbox，省去一次复制。

结果

计时方法与第 2 节相同：在 kernel 内部把 all-reduce 重复 64 次和 32 次，各读出 kernel 在 device 0 上的周期数，相减除以 32。各种设计的数值都与求和结果逐元素一致。“对角线最后发起”是只改一次交换中三个 DMA 的发起顺序，原因见下一小节。

单位是周期。预计值由上面的公式算出，环形包含每次从 HBM 重新读入的时间。汇合单独测得每次 1374 个周期（只做四方汇合、不传数据的 kernel），不在预计值中。除了原顺序的一次交换在 1 MiB 时多出 2.4 万个周期，实测与预计相差不大：环形与双向环形多约 1200–1700 个周期，与一次汇合相当；一次交换在小尺寸上也只多几百到一千个周期，可见这样把各部分相加的公式只是近似。本节没有再拆分。推导的结论都得到了验证：

• 16 KiB 时，环形的 6 次固定开销（约 11600 个周期）远大于数据量，一次交换快 3 倍多。
• 1 MiB 时，数据量占主导：一次交换的每个有向链路传 2M，与单向环形的 1.5M 加上 6 次固定开销相当；双向环形每个有向链路只传 0.75M，最快。
• 双向环形在 16 KiB 时没有收益：每步的数据减半，但固定开销不变。

拆开一次交换：发起顺序

原顺序的一次交换在 1 MiB 时是 78281 个周期，比链路负载推出的 53637 个多出近一半。为了找出原因，实验保留一次交换的写法，只改发给谁、按什么顺序发起：one_shot 的参数 offsets 列出每个 DMA 发给 mesh 中的 me + o，顺序就是发起的顺序。在物理环的 mesh 中，o = 1、3 是相邻芯片，o = 2 是对角线芯片：

copies = [pltpu.make_async_remote_copy(acc, inbox.at[k], send_sems.at[k], recv_sems.at[k], device_id={'device': (me + offset) % CHIPS, 'tc': 0}, device_id_type=pl.DeviceIdType.MESH) for k, offset in enumerate(offsets)]

单位是周期。前四行与链路负载吻合：一个有向链路传 M 时约 28600 个周期（1941 + 24.77 × 1024 是 27305，加上汇合的 1374），传 2M 时约 54000 个周期。两颗相邻芯片同时发送互不影响；对角线流与相邻流的时间几乎相同，带宽与一跳一样。

后五行的流量完全相同，差别只在发起顺序：只要有一个发往相邻芯片的 DMA 在发往对角线的 DMA 之后发起，1 MiB 就从约 5.4 万个周期变成约 7.8 万个；对角线放在最后发起，就回到链路负载的预计。“3、2”与“2、3”是同样的两条流，只是顺序相反。下一小节把流再拆细，找出原因。

一条链路上的流按顺序通过

本小节实验源码、输出。

为了看清发起顺序为什么起作用，实验不再让四颗芯片做同样的事，而是只发出指定的几条流。每条流写成（源, 目的），都是物理环 mesh 中的位置（位置 0、1、2、3 依次是 device 0、1、3、2）；kernel 用 pl.when(me == position) 让每颗芯片只执行属于自己的那一段：

for position in range(CHIPS):
    @pl.when(me == position)
    def _() -> None:
        outgoing = [flow(k) for k, (source, _) in enumerate(flows) if source == position]
        for copy in outgoing:
            copy.start()
        for copy in outgoing:
            copy.wait_send()
        for k, (_, destination) in enumerate(flows):
            if destination == position:
                flow(k).wait_recv()

实验只改 flows，每条流 1 MiB：

逐行读：

• 对角线流走哪边。 位置 0 的对角线流 (0, 2) 与发往位置 1 的相邻流一起发时要两倍时间，与发往位置 3 的相邻流一起发时不变：它经过位置 1，占用链路 0→1 和 1→2。位置 1 的对角线流 (1, 3) 与位置 0 发出的 (0, 3) 一起时也是两倍：它经过位置 0，占用 1→0 和 0→3。换成芯片坐标，两条都是先沿 x、再沿 y，证实了前面表中的推断。
• 同一个源的两条流共用一条链路时，顺序无关。 (0, 1) 与 (0, 2) 无论谁先，都是约 5.3 万个周期：这条链路总共要传 2 MiB。
• 三条流的最后两行是关键。 流量相同，只是位置 1 的两条流顺序相反，时间相差 2.4 万个周期。按链路数：0→3 上有 (0, 3) 和 (1, 3) 的第二跳，1→0 上有 (1, 0) 和 (1, 3) 的第一跳，每条链路都是 2 MiB，按链路负载应当都是约 5.4 万个周期。

7.7 万个周期是三条流一条接一条的时间。这说明一条链路上的流不是交替着分享带宽，而是按顺序一条一条通过，并且两跳的流在等第二跳时一直占着第一跳：

1. 位置 0 的 (0, 3) 和位置 1 的 (1, 3) 几乎同时发起。(0, 3) 就在链路 0→3 的起点，立刻占用它；(1, 3) 占用 1→0，等它到达位置 0 时，0→3 已经被占。
2. (1, 3) 停在位置 0 等 0→3，同时占着 1→0；排在它后面的 (1, 0) 也只能等。这是第一个 2.6 万个周期。
3. (0, 3) 传完，(1, 3) 通过，第二个 2.6 万个周期；然后才轮到 (1, 0)，第三个。

把位置 1 的顺序改为先发 (1, 0)：(1, 0) 与 (0, 3) 用的是不同的链路，同时进行；等它们传完，(1, 3) 的两跳都空了。只要两段。

回到一次交换。按 1、2、3 的顺序发起时，每颗芯片的对角线流都排在自己的一条相邻流之前、另一条相邻流之后，而它的第二跳会被中间芯片自己的相邻流抢先占用：对角线流占着第一跳等待，后面的相邻流被它挡住，于是成了三段。对角线最后发起时，所有相邻流先同时传完（每个有向链路一条），再同时传所有对角线流（每个有向链路一跳），是两段。

所以设计跨芯片通信时，除了数每条链路上的数据量，还要看流的先后：让只走一跳的流先发，多跳的流后发，避免多跳的流占着前一跳等后一跳。

对照：原生 XLA

本小节实验源码、输出。

在 shard_map 中调用 jax.lax.psum，XLA 生成一个 psum 段，其中有 16 条 dma.general、20 条 vsyncadd.remote，三种大小都相同。读清单的开头可以看出它和谁通信：它从 [smem:0x0] 读出本芯片的编号，拆出坐标的 x、y 位，算出 x + 1、x − 1、y + 1、y − 1 四个邻居的编号，汇合与数据都只发给这四颗芯片，从不直接发往对角线。在 2×2 中，x + 1 与 x − 1 是同一颗芯片，y 方向也一样，所以每个邻居要收两路数据。传数据的部分是几个循环，循环体中 4 条 dma.general 同时发往四个方向，ici_dest 分别由这四个编号拼成，每条 4 KiB 或 8 KiB。这符合沿 x、y 两个方向各做一次双向环形的结构，本节没有进一步逐条验证。

计时与 Pallas 版本一样取差值：用 fori_loop 循环 64 次与 32 次，用 KernelClock.time_ops 在整个程序的起止处插入 LCC 读数，取 device 0 的 TensorCore 0 上的周期数之差，除以 32：

XLA 对所有大小使用同一种算法，它在大数组上与双向环形相当，在小数组上的耗时是一次交换的 1.7–3.6 倍：后者正是固定开销占主导、一次交换更合适的区域。被计时的两种循环不完全相同（XLA 的循环在 kernel 之外，每次多一次乘法，数组可能被放进 CMEM，第 2 节），但小数组时差距主要来自算法本身的固定开销次数。

这个对照说明了手写的价值所在：XLA 必须选一种对所有情况都过得去的算法；知道硬件参数和自己的数据大小，就能选出更合适的那一种。本节的设计还有继续改进的余地，例如让每颗芯片的两个 TensorCore 各负责一半（第 1 节），或在中等大小时把一次交换与环形结合。在两个 host 的 TPU v4 (2x2x2) 上，芯片多了 z 方向的链路，研究报告 32 按同样的思路推导了三维的 all-reduce。