Ch 09 · 批不变与确定性
封面/第二部分 · 基础设施/Ch 09
第二部分 · 基础设施 · 09

批不变与确定性 — 训推一致的护身符

V4 如何让同一 token 的输出不随 batch 位置改变，并让关键训练 kernel 在相同输入下确定性执行；重点包括 attention、GEMM、sparse-attention backward、MoE backward 与 mHC。

名词速通 · 一分钟看懂"批不变 + 确定性"

批不变 = 同一 token 单独跑、混在 batch 里跑、不同 batch size 跑，输出逐 bit 一致；确定性 = 同一份输入跑两次输出逐 bit 一致

一句话：浮点归约顺序会影响末位结果，而 batch 组织和并行调度可能改变这个顺序。V4 通过固定计算路径、隔离累加 buffer 和确定性归约，提高预训练、后训练与推理 kernel 的位级对齐程度。

浮点结合律不成立
$(a + b) + c \ne a + (b + c)$ 在 FP32 / BF16 下都常出现 1ULP 差异。原因：每次加法做 round-to-nearest，累计误差和加法顺序绑死。这是后面所有 patch 的根源：只要顺序变了，结果就变了。
批不变（Batch Invariance）
性质 A：对相同 token 和上下文，输出不随它在 batch 中的位置改变。报告将这一性质用于训练复现和跨流水对齐；prompt cache 是否命中通常由 prefix、模型版本和配置键决定。
确定性（Determinism）
性质 B：同样代码 + 同样输入 + 同样 seed，跑两遍输出字节相同。比批不变更弱（不要求 batch 变化下也一致），但更基础 —— 没有它连"重启复现"都做不到。
split-K（被放弃的优化）
大 GEMM 可把 reduction 维 $K$ 切成多块并行计算，再合并部分和。不同切法会改变归约顺序，因此常破坏 batch invariance。V4 在大多数场景放弃 split-K，并用其他优化弥补性能损失；报告没有给出 60%→90% 的固定利用率。
atomicAdd（破坏确定性的元凶）
多个 thread 并发加同一个地址，硬件保证加完不丢，但谁先加谁后加由调度决定。结果非确定。Sparse Attn 反向、MoE 反向都天然多 writer，原本最爱用它。
DeepGEMM（V4 自研 GEMM）
V4 用 DeepGEMM 端到端替代传统 cuBLAS 路径，并针对小 batch 场景补偿不使用 split-K 的损失。报告称其在多数主要场景可匹配或超过标准 split-K 实现，但没有把原因归结为一组固定技巧。
长序列单 SM 解码 + 末尾分布式补 wave
Attention 推理的批不变方案：长 prefix 解码固定一个 SM 串行算（顺序固定），末尾 token 用 distributed shared memory 多 SM 并行补齐 —— 第二个 kernel 仍然按确定顺序归约。
SM-local 累加 + 全局确定性归约
Sparse Attn 反向的方案：每个 SM 在自己的 shared mem 里累加自己负责的部分，不跨 SM atomicAdd；最后所有 SM 各自的部分按 SM id 顺序归约。顺序固定就字节固定。
mHC 反向的 split-K（不得不用）
mHC 的 expansion 矩阵输出维只有 24，K 维 7168 → 必须 split-K。V4 的妥协：每个 split 单独 output 一个 buffer，再用确定性 reduce kernel 求和（顺序按 split id）。
一句话定位："批不变 + 确定性"通过固定并行路径和归约顺序，换取更强的可复现调试与训推对齐。主要做法包括大多数 GEMM 避免 split-K、Sparse Attention 反向避免跨 SM atomicAdd，以及 mHC 使用独立部分和与确定性归约。

1. 一切问题的根源：浮点加法不结合

并行归约常会重排加法顺序，而 IEEE-754 浮点加法通常不满足严格结合律：每次运算都要舍入，因而不同的分组方式可能得到不同的末位结果。

📖 公式白话翻译 · 1ULP 差异是怎么来的

BF16 在 $[1,2)$ 区间相邻可表示数相差 $2^{-7}=0.0078125$。令 $x=3\cdot2^{-10}=0.0029296875$，比较 $1+x+x$ 的两种顺序：

  • 顺序 A：$(1+x)+x$。$x$ 小于半个 ULP，第一次相加舍入回 1；第二次仍舍入回 1。结果 = 1.0。
  • 顺序 B：$1+(x+x)$。$x+x=0.005859375$，超过 1 附近的半个 ULP，最终舍入到下一个 BF16 数。结果 = 1.0078125。

这个例子说明归约顺序可以改变结果，但误差不会普遍按 $\sqrt{N}$ 固定增长，也不能据此断言一定会改变采样 token 或 advantage 符号。确定性 kernel 的价值在于去掉这类随调度变化的数值分支，使异常更容易复查。

Demo · 同一组数，不同求和顺序，结果差几个 ULP
交互
输入数组（混合大小数，N=16） 顺序 A：左到右 sequential add 0 顺序 B：tree reduction（pairwise） 0 顺序 C：split-K 模拟（4 块并行 + 归约） 0 顺序 D：随机 shuffle 后求和 0 最大差异 = 0 ULP（FP32） 点"重抽数据"换一组随机大/小数；同一组里四种顺序都用 FP32 模拟。

读图法：四个求和都用 FP32（IEEE-754 单精度）。同一组 16 个数，A/B/C/D 仅顺序不同。"差异 ULP" = (max - min) / ULP(max)。具体差异由输入值和顺序共同决定。
把 N 加到 64 可以观察更长归约链，但差异不保证单调增大。这个示意解释了 split-K 改变归约分组后为何可能改变末位结果；V4 的相关 kernel 通过固定计算路径和确定性归约消除这类调度差异。

