我将基于 PyTorch 2.12.1 的 bug 修复内容，写一篇关于 GPU 浮点不确定性的深度教程博客。

一句话总结

PyTorch 2.12.1 修复了 Flash Attention 在 NVIDIA Blackwell GPU 上的不确定性 bug——这个 bug 是理解”GPU 并行化 + 浮点数学”如何产生不可复现结果的绝佳教材。

为什么这个 Bug 值得深究？

2.12.1 是个纯粹的 bug-fix 版本，更新日志很短，但它修复的两个问题都指向同一个 ML 工程师容易忽视的深层问题：

你的模型，真的每次跑出一样的结果吗？

在 NVIDIA B200（Blackwell 架构，sm100）上，Flash Attention 会产生不确定性输出——相同的输入、相同的权重，两次计算却得到不同的数值。更危险的是，这种差异足够小，容易被误认为”正常训练波动”而被忽视。

这不是孤立 bug。它是 GPU 并行化与浮点数学必然碰撞的产物。

---

浮点不确定性的根源：加法不满足结合律

先看一个让很多人惊讶的基础事实：

import torch

a = torch.tensor(1e8,  dtype=torch.float32)
b = torch.tensor(-1e8, dtype=torch.float32)
c = torch.tensor(1.0,  dtype=torch.float32)

print((a + b) + c)  # tensor(1.) ← 正确
print(a + (b + c))  # tensor(0.) ← 精度丢失！

这不是 bug，是 IEEE 754 浮点标准的固有性质。GPU 的并行 reduction（softmax 求和、layer norm 等）在不同 SM 配置下以不同顺序完成浮点累加，结果因此出现微妙差异。

Blackwell（B200/B100）引入了全新 SM 架构（sm100），Triton 编译出的线程块布局与旧架构不同，改变了 reduction 的累加顺序。这正是 2.12.1 通过升级 Triton 到 3.7.1 修复的根本原因。

---

Flash Attention 里的不确定性藏在哪？

Flash Attention 的核心是分块计算 attention，为避免将整个 attention 矩阵存入 HBM，它使用 online softmax：

\[\text{output}_i = \frac{\sum_j e^{s_{ij} - m} \cdot v_j}{\sum_j e^{s_{ij} - m}}, \quad m = \max_j s_{ij}\]

分块计算时，每个 tile 独立维护局部 max 和 sum，最后做跨 tile 的 reduction：

def flash_attn_sketch(Q, K, V, TILE=64):
    """
    Flash Attention 分块骨架（教学版，省略 padding mask 和 causal mask）
    不确定性来源：最后的 cross-tile parallel reduction
    """
    T = K.shape[-2]
    tile_results = []

    for start in range(0, T, TILE):
        K_tile = K[..., start:start+TILE, :]
        V_tile = V[..., start:start+TILE, :]

        scores     = Q @ K_tile.transpose(-2, -1)   # [B, H, Tq, TILE]
        local_max  = scores.amax(dim=-1, keepdim=True)
        exp_scores = (scores - local_max).exp()
        local_sum  = exp_scores.sum(dim=-1, keepdim=True)
        local_out  = exp_scores @ V_tile             # [B, H, Tq, d]

        tile_results.append((local_max, local_sum, local_out))

    # ← 不确定性就藏在这里
    # 在 Python 里这是串行循环，但在真实 Triton 内核里
    # 各 tile 的 warp 是并行跑的，reduction tree 的形状取决于 GPU 调度
    # Blackwell 的 warp 调度策略与 Hopper 不同 → 浮点累加顺序不同
    global_max = torch.stack([r[0] for r in tile_results]).max(0).values
    rescaled_sums = [r[1] * (r[0] - global_max).exp() for r in tile_results]
    global_sum    = torch.stack(rescaled_sums).sum(0)
    rescaled_outs = [r[2] * (r[0] - global_max).exp() for r in tile_results]
    output = torch.stack(rescaled_outs).sum(0) / global_sum

    return output

torch.stack(...).sum(0) 在真实 Triton 内核里是跨 warp 并行完成的，浮点求和顺序由 warp 调度决定。Blackwell 改变了这一调度策略，暴露了原本隐藏的精度差异。

---

如何检测你的模型是否存在批次不变性问题？

这个 bug 是被 test_batch_invariance 测试发现的，而非普通的确定性测试。批次不变性的含义是：把 batch 里的每个样本单独跑，结果应该与批处理完全一致。

