Skip to content

ReduceScatter

二进制: extracted/libtpu-0.0.40-cp314-cp314-manylinux_2_31_x86_64/libtpu/libtpu.so(构建 ID 89edbbe81c5b328a958fe628a9f2207d,构建 libtpu_lts_20260413_b_RC00.text VMA == 文件偏移 0xe63c000)。 状态: 重新实现级 · 证据等级: 已确认(字节锚定)——RS 环路/成本分解和 SC 卸载扁平引脚与 IDA 反编译进行交叉检查 · 第 XIII 部分 — On-Pod 集体和障碍 /集体算法 · 返回索引

摘要

ReduceScatter 是执行归约的全归约的一半:它对每个进行求和 参与者的贡献,让每个参与者持有不同的“碎片” 减少结果。在物理环面上,这被实现为一个单向环,其中 每一步都会从上游邻居接收一个分片,将其就地减少到本地 累加器,并将结果转发给下游邻居。 P-1 步骤后(对于环 P-1 参与者的 )每个参与者都拥有一个完全缩减的 1/P 分片。

该二进制文件携带独立的手动调谐减少散射发射器。减少分散是 通过**AllReduce = ReduceScatter + AllGather**分解实现 TensorCore 集体堆栈围绕:专用 HLO 通行证(AllReduceReduceScatterReorderReduceScatterReassociateReduceScatterLegalizer)正是在这个边界上运行,并且 SparseCore-offload 基板构建了 ReduceScatterOffloadConfig,其 CollectiveIciStrategyConfig 是由 与 AllGather 和 AllReduce 相同的模板化环生成器。此页记录了戒指 归约分散循环,RS+AG = AR 分解恒等式,每步归约,以及 RS专用SC卸载路径;它链接了 AllGather 一半、分层 AllReduce 和 配置生成器而不是复制它们。

在二进制文件中观察到的ReduceScatter合约:

  • RS环是单向的,长度P(环参与者数),P-1步数。 每个 步骤精确传输一个 1/P 分片(= 结果形状大小),将其减少到本地 累加器,并转发。每步的传输量是输出形状,而不是操作数 形状 - 操作数是完整的预散射张量。
  • ReduceScatter 是 AllReduce 的归约阶段。 AllReduce(x) = AllGather(ReduceScatter(x)): RS 生成每个参与者的缩减分片,AllGather 将分片连接回完整的分片 复制张量。编译器将 RS 和 AR 视为一个操作的两个方面 - 专用 重新排序/重新关联/合法化 HLO 通行证在此边界上运行。 密集全面降低成本 支路充电B = 2 · operand_size(RS相1个operand_size + AG相1个 相),在活动环面轴上划分;收取SC-offload减少分散成本 分别在配置生成器发出的每个颜色环阶段 (§3)。
  • SC卸载ReduceScatter路径固定为平坦。 ConstructConfigForReduceScatterUniDirND @0x133ccbe0 使用 HierarchicalKind = 0x100 调用模板化构建器(engged + false = 显式平坦) - RS 在每个环面轴上获得一个显式相邻环,而不是多相 层次分解,在此构建中只有 AllReduce 可以采用。
  • RS 针对 1D/2D/3D ND 平面进行验证。 包装器拒绝任何其他内容 ("We only support 1D/2D/3D ReduceScatter for now.") 在到达环构建器之前。

概览