2. 批不变：Attention 与 GEMM 各下一刀

2.1 Attention：放弃 split-KV

推理时常见的优化是split-KV：把长 prefix 的 KV 维切成多段，每段一个 SM 算 attention，最后归约。问题：batch_size 不同 → SM 占用不同 → 切法不同 → 归约顺序不同 → 同一 token 的 logits 字节不一样。

V4 选择双 kernel 方案：

  • 长 prefix kernel：固定一个 SM 串行跑完整 prefix 的 attention。慢但顺序确定。
  • 末尾 wave kernel：剩下的几十个 token（远小于 SM 数）用 distributed shared memory 跨 SM 协同，但归约阶段按 SM id 排序，仍然确定。

固定并行与归约方式会限制部分调度自由。报告称批不变 kernel 仅带来很小开销，但没有给出“单 SM 慢 1.5×”或“整体慢约 10%”的数字。

2.2 GEMM：多数场景避免 split-K

报告指出传统 cuBLAS 不能满足 batch invariance，而小 batch GEMM 常借助 split-K 提高性能。V4 端到端使用 DeepGEMM，并在多数场景避免 split-K，再通过专门优化弥补性能损失。报告没有披露 $K\ge4K$ 的切换规则，也没有列出固定的三项实现技巧。

报告称这些优化使实现能在多数主要场景匹配或超过标准 split-K，同时保持 batch invariance，但没有公布统一速度区间。

3. 确定性：sparse attn 反向 / MoE 反向 / mHC 反向

3.1 Sparse Attn 反向：拒绝 atomicAdd

Sparse attention（CSA / HCA / Indexer）反向时，每个查询 token 的梯度要回写到它选中的那 $k$ 个 key 上。不同查询可能选中同一个 key → 多 writer → atomicAdd 是"自然"选择。但 atomicAdd 的顺序由调度决定，破坏确定性。

V4 的方案：SM-local 累加 + 全局确定性归约：

  1. 把 query token 按 SM id 划分；
  2. 每个 SM 在自己的 shared memory 里累加自己负责的查询对 keys 的贡献；
  3. 所有 SM 完成后，按 SM id 顺序把各 SM 的部分加到 global key buffer。

代价是显存占用增加（每个 SM 一份 key buffer 切片），但归约顺序由 SM id 决定，与调度无关，字节固定。

3.2 MoE 反向：token 排序预处理 + buffer 隔离

MoE 反向时，每个 expert 的梯度要回写给所有派发到它的 tokens。直接写会因为token 在 dispatch 时被打乱而非确定。V4 的做法：

  • 每个 rank 内：dispatch 阶段对落到该 rank 的 tokens 按 (expert_id, original_token_id) 排序，得到稳定顺序；
  • 每个 expert：单独的 grad buffer，写入位置由排序后下标决定；
  • 合并阶段：按排序顺序遍历 buffer 累加，无 atomicAdd。

3.3 mHC 反向：被迫 split-K，但自己写 reduce

mHC 的 expansion 矩阵 $W \in \mathbb{R}^{n \times r}$ 输出维 $r=24$，K 维 $n=7168$。这个 shape 极扁，不切 K 的话单 SM 算太久 GPU 闲。所以 mHC 不得不 split-K。妥协：

  • 每个 K-split 单独 output 到独立 buffer（不直接归约）；
  • 归约由专门的 deterministic reduce kernel 完成，按 split id 顺序累加；
  • 各部分和的缓冲区和归约顺序固定，从而避免原子加法带来的非确定到达顺序。
数值演练 · 这些 patch 的代价
  • Attention：固定并行方式可能降低部分 shape 下的峰值吞吐；
  • GEMM：避免依赖 batch size 改变 split-K 策略；
  • Sparse Attention / MoE / mHC backward：通过固定归约顺序和独立 buffer 换取确定性；
  • 报告将这些设计的整体开销描述为 minimal 或 negligible，没有逐项给出毫秒、显存或百分比。
这些设计的主要价值是复现与调试。不能由此推出 prompt cache “全命中”；缓存键通常由 token prefix、模型版本和配置决定。

4. 为什么这件事重要

主要工程价值
  1. 调试：相同输入和执行条件下得到相同结果，有助于复现硬件、软件故障和 loss spike，并缩小数值问题的排查范围。
  2. 流水对齐：预训练、后训练和推理采用位级对齐的关键 kernel，可减少实现差异带来的额外分布偏移；这不意味着所有训练与推理 logits 在任意配置下都自动相同。
  3. 缓存复用：确定性有助于验证复用结果一致，但 prompt cache 通常按 token prefix 和模型配置建立键。报告没有声称命中率提升 30%。

批不变与确定性 kernel 的直接价值是：相同输入在不同 batch 组织或重复运行时保持一致，便于定位训练异常并减少训推实现差异。它是一项基础设施能力，不需要用“护城河”来放大结论。

本章小结

  • 浮点加法通常不满足严格结合律；并行归约一旦改变加法分组，末位结果就可能变化。
  • 批不变 = 同 token 跨 batch_size 字节一致；确定性 = 同输入跨次运行字节一致。两者都不是免费的。
  • V4 在 Attention（双 kernel）、GEMM（DeepGEMM 不 split-K）、Sparse Attn 反向（SM-local + 顺序 reduce）、MoE 反向（排序 + buffer）、mHC 反向（split-K + 自定 reduce）逐项守住。
  • 报告称整体额外开销很小；收益主要是可复现调试与更严格的训推一致性。