Skip to content

低精度/量化集合

此页面上的所有地址、符号、偏移量和 .rodata 常量适用于 libtpu-0.0.40-cp314 轮中的 libtpu.so(build-id 89edbbe81c5b328a958fe628a9f2207d,781,691,048 字节,未剥离,.text VA == 文件偏移)。其他版本会有所不同;将每个 VA 视为版本固定。

摘要

本页记录了 TPU 集体的 缩减精度 - libtpu 更改用于“缩减”集体的数字格式的两种不同的独立机制,以及二进制文件实际包含内容的诚实存在/不存在清单。这两种机制很容易混淆,但它们是通过不同的通道、在不同的基材上、具有不同的目标而发射的:

  • BF16 累积精度提升(upcast)。 SPMD 分区器可以将 BF16 all-reduce 重写为 convert(bf16→f32) → all-reduce-in-f32 → convert(f32→bf16),因此减少累积在 F32 而不是 BF16 中。这是自动分片-spmd.md引用的MayIncreaseBF16AllReduceAccumulationAccuracy门(0x127a22c0)。它会花费更多的线路带宽(F32 有效负载),但给出的总和更准确。它由 xla_tpu_spmd_f32_accum_for_bf16_ar 和伴随的最小子组大小标志进行门控。
  • 8 位在线量化全归约(向下转换)。与上述不同的是,libtpu 具有完整的 RotatedPincerQuantizedEmitter 系列,在线路写入之前将每个分片对称量化为 8 位,逐步减少,并在最后进行反量化 - 将数值范围换为 ICI 有效负载的 一半/四分之一。支持的 8 位格式为 S8F8E5M2F8E4M3B11FNUZ。这是由 xla_tpu_quantized_all_reduce_level 下的 TpuQuantizedAllReduceBackendConfigSetter HLO pass (0x11107b00) 和大小阈值选择的。

存在/不存在的库存是此页面的中心结果(参见§3)。头条事实:fp8 集体压缩 存在 - 但它是 仅全归约(二进制中没有量化的全收集/全对所有/归约分散符号),并且量化集正是 {S8, F8E5M2, F8E4M3B11FNUZ} - F8E4M3Fn 在集体量化器中不存在,并且没有零点(设备量化器严格)对称)。

在二进制中观察到的量子化集体表面的契约:

  • 两个精度旋钮是正交且独立门控的。 F32 累积提升在 SPMD 分区器中运行(主机端 HLO 重写,在降低之前); 8 位在线量化作为单独的 HLO 后端配置通道运行,并由 TensorCore 降低发射器实现。两者都不暗示对方。
  • 量化全归约在 F32 中进行归约,而不是在 8 位中进行。 8 位格式仅是 wire 格式。每个环步将两个分片反量化为 F32,合并到 F32,并在重新量化下一跳之前重新跟踪正在运行的 absmax — 整数/fp8 字节从不参与算术。
  • 设备量化器是对称的absmax,每个分片有一个比例,零点固定为0scale = absmax / qmax,其中 qmax ∈ {127.0 (S8), 57344.0 (F8E5M2), 30.0 (F8E4M3B11FNUZ)}
  • 压缩在硬件生成和形状上进行门控。 CanLowerToQuantizedAllReduce (0x13798420) 需要 TpuVersion ∈ {3, 4} (viperfish / Ghostlite) 和一个形状,其次要/第二次要尺寸是车道/子车道计数的倍数,其元素类型是 BF16 或 F32。

概览

方面BF16→F32累积提升8位在线量化全归约
方向upcast(更准确,更多带宽)向下转换(带宽较少,有损)
门功能MayIncreaseBF16AllReduceAccumulationAccuracy 0x127a22c0RotatedPincerQuantizedEmitter::CanLowerToQuantizedAllReduce 0x13798420
实现者SPMD 创建者包装器 $_0 0x127a4340(克隆 + set_element_type + CreateConvert)RotatedPincerQuantizedEmitter 系列 + pincer_utils
选择通行证TpuSpmdPartitioner::AllReduceAlongShardingDims 0x127a28c0TpuQuantizedAllReduceBackendConfigSetter::RunImpl 0x11107b00
在哪里运行主机HLO重写(预降)HLO 后端配置标签 → TensorCore 降低
格式F32累加,BF16端点S8(2)、F8E5M2(19)、F8E4M3B11FNUZ(23)
主标志xla_tpu_spmd_f32_accum_for_bf16_arxla_tpu_quantized_all_reduce_level
涵盖的集体all-reduce(BF16还原)仅限 all-reduce
零点不适用(格式更改)无(对称)