方面源码
HLO 操作码reduce-scatter = 93成本模型跳转表@0x130abfc0(v96 == 93)
分解标识AllReduce = AllGather ∘ ReduceScatterAllReduceReduceScatterReorder / ReduceScatterReassociate / ReduceScatterLegalizer 通过(纳米确认);密集AR分支收费B = 2·operand_size
成本分支SC卸载环电荷(不是密集的AR分支)GetCollectiveCycles 第二个开关 case 93 @ 线 902;操作码 93 在第一个(密集)开关中缺失,并落至 goto LABEL_447
成本分子v92 = ShapeSize(operand0)(v96 == 93 分支读取操作数 0),在发出的每色环相位上充电GetCollectiveCycles(v92 = GetShapeSize(v93+88)@第868行,v93 = operand(a2,0)@第853行)
环形单向,长度PP-1步骤,1/P分片/步骤AR系列分解 + SC IciStrategyRingType UNIDIR
SC卸载构建器ConstructConfigForCollectiveUniDirNDGroups<ReduceScatterOffloadConfig,…> @0x133cd800nm确认符号
SC卸载ND包装器ConstructConfigForReduceScatterUniDirND @0x133ccbe0使用 HierarchicalKind = 256 (FLAT) 调用构建器
SC卸载分相固定平面 (0x100) — 从不分层包装 …UniDirNDGroups<…>(…, 256, …) @ 第 80 行
SC后端配置ReduceScatterOffloadConfig(0x48的大小,vtable 0x21ce1c60)ctor @0x1d6eebe0;成本探针has_ici_strategy_config(hasbit 0x2)
ND 平面验证仅 1D / 2D / 3DGetCollectiveNDPlaneDimensionCount; ValidateReduceScatterReplicaGroupsOrderOnNDPlane @0x133cce40

1.环形减散环

ReduceScatter 是规范的 reduce-in-place 环。考虑一圈 P 参与者 r0, r1, …, r(P-1) 单向连接(r_i → r_{i+1 mod P}),各自持有一份本地副本 完整操作数 r0, r1, …, r(P-1) 的 ,概念上分为 P 分片 x_i[0], x_i[1], …, x_i[P-1]。 缩减分散运行 P-1 步骤。在步骤 t (0 ≤ t < P-1),参与者 r_i

  1. 从其上游邻居 r_{i-1} 接收一个分片 - 的部分累积 分片索引 (i - t - 1) mod P
  2. 就地减少:应用集体的减少操作(总和,对于嵌入梯度 case) 到它自己的同一分片索引的副本中,将邻居的部分总和累加到 本地一号。
  3. 转发现在进一步积累的分片到下游邻居 r_{i+1} 下一步。
text
ReduceScatter over a unidirectional ring of P participants (reduction op = ⊕)
each participant starts with the FULL operand, partitioned into P shards [0..P-1]

step t          r_i receives shard s = (i - t - 1) mod P from r_{i-1}
                acc_i[s] ⊕= recv                       (REDUCE IN PLACE)
                send acc_i[s] to r_{i+1}               (FORWARD)

after P-1 steps  r_i holds the fully-reduced shard index (i + 1) mod P :
                 acc_i[(i+1) mod P] = ⊕_{j=0..P-1} x_j[(i+1) mod P]

  invariants
    • per-step transfer volume = ONE shard = (operand_size / P) = OUTPUT-shape size
    • total bytes moved per participant = (P-1)/P · operand_size  ≈ operand_size
    • each participant ends owning exactly ONE 1/P shard of the reduced result