import torch
import torch.nn as nn

def check_batch_invariance(model: nn.Module, x: torch.Tensor,
                            atol: float = 1e-5) -> bool:
    """
    测试批次不变性：逐样本处理结果是否与批处理一致？
    Flash Attention 的 Blackwell bug 正是通过这类测试发现的。
    """
    model.eval()
    with torch.no_grad():
        batch_out = model(x)

        individual_outs = [model(x[i:i+1]) for i in range(x.shape[0])]
        sequential_out  = torch.cat(individual_outs, dim=0)

    max_diff     = (batch_out - sequential_out).abs().max().item()
    is_invariant = torch.allclose(batch_out, sequential_out, atol=atol)

    print(f"最大差异: {max_diff:.2e}")
    print(f"批次不变性: {'✓ 通过' if is_invariant else '✗ 失败 — 存在不确定性！'}")
    return is_invariant


# 测试含 Flash Attention 的 Transformer
# PyTorch 2.x 在支持的 GPU 上默认启用 Flash Attention
model = nn.TransformerEncoder(
    nn.TransformerEncoderLayer(d_model=256, nhead=8, batch_first=True),
    num_layers=4
).cuda().eval()

x = torch.randn(4, 128, 256, device='cuda')
check_batch_invariance(model, x)

如果你使用 Blackwell GPU 且 PyTorch < 2.12.1，大概率会看到 ✗ 失败。

---

第二个 Bug：Triton 卷积核的非法内存访问

convolution2d_bwd_weight 在 sm100 上出现非法内存访问（illegal memory access）。这类 bug 的经典成因是 tile 边界假设失效：内核假设某个维度总是 16 的倍数（Ampere/Hopper 上通常如此），而 Blackwell 的线程映射改变了这一隐式约束。

Triton 内核中正确的防御性边界检查：

import triton
import triton.language as tl

@triton.jit
def safe_reduce_kernel(ptr, out_ptr, N: tl.constexpr, BLOCK: tl.constexpr):
    pid     = tl.program_id(0)
    offsets = pid * BLOCK + tl.arange(0, BLOCK)

    # 没有 mask：N % BLOCK != 0 时必然越界
    # x = tl.load(ptr + offsets)  ← 危险！

    # 有 mask：越界位置填 other=0.0，安全
    mask = offsets < N
    x    = tl.load(ptr + offsets, mask=mask, other=0.0)

    tl.store(out_ptr + pid, tl.sum(x, axis=0))

Blackwell 不再像旧架构那样”静默忽略”越界读取，而是直接报错崩溃。这实际上是一种进步——让潜伏多年的隐患浮出水面。

---

实践：强制使用确定性算法

如果你的实验对可复现性有严格要求（消融分析、对比实验）：

import torch

torch.use_deterministic_algorithms(True)   # 无确定性实现的算子会抛 RuntimeError
torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark    = False  # 禁止自动选择"最快但可能不确定"的算法

# 如果某个算子抛出 RuntimeError: no deterministic implementation
# 可对 attention 单独降级到数学实现（慢但确定）：
with torch.backends.cuda.sdp_kernel(enable_flash=False, enable_math=True):
    output = model(x)

代价是性能下降约 10-30%，在基准测试和可复现性实验中这个代价值得付。

---

适用边界

场景	建议
Blackwell GPU（B100/B200），任何用途	立即升级，之前结果可能有误
Hopper/Ampere + attention-heavy 模型	建议升级并运行批次不变性测试
训练卷积网络（conv2d backward 路径）	升级，排除内存安全隐患
纯推理，无 Blackwell GPU	低优先级，按常规节奏升级即可

我的观点

这个 release 的价值不在于修复本身，而在于它揭示的规律：新 GPU 架构是最好的压力测试，它暴露了上游编译器（Triton）在线程调度和边界处理上的隐式假设。

这些 bug 在 Ampere、Hopper 上”正常工作”，是因为那些架构恰好满足了代码的隐含约束，而不是代码本身是正确的。Blackwell 让沉默多年的技术债务被迫还清。

对工程师的实际启示：在新硬件上运行第一个实验之前，先跑批次不变性和确定性测试，再信任数字结果。 数值上的微小不确定性在训练中会被梯度放大，在长时间训练后可能导致不同的收敛路径，而这个差异极难在事后追溯。