1. BF16→F32累加-精度提升

1.1 它是什么

默认 BF16 全归约 BF16 中的“传输”和“累积”:每个环步骤将接收到的 BF16 分片添加到 BF16 部分和中,因此舍入误差会在环上复合。提升将其替换为 F32 累加:操作数向上转换为 F32,F32 中的全归约和(归约计算本身被克隆并重新键入为 F32),结果转换回 BF16。电线有效负载加倍;总和更准确。

1.2 大门 — MayIncreaseBF16AllReduceAccumulationAccuracy

xla::jellyfish::(anonymous namespace)::MayIncreaseBF16AllReduceAccumulationAccuracy (0x127a22c0) 是从 TPU SPMD 分区器覆盖 AllReduceAlongShardingDims (0x127a28c0) 到达的;请参阅 TPU 覆盖 表。它需要 ObjectView<TpuCompilationEnvironment>SPMDCollectiveOpsCreator,并返回(可能已包装的)创建者

反编译是一种干净的双臂结构,以一个配置字节为关键:

c
// MayIncreaseBF16AllReduceAccumulationAccuracy — 0x127a22c0
if ( *(_BYTE *)(tpu_comp_env + 4368) )      // xla_tpu_spmd_f32_accum_for_bf16_ar
{
    // TRUE arm: build a NEW SPMDCollectiveOpsCreator whose all-reduce callback
    // is the $_0 wrapper (0x127a4340). The companion field at +4632 (the
    // min-subgroup-size threshold) is captured into the wrapper's closure.
    creator.all_reduce_cb = &__call_func<...::$_0>;       // 0x127a4340
}
else
{
    // FALSE arm: pass the original creator through unchanged (no promotion).
}
```text

`tpu_comp_env + 4368`处的字节是`xla_tpu_spmd_f32_accum_for_bf16_ar`标志; `+4632` 处捕获的四字是 `xla_tpu_spmd_f32_accum_for_bf16_ar_min_subgroup_size` 同伴。设置该标志后,原始创建者的 all-reduce 回调将被执行实际重写的 `$_0` 闭包替换。

### 1.3 重写 — `$_0` 包装器

`__call_func<…MayIncreaseBF16AllReduceAccumulationAccuracy…::$_0>` (`0x127a4340`) 是包装的创建者。它的签名是集体创建者的全归约槽——`HloInstruction*(SpmdBuilder*, HloInstruction* operand, HloComputation* reduction, CollectiveDeviceListBase const&, long)`。解码后的主体是教科书的 upcast-reduce-downcast:

```c
// $_0 — 0x127a4340  (the F32-accumulation wrapper)
//  [1] subgroup-size short-circuit:
if ( creator.num_devices() > 0 && subgroup_min < creator[+32] )   // line 66
    return original_creator(operand, reduction, ...);             // too small -> passthrough

//  [2] dtype short-circuit:
if ( element_type(operand) != BF16 /*16*/ || rank-flag == 55 )    // line 75
    return original_creator(operand, reduction, ...);             // not BF16 -> passthrough

//  [3] clone the reduction computation and re-type its BF16 nodes to F32:
clone = HloComputation::Clone("clone");                           // line 84
AddEmbeddedComputation(clone);
for each instr in clone:                                          // lines 90..218
    if element_type(instr) == BF16: set_element_type(instr, F32); // line 159-160
    if shape is TUPLE (prim 4): ForEachMutableSubshapeHelper(... promote BF16->F32);

//  [4] upcast operand, run the F32 all-reduce, downcast result:
f32_in  = CreateConvert(operand, ChangeElementType(shape, F32 /*11*/));   // line 222-223
f32_ar  = original_creator(f32_in, clone, ...);                          // line 227 (F32 reduce)
bf16_out= CreateConvert(f32_ar, ChangeElementType(shape, BF16 /*16*/));  // line 228-230
return bf16_out;