```text

RS 与 all-gather 的区别特征是每次都进行 **就地归约**
步骤(步骤 2):接收到的分片不仅仅被暂存到不断增长的缓冲区中(即全部收集)
但通过归约操作折叠到本地累加器中。这就是成本模型处理的原因
RS 阶段和 AG 阶段作为**相同的传输量** — 均移动 `(P-1)/P · operand_size` 环上的
字节 — 但只有 RS 对每个登陆的分片执行缩减算法。

在 TPU 环面上,环放置在环面轴 (X/Y/Z) 上。对于多维减少分散
环运行**每个活动环面轴**(副本组跨度的每个维度一个环) - 相同
每轴结构密集全降低成本分支数与其`num_dims = popcnt(active axes)`
术语 (`§2.1`)。每轴环方向为
单向 — SC 卸载基板明确命名它 (`ICI_RING_TYPE_UNIDIR_CW` /
`ICI_RING_TYPE_UNIDIR_CCW`),并且致密基板通过以下方式运行相同的每色环分解
`StrategyND` 的 `UniDirectionNDRingStrategy`。

> RS 操作码 (`93`) 及其形状处理以字节锚定在
> `CostModel::GetCollectiveCycles` @`0x130abfc0`:`v96 = *((_BYTE *)a2 + 12)`(操作码),
> `if (v96 == 93)` 分支@第859行读取**操作数0的**形状(`v154 = *(v93+88)`,
> `v93 = operand(a2,0)` @ 第 853 行; `v92 = GetShapeSize(v154)` @ 第 868 行) 为成本分子,
> 然后 `goto LABEL_113`。 `v96 == 9`(全归约)分支改为读取指令自己的
> 结果形状位于 `a2+88` (`*((_QWORD *)a2 + 11)`) — 对于 all-reduce 来说等于操作数。的
> 单向每轴环分解由 SC 卸载环配置 (`§4`) 承载。
> **[LOW]** 确切的 `(i - t - 1) mod P` 分片索引计划和 `P-1` 步数是
> 标准环减少分散算法 — *结构上*得到确认(UNIDIR 环 + 每个分片
> 输出大小传输 + `2·operand_size` AR 系列成本),但精确的每步分片
> 旋转索引不会作为此处跟踪的成本路径中的常量单独发出;是的
> 由每颜色 `RingLocation` 邻居调度携带(请参阅 [路由](../routing/overview.md))。

---

## 2. AllReduce = ReduceScatter + AllGather 恒等式

集体堆栈是围绕教科书带宽最优全归约分解构建的:

```text
AllReduce(x)  ≡  AllGather( ReduceScatter(x) )

  ReduceScatter :  P participants, each holds full x  ──►  each holds reduced shard s_i = ⊕_j x_j[i]
  AllGather     :  each holds shard s_i               ──►  each holds the concatenation ⊕_j x_j (full)

  total ring traffic per participant = 2 · (P-1)/P · operand_size   ≈  2 · operand_size
                                       └── RS phase ──┘ └── AG phase ──┘

该标识不仅仅是描述性的 — 它在二进制文件中编码在三个位置

2.1 密集全归约分支编码 B = 2 · operand_size

all-reduce (9) 和 all-reduce-start (11) 共享一个跳转表分支 (case 9:/case 11:) 位于 GetCollectiveCycles第一个(密集)开关中。计算出的带宽项为“B = 2· operand_size,RS阶段体积和AG阶段体积的代数和——成本模型 的AllReduce = AllGather ∘ ReduceScatter`分解表达式:

text
GetCollectiveCycles @0x130abfc0  (dense all-reduce branch, case 9/11 — NOT reduce-scatter)
    ShapeSize = GetShapeSize(operand0);   // @line 582 (case 9/11), operand 0 size
    v460 = 2 * ShapeSize;                 // @line 674: B = 2 · operand_size  (RS phase + AG phase)
    ...
    ShapeSize *= 2;                       // @line 707: the doubling preserved through the ND path
    __popcnt((unsigned __int8)v465 & 7);  // @line 727: num_dims = popcount of active axes (mask 0x7 = X/Y/Z)
    // cycles ≈ B / (num_dims · eff_Bps)
```text

减少分散(操作码 `93`)**不**进入此分支:第一个开关中不存在 `93`
并将落至其 `goto LABEL_447` 出口。 RS 仅作为 SparseCore 卸载才达到成本模型
指令,通过*第二个*开关(`case 93:`,`§3`),其中成本分子为
`v92 = ShapeSize(operand0)` 在发射的每色环相位上充电 - 没有 `2·` 因子 RS 路径上的
。 AR 路径上的 `2·` 因子仍然是
分解:全归约为归约阶段支付一个操作数大小的环流量,并为归约阶段支付一个操作数大小的环流量
收集阶段,分布在 `num_dims` 活动环面轴上。 **无添加剂 AR 分支中的
延迟术语** — 它是纯带宽(与概述一致)
“任何集合分支中都没有附加延迟项”)。完整的每类公式和 ICI
资源槽存款位于 [SPMD 链路计数成本](spmd-link-count-cost.md)。

### 2.2 HLO 跨界传递重新排序/合法化

