ReduceScatter
二进制:
extracted/libtpu-0.0.40-cp314-cp314-manylinux_2_31_x86_64/libtpu/libtpu.so(构建 ID89edbbe81c5b328a958fe628a9f2207d,构建libtpu_lts_20260413_b_RC00;.textVMA == 文件偏移0xe63c000)。 状态: 重新实现级 · 证据等级: 已确认(字节锚定)——RS 环路/成本分解和 SC 卸载扁平引脚与 IDA 反编译进行交叉检查 · 第 XIII 部分 — On-Pod 集体和障碍 /集体算法 · 返回索引
摘要
ReduceScatter 是执行归约的全归约的一半:它对每个进行求和 参与者的贡献,让每个参与者持有不同的“碎片” 减少结果。在物理环面上,这被实现为一个单向环,其中 每一步都会从上游邻居接收一个分片,将其就地减少到本地 累加器,并将结果转发给下游邻居。 P-1 步骤后(对于环 P-1 参与者的 )每个参与者都拥有一个完全缩减的 1/P 分片。
该二进制文件不携带独立的手动调谐减少散射发射器。减少分散是 通过**AllReduce = ReduceScatter + AllGather**分解实现 TensorCore 集体堆栈围绕:专用 HLO 通行证(AllReduceReduceScatterReorder、 ReduceScatterReassociate、ReduceScatterLegalizer)正是在这个边界上运行,并且 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 ∘ ReduceScatter | AllReduceReduceScatterReorder / 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行) |
| 环形 | 单向,长度P,P-1步骤,1/P分片/步骤 | AR系列分解 + SC IciStrategyRingType UNIDIR |
| SC卸载构建器 | ConstructConfigForCollectiveUniDirNDGroups<ReduceScatterOffloadConfig,…> @0x133cd800 | nm确认符号 |
| 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 / 3D | GetCollectiveNDPlaneDimensionCount; 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:
- 从其上游邻居
r_{i-1}接收一个分片 - 的部分累积 分片索引(i - t - 1) mod P。 - 就地减少:应用集体的减少操作(总和,对于嵌入梯度 case) 到它自己的同一分片索引的副本中,将邻居的部分总和累加到 本地一号。
- 转发现在进一步积累的分片到下游邻居
r_{i+1}下一步。
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`分解表达式:
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+88、v93 = operand(a2,0))并跳转 直接到 LABEL_113 — 减少分散的成本分子是完整的操作数大小, 稍后的循环对发出的环相位进行划分。 (All-reduce 读取 a2 自己的结果形状 a2+88;对于等于操作数的 AR。)operand0_size · (v98 − 1) 分片分数形式是 全聚集腿(操作码 6),仅当 93 和 9 都不匹配时才到达 - 它不在 减少散射路径。
对于 SparseCore-offload 路径,成本模型探测 GetCollectiveOffloadConfig @0x133e1740,对于操作码 93,取消引用 reduce_scatter_offload_config() 并断言 has_ici_strategy_config()(配置消息上有 0x2):
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."),IsSupportedReduceScatter,GetCollectiveNDPlaneDimensionCount和ValidateReduceScatterReplicaGroupsOrderOnNDPlane@0x133cce40都存在于包装中。 RS 构建者 @0x133cd800已得到 nm 确认并且 共享 AllReduce 构建器主体(相同的HierarchicalKind参数、双端队列、扭曲环面路径、across_cores_on_chipD2D 环)。
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 卸载基板,它是共享配置生成器的扁平引脚实例化。要么 减少分散原语本身的方式 - 接收分片,就地减少,转发 - 是 不变式,周围的布线、扭曲和屏障子系统围绕其进行调度。
交叉引用
全归约系列
- AllGather ND 环 — 分解的全收集一半(
GetShardIndex/GetOffset,2D/3D 选择器) - AllReduce 分层/钳形 — 多相
0x101全减+钳融合 - 二项式/递归加倍 — 延迟受限的全缩减发射器
- ICI 全归约原语 — 织物上密集的 ICI 全还原
SparseCore 卸载基板
- SC-卸载配置生成器 —
ConstructConfigForCollectiveUniDirNDGroups<*>和*OffloadConfig原型 /IciStrategyRingConfig字段图 - 分层种类 —
AutoOr<bool>平面 (0x100) 与分层 (0x101) 相分离
成本和选择
- SPMD 链路计数成本 —
GetCollectiveCycles每种公式、2·operand_sizeAR 系列术语、ICI 资源槽 - SelectND策略 — 密集的 ND 策略选择器,可发出每色环分解
剖面图
- On-Pod 集体 — 剖面图 — 基板分割、端到端流程和操作系列调度
- 返回索引