因此,提升不是硬件累加器宽度开关 - 它是 HLO 图重写,将全归约包装在 BF16↔F32 convert 操作中,并重新输入归约体。 PrimitiveType 使用的整数:BF16 = 16F32 = 11TUPLE = 4。两次短路(子组太小;操作数不是 BF16)使原始创建者保持不变。

注意 — 子组大小短路 (creator.num_devices() > 0 && subgroup_min < creator[+32]) 是 xla_tpu_spmd_f32_accum_for_bf16_ar_min_subgroup_size 的运行时使用:当 all-reduce 的子组低于阈值时,将跳过 F32 提升(小缩减不会积累足够的错误以值得双倍的有效负载)。阈值是否比较子组大小计数未进行位跟踪(低)。

1.4 在线简化型配套产品 — bf16_inside_cross_replica_sum

在降低时间存在第二个较低级别的决策:AllReduceEmitter::bf16_inside_cross_replica_sum (0x1373ca60) 决定 ICI 环减少是否在 BF16 与 F32 中“在线”运行,独立于上面的 SPMD 级别提升。

c
// AllReduceEmitter::bf16_inside_cross_replica_sum — 0x1373ca60
reduce_scatter = ExtractInstruction(fusion, {9, 11, 0x3d});      // AR / AR-start / reduce-scatter
CHECK(reduce_scatter != nullptr);                                // all_reduce_emitter.cc:585
if ( GetTpuCompEnv()[+3789] )    return true;                    // xla_jf_bf16_inside_cross_replica_sum
else                            return element_type(to_apply.root) == BF16 /*16*/;
```text

当 `xla_jf_bf16_inside_cross_replica_sum` 标志(`TpuCompEnv + 3789` 处的配置字节)置位时,强制进行 BF16 环内求和;否则它遵循归约计算的根元素类型。这是与 SPMD 级 F32 提升(第 1.3 节)相反的旋钮 — 提升的存在正是为了覆盖 F32 的默认值,以实现对精度敏感的减少。

---

## 2. 8位在线量化全归约

这是真正的**有损压缩**路径:它在传输之前将每个分片量化为 8 位格式,以将每个环步移动的 ICI 字节减半(或四分之一)。它是由 `RotatedPincerQuantizedEmitter`(双向 [钳](allreduce-hierarchical-pincer.md) 全归约的量化特化)和 `pincer_utils` 对称量化助手实现的。

### 2.1 选择 — `TpuQuantizedAllReduceBackendConfigSetter`

传递 `TpuQuantizedAllReduceBackendConfigSetter::RunImpl` (`0x11107b00`) 遍历每个非融合计算,找到 `all-reduce` 操作(操作码 `9`,反编译第 182 行),并使用 `QuantizedAllReduceConfig` 后端配置标记符合条件的操作。解码后的门序列:

```c
// TpuQuantizedAllReduceBackendConfigSetter::RunImpl — 0x11107b00
level = GetTpuCompEnv()[+5600];  if (level >= 4) level = 0;   // xla_tpu_quantized_all_reduce_level, clamp 0..3
thr_f = GetTpuCompEnv()[+0x15dc];                            // size threshold (MiB) -> elements via 0x84a2c94
thr_elems = vcvttss2si( 0x84a2c94 * thr_f * 0x84a2c94 );     // line 153-156
for each all-reduce (opcode 9):
    // (a) optional frontend-attribute override of dtype/stage:
    if frontend_attr[kQuantizeAllReduceDtypeFrontendAttribute] present:
        dtype = match{ "S8"->2, "F8E5M2"->19, "F8E4M3B11FNUZ"->23 }  // lines 241..327
                else FAIL "Unsupported quantized type: <name>"
    if frontend_attr[kQuantizeAllReduceStageFrontendAttribute] present:
        stage = match{ name -> 1 | 2 | 3 } else FAIL "Unsupported quantized stage: <name>"
    // (b) size gate (only when level != 0):
    if (level != 0 && GetSizeInBytes(operand) < thr_elems) continue;   // line 565-567
    // (c) emitter feasibility gate:
    if (!CanLowerToQuantizedAllReduce(instr, target))                  // line 581
        { VLOG "Not using quantized all-reduce for <hlo> ..."; continue; }
    // (d) write QuantizedAllReduceConfig { dtype, ?, stage } into backend config:
    cfg.dtype = dtype; cfg[+28] = ...; cfg.stage = stage;              // lines 710..735