编译器携带专用的 HLO 通道,这些通道恰好在此分解边界上运行,
确认堆栈将 RS 和 AR 视为一个操作的两个面:

- **`AllReduceReduceScatterReorder`** (`HloPassFix` @`0x109611a0` / `0x10961100`) — 重新订购
`all-reduce` 后跟一个切片到 `reduce-scatter` (规范的重新关联,将
`AllReduce` 然后每个分片使用更便宜的 `ReduceScatter`)。
- **`ReduceScatterReassociate`** (`HloPassFix` @`0x109603c0` / `0x10960320`) — 提升机减少分散
通过联想运算。
- **`ReduceScatterLegalizer`**(管道传递@`0x10969540`)-水母后端合法化器
将 `reduce-scatter` 降低为 TensorCore 路径消耗的环可发射形式。

### 2.3 SC-offload 后端配置在结构上与 AllReduce/AllGather 相同

三个卸载后端配置消息 — `AllGatherOffloadConfig`、`AllReduceOffloadConfig`、
`ReduceScatterOffloadConfig` — **在布局中字节相同**(`0x48` 的大小,从
相同的字段集),仅在其 vtable 上有所不同(`0x21ce1ce0` / `0x21ce1ca0` / **`0x21ce1c60`** 为
RS)和类型信息(RS 为 `0x21ce6a90`)。所有三个都嵌套相同的 `ici_strategy_config :
CollectiveIciStrategyConfig` (field 2) carrying the same per-color `phase_rings:
IciStrategyRingConfig`树。结构同一性是原型水平的表达
RS/AG/AR 系列共享一环表示。完整布局已开启
[SC-卸载配置生成器](sc-offload-config-builder.md)。

> `2·` 带宽因子(`v460 = 2 * ShapeSize` @ 第 674 行,则 `ShapeSize *= 2`
> @第707行)和`__popcnt(… & 7)` num-dims术语@第727行是字节锚定的
> `GetCollectiveCycles` @`0x130abfc0` — 在密集 `case 9:`/`case 11:` (全归约)分支中,不在
> 减少散射路径。 `AllReduceReduceScatterReorder` / `ReduceScatterReassociate` /
> `ReduceScatterLegalizer` 通证是 nm 确认的符号。字节相同
> `ReduceScatterOffloadConfig` 布局(`0x48` 的大小,vtable `0x21ce1c60`)通过其 ctor 确认
> @`0x1d6eebe0`。

---

## 3. 每步减少和成本形状

使 RS *reduce* 而不是 *gather* 的算术,是在每个落地 shard(`§1` 步骤 2)处进行就地折叠。逐元素 reduction op 运行在 TensorCore datapath 上,而不是作为离散成本项出现;成本模型只统计传输量。`GetCollectiveCycles` 中的 size-derivation `switch` 由三个 SC-offload collective 共享(opcodes `93`、`9`、`6`),reduce-scatter 采用其中最简单的分支:

```text
GetCollectiveCycles @0x130abfc0  (size derivation, lines 853-888)
    v93 = HloInstruction::operand(a2, 0)                   // operand 0          (line 853)
    v94 = v95 = GetShapeSize(operand0)                     // full pre-scatter size (line 854)
    v96 = opcode                                           // a2+12              (line 858)
    if (v96 == 93)       v92 = GetShapeSize(v93+88) ──────►// RS: OPERAND 0 size  (line 868) ; goto LABEL_113
    else if (v96 == 9)   v92 = GetShapeSize(a2+88)  ──────►// AR: result shape (== operand) ; goto LABEL_113
    else if (v96 != 6)   FATAL("Unsupported collective opcode")
    // opcode 6 = all-gather only — falls through to the shard-fraction path:
    v97 = GetShapeSize(a2+88)                              // AG: result shape   (line 883)
    v98 = v97 / v95                                        // shard fraction = result / operand
    v92 = GetShapeSize(operand0) · (v98 - 1)               // AG ring traffic    (line 888)

v96 == 93 分支读取操作数 0 的形状(v93+88v93 = operand(a2,0))并跳转 直接到 LABEL_113 — 减少分散的成本分子是完整的操作数大小, 稍后的循环对发出的环相位进行划分。 (All-reduce 读取 a2 自己的结果形状 a2+88;对于等于操作数的 AR。)operand0_size · (v98 − 1) 分片分数形式是 全聚集腿(操作码 6),仅当 939 都不匹配时才到达 - 它在 减少散射路径。

对于 SparseCore-offload 路径,成本模型探测 GetCollectiveOffloadConfig @0x133e1740,对于操作码 93,取消引用 reduce_scatter_offload_config() 并断言 has_ici_strategy_config()(配置消息上有 0x2):

text
GetCollectiveCycles @0x130abfc0  (SC-offload probe, case 93, lines ~890-908)
    GetCollectiveOffloadConfig(&cfg, a2)
    CHECK(collective_offload_config != nullopt)            // line 688: "!= std::nullopt"
    switch (opcode) {
      case 93:                                              // reduce-scatter
        cfg = cfg ?: &ReduceScatterOffloadConfig_globals_
        if ((cfg[2] & 2) == 0)                              // has_ici_strategy_config hasbit
          CHECK("…reduce_scatter_offload_config().has_ici_strategy_config()")
        break;                                              // charge the SC ring operating point
      case 9: … // all-reduce reads AllReduceOffloadConfig
    }
```text

当 SC 配置存在时,生成器发出的每颜色 UNIDIR 环集的成本费用,
不是密集 TC 工作点 — 与 SC 概述中描述的相同探头和充电
基板。

> RS尺寸分支(`v96 == 93`→`GetShapeSize(operand0)`,第868行,跳转
> 直接到 `LABEL_113`) 字节锚定在 `GetCollectiveCycles` @`0x130abfc0` 第 859/868 行;
> `operand0_size · (v98 - 1)` 分片分数形式 @ 第 888 行属于 **全聚集**
> (操作码 `6`)腿,不是 RS。 `case 93:` SC 卸载探头读数 `ReduceScatterOffloadConfig`
> 与 `& 2` (`has_ici_strategy_config`) 的 hasbit 位于反编译行 902(源 `cost_model.cc`
> 702 行)。

---

## 4. SC卸载ReduceScatter路径(固定平面)

SparseCore 卸载基板通过**相同的模板构建 `ReduceScatterOffloadConfig`
构建器** 作为 AllGather 和 AllReduce —
`ConstructConfigForCollectiveUniDirNDGroups<ReduceScatterOffloadConfig, HloReduceScatterInstruction>`
@`0x133cd800` — 由公共 ND 包装器 `ConstructConfigForReduceScatterUniDirND` 驱动
@`0x133ccbe0`。包装器的工作是验证指令,然后使用
**已修复 `HierarchicalKind`**。

### 4.1 包装器:验证然后平固定构建器调用