因此,配置带有 量化 dtype阶段 枚举(QuantizedAllReduceStage,值 1/2/3)。该过程建立基线 {dtype=S8, stage=2} 配置(第 234 行 v140 = 0x200000001 = 具有默认 dtype 的打包 {1, 2} 字段对),然后从前端属性(如果存在)覆盖。两个故障字符串 — "Unsupported quantized type: ""Unsupported quantized stage: " — 字节锚定在 TU tpu_quantized_all_reduce_backend_config_setter.cc 中。

前端属性/配置字段源码解码值
量化数据类型kQuantizeAllReduceDtypeFrontendAttributeS8=2、F8E5M2=19、F8E4M3B11FNUZ=23
量化级kQuantizeAllReduceStageFrontendAttributeQuantizedAllReduceStage ∈ {1, 2, 3}
级别xla_tpu_quantized_all_reduce_level (TpuCompEnv+5600)钳位至 0..3
尺寸阈值xla_tpu_quantized_all_reduce_size_threshold_mib(TpuCompEnv+0x15dc)MiB → 元素
操作数组合xla_tpu_combine_quantized_all_reduce_operandsbool(默认生成存在)

2.2 可行性门 — CanLowerToQuantizedAllReduce

RotatedPincerQuantizedEmitter::CanLowerToQuantizedAllReduce (0x13798420) 是每操作 + 每目标检查。两个连词:

c
// CanLowerToQuantizedAllReduce — 0x13798420
if ( !IsShapeSupported(shape, target) )      // 0x13798560
    { VLOG "Minor/second-minor dimension is not a multiple of lane/sublane count"
          " or type is neither BF16 nor F32. Shape not supported ..."; return 0; }
if ( (unsigned)(TpuVersion(target) - 3) < 2 )  return 1;   // TpuVersion in {3,4}
VLOG "Target is not supported by RotatedPincerQuantizedEmitter. Target: <version>";
return 0;
```text

- **形状支持** (`IsShapeSupported` `0x13798560`):次要维度和第二次要维度必须是通道/子通道向量计数的倍数(以便分片干净地平铺),并且 *操作数* 元素类型必须是 **BF16 或 F32** (正在量化的值是 BF16/F32;8 位type 是有线格式,而不是操作数类型)。
- **目标生成**:`TpuVersion - 3 < 2`,即 `TpuVersion ∈ {3, 4}` = **viperfish / Ghostlite**。老一代被字节锚定的“不支持目标”VLOG(以及插入的 `TpuVersion` 枚举器)拒绝。

> **注意 —** `TpuVersion` 编号与 [集体概述](overview.md) §5 内部枚举 (`0 jellyfish, 1 dragonfish, 2 pufferfish, 3 viperfish, 4 ghostlite, 5 …`) 匹配。量化全归约准确地承认 `{viperfish, ghostlite}` — `Target::SupportsVectorConvertF32Stochastic` 集携带 `{F8E5M2, F8E4M3B11FNUZ}` 的相同两代。这种对齐方式(支持 SR 的 fp8 集 == AR 可量化的 fp8 集)是一致的,但在此并未证明因果关系(低)。

### 2.3 对称量子内核 — `pincer_utils`

实际的数字内核位于 `xla::jellyfish::pincer_utils` 中,是二进制文件中最干净的对称absmax-8 位引用。比例公式是字节精确的:

```c
// pincer_utils::UpdateScale — 0x137b75c0
//   scale_out = qmax(dtype) / absmax        (the *quant* multiplier reciprocal)
qmax = switch(dtype) {                        // .rodata f32 constants:
   S8           (2):  127.0    // dword_84a2a28
   F8E4M3B11FNUZ(0x17): 30.0   // dword_84a27fc
   F8E5M2       (0x13): 57344.0 // dword_84a2530
   default: FATAL "Unsupported quantized shard type: %s"   // pincer_utils.cc:200
};
m       = VimmF32(qmax);
absmax  = Vld(max_abs_addr);                  // both addrs CHECKed to be in VMEM
scale   = VdivF32(m, absmax);                 // scale = qmax / absmax
Vst(scale_addr, scale);