```text
ConstructConfigForReduceScatterUniDirND @0x133ccbe0
  1. RetCheck(reduce_scatter != nullptr)                              // line 1832
  2. IsSupportedReduceScatter(target, hlo)         else bail          // line 1833
  3. GetCollectiveNDPlaneDimensionCount(...) == 1   (must be a single ND-plane)
       else bail (line 1837)
  4. dim_count ∈ {1, 2, 3}   else RetCheck                            // line 1838
       "We only support 1D/2D/3D ReduceScatter for now."
  5. ValidateReduceScatterReplicaGroupsOrderOnNDPlane(...)  @0x133cce40   else bail (line 1843)
  6. ConstructConfigForCollectiveUniDirNDGroups<ReduceScatterOffloadConfig,…>(
         …, /* HierarchicalKind */ 256 /* = 0x100 = engaged+false = FLAT */, …)

第 6 步中的文字 256 (= 0x100) 是 RS 被 固定平坦 的字节锚定证明: HierarchicalKind & 0x101 == 0x100 为“接合 + 假” — 明确为单相扁平环 路径。因此,构建器的扁平与分层调度 (and $0x101 ; cmp $0x100) 始终 选择 FLAT 来减少分散。 RS永远无法达到多阶段分层双端队列行走; 仅 AllReduce 可以,并且仅当 xla_tpu_enable_sparse_core_hierarchical_all_reduce 启用时

  • 正确。 HierarchicalKind 包装记录在 分层种类 上。

4.2 扁平RS生成器发出什么

因为 RS 是平坦的,所以构建器会发出 每个环面轴一个 EXPLICIT 邻居环(无 D2D 片内多相分割超出了所有三种共享的兆核门控 0 相环)。 发射的 IciStrategyRingConfig 每轴携带:

字段(偏移)RS平值含义
ring_neighbor (0x24)ICI_RING_NEIGHBOR_EXPLICIT = 1通过预先计算的表进行邻居调度
core_count (0x18)计算环长度每轴环参与者计数
ring_neighbor_table_offset (0x28)来自 RingConfigAttributes[ringDim]邻居重排序表中的偏移量
has_reordering_map (0x3d)来自 RingConfigAttributes[ringDim]是否适用重新排序映射
explicit_strategy_ring_dim (0x48)每轴 ringDim该环经过哪个圆环轴
partner_transfers_outside_the_ring (0x3e)false

构建者共享相同的主体,双端队列-tuple<IciStrategyRingDim,long,long>,扭曲环面 门,以及带有 AllGather/AllReduce 的每颜色附加器 lambda(RS 的 @0x133e0c00) 实例化。 RS实例化仅在其构造的*OffloadConfig类型上有所不同 (vtable 0x21ce1c60)并且只能通过平销包装器到达。完整的 场图、每颜色发射环和 GetDimensionRings 每轴分区器已打开 SC-卸载配置生成器

ConstructConfigForReduceScatterUniDirND @0x133ccbe0 调用模板化 构建器,带有文字 256 (FLAT) — 字节锚定在反编译的调用站点。 1D/2D/3D 门("We only support 1D/2D/3D ReduceScatter for now."),IsSupportedReduceScatterGetCollectiveNDPlaneDimensionCountValidateReduceScatterReplicaGroupsOrderOnNDPlane @0x133cce40 都存在于包装中。 RS 构建者 @0x133cd800 已得到 nm 确认并且 共享 AllReduce 构建器主体(相同的 HierarchicalKind 参数、双端队列、扭曲环面路径、 across_cores_on_chip D2D 环)。


5. 与集体堆栈其余部分的关系

组件本页同级页面
减少分散环循环 + 每步减少此处(§1,§3)
AllReduce = RS + AG 分解恒等式此处 (§2)
全聚半(GetShardIndex/GetOffset,2D/3D环)已链接AllGather ND 环
分层多相AllReduce(0x101路径,钳形融合)已链接AllReduce 分层/钳形
延迟受限的 AR 发射器(二项式/递归加倍)已链接二项式/递归加倍
SC 卸载配置构建器(模板化环构建器 + 原型)已链接SC-卸载配置生成器
HierarchicalKind 扁平与分层标志已链接分层种类
密集ICI全约简原语已链接ICI 全归约原语
每种成本公式+ICI资源槽已链接SPMD 链路计数成本

ReduceScatter 位于 all-reduce 系列的中心:它是输出的 reduce 一半 对 AllGather 一半进行分片,然后重新连接。在密集的 TensorCore 基底上,它被选择并 由共享 StrategyND 机械发出(参见 SelectND策略);于 SparseCore 卸载基板,它是共享配置生成器的扁平引脚实例化。要么 减少分散原语本身的方式 - 接收分片,就地减少,转发 - 是 不变式,周围的布线、扭曲和屏障子系统围绕其进行调度。


交叉引用

全归约系列

SparseCore 卸载基板

  • SC-卸载配置生成器ConstructConfigForCollectiveUniDirNDGroups<*>*OffloadConfig 原型 / IciStrategyRingConfig 字段图
  • 分层种类AutoOr<bool> 平面 (0x100) 与分层 (0x101) 相分离

成本和选择

剖面图