内核家族及其在一环步骤中的作用:

pincer_utils 内核地址角色
UpdateMaxLocalChunk0x137b73a0在分片上运行absmax:acc = max(acc, |x|)
UpdateScale0x137b75c0scale = qmax / absmax(上述每个 dtype 的 qmax)
SymmetricallyQuantizeShardInPlaceTo8Bits0x137b7740q = round(x · scale) 然后通道打包为 8 位
SymmetricallyDequantizeShardInPlace8Bit0x137b7fc0f = q / scale(解包8位→F32)
ReduceSymmetricallyQuantized8BitShardInPlace0x137b8880每步 F32 取消量化 → 合并 → 重新跟踪 absmax

决定性的结构事实,来自 ReduceSymmetricallyQuantized8BitShardInPlace:每步环缩减将本地分片和接收分片反量化为 F32,合并到 F32(缩减函子)中,然后重新跟踪absmax以进行下一跳的重新量化。 8 位整数/fp8 字节永远不会直接求和——它们纯粹是线表示。这就是为什么量化路径需要每个分片的规模与数据一起传送以及为什么它可以携带任意约简函子(sum / max / …)。

2.4 发射器 — RotatedPincerQuantizedEmitter

发射器是双向 钳形全减 的量化专业化:它在两个方向上运行带宽最佳旋转环,但在每次发送之前插入一个量化,在每次接收之后插入一个反量化。其恢复的表面:

方法地址角色
Init0x13797700分配刻度缓冲区,暂存
QuantizeShard0x1379c5e0线写前按分片量化
DequantizeShard0x1379c940每个分片接收后反量化
ComputeScaleValue / ComputeScaleFactor0x1379e0a0 / 0x137a2440每分片/每步规模
ReductionLoop0x13798ec0F32-合并环路
DequantAndReduceShardInPlace0x1379dd00反量化 + 合并融合
SendOrWaitForShards / WaitForShardScaleFactor0x1379e460 / 0x1379a500ICI握手(数据+规模)
SetSummationPrecision0x137a2400设置 LocalDmaPipe::Precision(代表基地)
CanLowerToQuantizedAllReduce0x13798420§2.2门

SetSummationPrecision (0x137a2400) 是 RotatedPincerEmitterBase::SetSummationPrecision 的精简版 — 确认求和精度(第 2.3 节的 F32 合并)是来自 8 位线格式的独立 LocalDmaPipe::Precision 枚举旋钮。量化/反量化/打包/解包叶操作(VcvtF32ToS32VpackcB8VunpackCS8 / VunpackCF8E5M2 / VunpackCF8E4M3B11VdivF32)是共享的 TensorCore 转换表面,并使用 matmul-epilogue 数字进行记录,而不是在此处记录。


3. 存在/不存在库存

这是该页面的主要交付内容。每行都是字节锚定的,并通过反编译中的符号扫描来确认(或驳斥)。

能力状态证据
BF16→F32全减累积提升现在MayIncreaseBF16AllReduceAccumulationAccuracy 0x127a22c0;包装器 $_0 0x127a4340(克隆 + set_element_type + CreateConvert BF16↔F32)
在线 BF16-vs-F32 减速拨动开关现在AllReduceEmitter::bf16_inside_cross_replica_sum 0x1373ca60(标志xla_jf_bf16_inside_cross_replica_sum)
8 位在线 量化全归约现在RotatedPincerQuantizedEmitter(全方法面); TpuQuantizedAllReduceBackendConfigSetter::RunImpl 0x11107b00; QuantizedAllReduceConfig原型
量化格式 S8F8E5M2F8E4M3B11FNUZ现在pincer_utils::UpdateScale 开关盒 {2, 0x13, 0x17}; setter dtype 匹配 {2, 19, 23}
对称绝对最大刻度 (qmax/absmax)现在UpdateScale 0x137b75c0; .rodata qmax {127.0, 57344.0, 30.0}
F32 飞行中还原(8 位 = 仅有线)现在ReduceSymmetricallyQuantized8BitShardInPlace 0x137b8880(反量化→F32合并→重新absmax)
量化格式 F8E4M3Fn (PrimitiveType 20)不存在不在 UpdateScale 开关中,也不在 setter dtype 匹配中(仅 {2,19,23})
零点 /非对称集体量化不存在设备量化器是对称的(仅比例); pincer_utils 中无零点场
量化全聚集不存在Quantized*Gather 符号(仅 QuantizedAllReduce*)
量化 全部到全部 / 减少分散不存在Quantized{AllToAll,ReduceScatter} 符号
TpuVersion ∉ {3,4} 上的量化 AR不存在(门控关闭)CanLowerToQuantizedAllReduce 0x13798420TpuVersion - 3 < 2

注意 - “fp8 量化集合”在这里的含义和不含义。 期望通用“fp8 压缩每个集合”功能的读者将找不到它:唯一具有 8 位在线路径的集合是 all-reduce,通过旋转钳子系列。全部收集、减少分散和全部到全部移动未压缩的数据。 F8E4M3Fn 格式(TPU 转换表面上其他地方更常见的 fp8 变体)不是可接受的集体量化类型;集体量化器仅允许 B11FNUZ fp8 变体加上 F8E5M2S8。而且量化器是对称的 - 无零点 - 因此它不适合非对称/无符号激活分布。

注意 — QuantizedAllReduceStage 枚举(值 1/2/3)几乎可以肯定区分多通道量化减少(例如减少分散阶段与全收集阶段与组合),匹配 钳融合 的可分离臂结构,但每级发射器发散并未展开(低)。确认 xla_tpu_combine_quantized_all_reduce_operands 标志(操作数批处理旋钮,默认生成存在)存在,但未跟踪其组合逻辑(低)。


4. 两种精密机构的关系

这两种机制位于精度/带宽交易的两端,并且在意图上是互斥的,尽管二进制文件中没有任何内容禁止设置这两个标志:

text
operand (BF16)

   ├── xla_tpu_spmd_f32_accum_for_bf16_ar (SPMD, §1)
   │        convert BF16->F32 ─► all-reduce in F32 ─► convert F32->BF16
   │        (MORE accurate, 2x wire bytes; subgroup-size gated)

   └── xla_tpu_quantized_all_reduce_level (backend-config, §2)
            quantize BF16->8bit ─► pincer ring (F32 merge per step) ─► dequant 8bit->BF16
            (LESS bandwidth, lossy; TpuVersion {3,4} + shape + size gated)
```text

累积提升是 SPMD 分区器用于*正确性敏感*BF16 减少的杠杆(在分片期间选择,其中 `GetCommunicationTimeInMilliSec` 占双倍的有效负载 - 请参阅 [SPMD 链路计数成本](spmd-link-count-cost.md))。量化路径是后端在最新硬件上减少“带宽限制”的杠杆。重新实现者必须将它们保留在单独的通道中:F32 提升是在降低之前发出的 HLO 图重写;量化是 TensorCore 降低发射器消耗的后端配置标签。

---

## 交叉引用

- [集体概述](overview.md) — 集体家族分类、基质分割和 `TpuVersion` 内部枚举。
- [AllReduce 分层/钳形](allreduce-hierarchical-pincer.md)——`RotatedPincerQuantizedEmitter`专业的双向钳系列;可分离的减少分散/全聚集臂。
- [二项式/递归加倍](binomial-recursive-doubling.md) — 延迟受限的全缩减发射器(非钳形拓扑,在此版本中不可量化)。
- [SPMD 链路计数成本](spmd-link-count-cost.md) — `GetCommunicationMultiplier` 以及 SPMD 分区器在选择 F32 累积重写时权衡的每种成本。
- [自动分片/SPMD](../compiler/auto-sharding-spmd.md) — `TpuSpmdPartitioner::AllReduceAlongShardingDims` 覆盖和此页面展开的 `MayIncreaseBF16AllReduceAccumulationAccuracy` 门。
- [返回索引](../index.md)