第三方并行库"静默算错"：一次 ring attention 的排查记录

背景

我最近在做一件很具体的事：把一个扩散模型（图像生成模型）的推理框架，从 CUDA 移植到另一种加速卡上。

这个模型是 20B 级别的 Transformer，出一次图要跑几十步迭代。序列长度也很夸张------图像不是按像素处理，而是先压到潜空间再切成小块，一张 1024 级别的图展开就是几千个 token。结果是两个问题一起压过来：

• 单卡太慢：出图要几十秒到几分钟
• 单卡装不下：模型权重加上中间激活，显存压力很大

解法是序列并行：把序列切成几段，分给多张卡同时算，算完再拼回来。

其中注意力那一段（序列并行里最关键的部分）我没有自己写，直接用了业界的开源实现------这也是常规做法。ring attention 涉及环形通信和流式 softmax，自己从零写很容易出错，用经过验证的实现本来是最省事的路径。

但接入之后，端到端结果不对。 而且诡异的是：不报错、不崩溃、日志干净、跑完 40 步正常出图------只是出的图跟预期完全是两张。

这篇文章记录我怎么把它揪出来的。核心结论是：这个库的注意力实现压根没做通信，它只算了本地那一段。

下文只谈技术，代码为示意。

---

一、这类 bug 为什么难发现

普通的 bug 会崩，会抛异常，至少给你一个堆栈。

但分布式并行库的 bug 往往长这样：

• 通信确实发生了，只是算的是错的东西
• 或者通信压根没发生，但代码路径看起来是通的
• 张量形状对、dtype 对、内存不越界
• 跑完 40 步，出图流程正常结束，日志干干净净

你唯一能察觉的是：结果不对。而"结果不对"这件事，在扩散模型里可以有一百种原因------采样随机性、精度差异、参考图过期、seed 不同......

所以第一步不是查库，是排除其他所有可能。

---

二、排查路径：从"怀疑自己"到"怀疑库"

我当时的顺序是这样的，后来证明这个顺序是对的------由外向内，一层层缩小范围。

第 1 步：确认单卡是对的

先把并行关掉，单卡跑一遍，和参考结果比。单卡能对上（或者差距在已知范围内），说明模型实现本身没问题，问题出在并行路径上。

这一步很重要，如果单卡就不对，那跟并行库没关系。

第 2 步：确认切分逻辑是对的

把"每张卡应该拿到哪一段数据"打印出来，人工核对。

具体做法：给输入序列打上标记（比如第 i 个位置填 i），然后看每张卡拿到的片段是不是连续的、有没有重叠或遗漏。

这一步能排除"分片算错了但注意力是对的"这种情况。

第 3 步：单层对比（关键一步）

不看整条链路，只拿出一个注意力层，喂小输入，把并行实现的输出和朴素实现的输出逐元素比。

# 朴素实现：完整序列一次算完
def naive_attention(q, k, v):
    scale = q.shape[-1] ** -0.5
    s = torch.matmul(q, k.transpose(-2, -1)) * scale
    p = torch.softmax(s, dim=-1)
    return torch.matmul(p, v)

判定标准：

对比对象	指标	正常值（bf16）	说明
vs 全序列参考	max_abs_err	1e-3 ~ 1e-2	舍入误差，正常
vs 全序列参考	余弦相似度	> 0.999	同上
vs 只用本地数据	max_abs_err	必须 > 0	等于 0 说明根本没通信

最后一行是最有价值的判据，下一节专门讲。

第 4 步：验证"到底有没有通信"

如果第 3 步发现输出和朴素实现差很多，先别急着改算法，先问一句：它真的通信了吗？

验证方法极简单但极有效：

# 并行实现的输出
out_sp = sp_attention(q_local, k_local, v_local)

# 只用本地分片算（完全不通信）
out_local = naive_attention(q_local, k_local, v_local)

print("max_err =", (out_sp - out_local).abs().max().item())

如果 max_err == 0，说明并行实现返回的就是"只看本地数据"的结果------它压根没把别的卡的数据算进去。

我第一次跑出这个 0 的时候，基本就确定问题在哪了。

---

三、验证脚本（可直接用）

下面这份脚本脱胎于当时的排查工具，做了通用化处理，不依赖特定硬件------把后端的 gloo 换成 nccl 或各家加速卡的通信后端就能直接跑。

"""
verify_sp_attention.py

用途：验证序列并行注意力实现的数值正确性。

用法：
    torchrun --nproc_per_node=2 verify_sp_attention.py
    torchrun --nproc_per_node=2 verify_sp_attention.py --backend nccl

会依次执行三项检查：
    [1] 分片正确性：每张卡拿到的序列片段是否连续无重叠
    [2] 数值正确性：并行输出 vs 朴素实现
    [3] 通信真伪：并行输出是否等价于"只用本地数据"
"""
import argparse

import torch
import torch.distributed as dist


def naive_attention(q, k, v):
    """参考实现：完整序列的朴素注意力。q/k/v: (B, S, H, D)"""
    scale = q.shape[-1] ** -0.5
    scores = torch.matmul(q, k.transpose(-2, -1)) * scale
    probs = torch.softmax(scores, dim=-1)
    return torch.matmul(probs, v)


def make_inputs(b, s, h, d, seed=0, device="cpu"):
    """同一个种子，保证每张卡造出一模一样的完整张量"""
    g = torch.Generator().manual_seed(seed)
    q = torch.randn(b, s, h, d, generator=g)
    k = torch.randn(b, s, h, d, generator=g)
    v = torch.randn(b, s, h, d, generator=g)
    return q.to(device), k.to(device), v.to(device)


def check_shard_cover(rank, world, seq_len):
    """[1] 每张卡分到的 [start, end) 是否连续、无重叠、无遗漏"""
    per = seq_len // world
    start, end = rank * per, (rank + 1) * per if rank < world - 1 else seq_len
    print(f"  rank{rank}: [{start}, {end})  len={end - start}")
    return start, end


def check_numeric(out_par, out_ref, tag):
    """[2] 数值对比"""
    a, b = out_par.float(), out_ref.float()
    max_err = (a - b).abs().max().item()
    cos = torch.nn.functional.cosine_similarity(
        a.flatten(), b.flatten(), dim=0
    ).item()
    ok = max_err < 1e-2 and cos > 0.999
    print(f"  {tag:<28} max_err={max_err:.6f}  cos={cos:.8f}  {'OK' if ok else 'FAIL'}")
    return ok


def check_did_communicate(out_par, out_local, tag):
    """[3] 并行输出如果和"只用本地数据"完全相同，说明没通信"""
    diff = (out_par.float() - out_local.float()).abs().max().item()
    printed = "完全相同 -> 疑似未通信" if diff == 0 else f"有差异({diff:.6f})"
    print(f"  {tag:<28} vs 本地结果: {printed}")
    return diff > 0


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--backend", default="gloo")
    ap.add_argument("--seq", type=int, default=64)
    ap.add_argument("--heads", type=int, default=2)
    ap.add_argument("--dim", type=int, default=32)
    args = ap.parse_args()

    dist.init_process_group(backend=args.backend)
    rank, world = dist.get_rank(), dist.get_world_size()
    device = torch.device("cpu")

    b, s, h, d = 1, args.seq, args.heads, args.dim
    q, k, v = make_inputs(b, s, h, d, device=device)
    ref = naive_attention(q, k, v)          # 全序列参考结果

    if rank == 0:
        print(f"\n[1] 分片覆盖检查 (world={world}, seq={s})")
    dist.barrier()
    start, end = check_shard_cover(rank, world, s)
    dist.barrier()

    # 本卡分片
    q_l = q[:, start:end].contiguous()
    k_l = k[:, start:end].contiguous()
    v_l = v[:, start:end].contiguous()

    # ---- 这里替换成你正在验证的那个库的实现 ----
    # 约定：输入本卡的 q/k/v 分片，返回本卡分片对应的输出 (B, S_local, H, D)
    from your_lib import ring_attention          # noqa: F401
    out_par = ring_attention(q_l, k_l, v_l, dist.group.WORLD)
    # ------------------------------------------

    out_ref = ref[:, start:end]             # 参考结果里对应的那一段
    out_local = naive_attention(q_l, k_l, v_l)   # 只用本地数据

    if rank == 0:
        print("\n[2] 数值正确性")
    dist.barrier()
    ok_num = check_numeric(out_par, out_ref, f"rank{rank} vs 全序列参考")
    dist.barrier()

    if rank == 0:
        print("\n[3] 通信真伪")
    dist.barrier()
    ok_comm = check_did_communicate(out_par, out_local, f"rank{rank} 并行 vs 本地")
    dist.barrier()

    dist.destroy_process_group()

    if rank == 0:
        print("\n结论：")
        print(f"  数值      {'通过' if ok_num else '失败'}")
        print(f"  通信      {'发生' if ok_comm else '未发生（几乎可以确定是 bug）'}")


if __name__ == "__main__":
    main()

怎么读这份脚本的输出

跑两卡，正常的情况应该是这个样子：

[1] 分片覆盖检查 (world=2, seq=64)
  rank0: [0, 32)  len=32
  rank1: [32, 64)  len=32
[2] 数值正确性
  rank0 vs 全序列参考        max_err=1e-3 量级   cos>0.999   OK
  rank1 vs 全序列参考        max_err=1e-3 量级   cos>0.999   OK
[3] 通信真伪
  rank0 并行 vs 本地          vs 本地结果: 有差异
  rank1 并行 vs 本地          vs 本地结果: 有差异

我当时跑出来是这样（只贴 rank0，最直接的证据在这里）：

[2] 数值正确性
  rank0 vs 全序列参考        FAIL（偏差远超 bf16 量级）
[3] 通信真伪
  rank0 并行 vs 本地          vs 本地结果: 完全相同 -> 疑似未通信

两个信号叠在一起，结论就很清楚：rank0 从头到尾没看过 rank1 的数据。 它既对不上全序列参考，又和自己的本地结果逐位相等------这两件事同时成立，只可能是"没通信"。

---

四、根因：两个必须做对的事

翻实现的时候，问题出在两处，都很典型。

问题 1：跳过条件漏了一个前置条件

环形注意力的循环大概长这样：

for step in range(1, world_size):
    if step <= comm.rank:      # ← 这里
        continue
    ...

这行的本意是处理 causal（因果）注意力：因果掩码下，位置 i 只能看到 j ≤ i，所以"后面的块"可以跳过。

但它只对 causal 成立。扩散模型的注意力是非因果的（每个位置都能看到全序列），这个"跳过"就变成了"少算"。

正确写法是：

if is_causal and step <= comm.rank:
    continue

少一个 is_causal，非因果场景下每张卡会静默地少算若干块，且不报错。

问题 2：分块结果没有正确合并

这是更本质的一个。

注意力输出的本质是加权平均，权重来自 softmax。分块计算时，每块只能算出"以本块为分母"的局部 softmax，直接相加是错的------因为各块的归一化因子不同。

正确做法是用 log-sum-exp 逐步合并（online softmax）：

def combine(out1, lse1, out2, lse2):
    """把两块的结果按 softmax 权重合并"""
    lse = torch.logaddexp(lse1, lse2)
    w1 = torch.exp(lse1 - lse)
    w2 = torch.exp(lse2 - lse)
    return out1 * w1 + out2 * w2, lse

完整循环：

# 初始化：先用本地块
out, lse = local_attn_with_lse(q, k_l, v_l)

for step in range(1, world_size):
    k_r, v_r = ring_send_recv(k_l, v_l)        # 收上一张卡的，发给下一张
    o_r, l_r = local_attn_with_lse(q, k_r, v_r)
    out, lse = combine(out, lse, o_r, l_r)     # ← 少了这行就只保留最后一块

return out

那个库的问题就是从来没有调用过合并函数。所以每张卡的输出要么是"只算了本地块"，要么是"只保留了循环最后一次的结果"------两种都能跑完、都不报错。

---

五、给合并逻辑补一个单元测试

上面说"要合并"，但怎么验证你的合并写对了？

不需要多卡环境，单机就能测。思路是：分块算的结果，必须等于一次算完的结果。

"""
test_combine.py ------ 验证分块 LSE 合并与一次性计算等价
纯单机、CPU 即可运行，不需要分布式环境。
"""
import torch


def naive_attn(q, k, v):
    """一次算完的参考实现"""
    scale = q.shape[-1] ** -0.5
    s = torch.matmul(q, k.transpose(-2, -1)) * scale
    return torch.softmax(s, dim=-1) @ v


def local_attn_lse(q, k, v):
    """算局部注意力，同时返回 softmax 的 log-sum-exp

    q: (B, Sq, H, D) 本卡的 query（完整或分片）
    k/v: (B, Sk, H, D) 当前这一块的 KV
    """
    scale = q.shape[-1] ** -0.5
    s = torch.matmul(q, k.transpose(-2, -1)) * scale
    m = s.max(dim=-1, keepdim=True).values          # 减去最大值防溢出
    p = torch.exp(s - m)
    l = p.sum(dim=-1, keepdim=True)
    o = torch.matmul(p, v) / l
    lse = (m + torch.log(l)).squeeze(-1)            # (B, Sq, H)
    return o, lse


def combine(o1, lse1, o2, lse2):
    """按 softmax 权重合并两块结果（online softmax 的核心）"""
    lse = torch.logaddexp(lse1, lse2)
    w1 = torch.exp(lse1 - lse).unsqueeze(-1)
    w2 = torch.exp(lse2 - lse).unsqueeze(-1)
    return o1 * w1 + o2 * w2, lse


def chunked_attn(q, k, v, n_chunks):
    """模拟 ring：q 固定，KV 分块遍历并逐步合并"""
    per = k.shape[1] // n_chunks
    out = lse = None
    for i in range(n_chunks):
        sl = slice(i * per, (i + 1) * per)
        oi, li = local_attn_lse(q, k[:, sl], v[:, sl])
        out, lse = (oi, li) if out is None else combine(out, lse, oi, li)
    return out


def test_chunked_equals_full(dtype=torch.float32):
    torch.manual_seed(0)
    B, S, H, D = 1, 128, 2, 32
    q, k, v = (torch.randn(B, S, H, D, dtype=dtype) for _ in range(3))
    ref = naive_attn(q, k, v)

    print(f"\n精度 = {dtype}")
    for n in (1, 2, 4, 8, 16):
        got = chunked_attn(q, k, v, n)
        err = (got.float() - ref.float()).abs().max().item()
        cos = torch.nn.functional.cosine_similarity(
            got.float().flatten(), ref.float().flatten(), dim=0
        ).item()
        flag = "OK" if cos > 0.999 else "FAIL"
        print(f"  {n:>2} 块  max_err={err:.2e}  cos={cos:.8f}  {flag}")
        assert cos > 0.999, f"{n} 块合并后与整体计算不一致"


def test_wrong_merge_is_caught(dtype=torch.float32):
    """反例：如果不用 LSE 加权，而是简单平均，误差有多大"""
    torch.manual_seed(0)
    B, S, H, D = 1, 128, 2, 32
    q, k, v = (torch.randn(B, S, H, D, dtype=dtype) for _ in range(3))
    ref = naive_attn(q, k, v)

    per, outs = S // 4, []
    for i in range(4):
        sl = slice(i * per, (i + 1) * per)
        oi, _ = local_attn_lse(q, k[:, sl], v[:, sl])
        outs.append(oi)

    wrong_avg = sum(outs) / len(outs)            # 错法 A：简单平均
    wrong_last = outs[-1]                        # 错法 B：只用最后一块

    for name, o in (("简单平均", wrong_avg), ("只用最后一块", wrong_last)):
        err = (o.float() - ref.float()).abs().max().item()
        cos = torch.nn.functional.cosine_similarity(
            o.float().flatten(), ref.float().flatten(), dim=0
        ).item()
        print(f"  错误做法[{name}]  max_err={err:.3f}  cos={cos:.4f}  "
              f"{'看起来还行（危险）' if cos > 0.99 else '明显错误'}")


if __name__ == "__main__":
    test_chunked_equals_full(torch.float32)
    test_chunked_equals_full(torch.bfloat16)
    print()
    test_wrong_merge_is_caught()

这个测试为什么有用

第一，它不需要多卡。 合并逻辑的正确性和通信无关------把 KV 切成几块、逐块算、逐步合并，单进程就能完整复现 ring attention 的数据流。这样 CI 里随时能跑，几十毫秒出结果。

公式本身我用纯 Python 单独验证过一遍：切 1 / 2 / 4 / 8 块，合并结果和一次性计算的误差都在 1e-16 量级（float64 的机器精度）。说明这套合并公式是严格的数学等价，不是近似。

第二，它能抓住"看起来还行"的错误。 用同样的方法跑两种错误写法，相对误差是这样的：

简单平均         max_err=0.620   相对误差=58.5%
只用最后一块      max_err=1.255   相对误差=118.3%

注意第二个------比不合并更离谱，因为它把前面所有块的信息全丢了。

而这个量级的错误，在单层上看起来往往是"结果还算合理"：张量形状对、数值范围正常、没有 NaN。既不是"明显崩了"（那样一眼就能发现），也不是"接近正确"。它就落在那个最难受的区间里。

这也解释了为什么端到端要跑几十步才暴露：单层几成的相对误差，经过 40 步迭代放大，就变成了完全不同的一张图。

第三，它给了你一个可以钉死的验收标准。 以后换库、升级版本、改实现，先跑这个测试，cos > 0.999 过不了就别往下接。

---

六、一份可复用的检查清单

这类"静默算错"的问题，我总结了几条通用检查法，遇到新的并行库可以直接套：

数值层

• 小输入（seq=64，heads=2，dim=32）先跑通，别一上来就上真实尺寸
• 和朴素实现逐元素比，容差设成 bf16 级别（1e-2）
• 分块数与整体结果必须等价（第 5 节那个单测）
• 检查每张卡的误差是否一致------如果只有部分 rank 错，基本能定位到通信路径

通信层

• "并行结果 vs 只用本地数据的结果"必须有差异
• 把远端数据换成特殊值（比如全 1e4），看输出是否随之改变
• 检查通信次数：world_size - 1 次还是 world_size 次，写错也能跑

结构层

• 分片是否连续、无重叠、无遗漏（打印区间核对）
• 输出拼回去的顺序对不对（常见错误：按 rank 顺序拼 vs 按原始位置拼）
• 有没有 causal 掩码相关的条件分支，前置条件是否完整

环境层

• 单卡先验证库本身（有些库单卡路径和多卡路径是两套代码）
• 换后端（gloo 换 nccl/各家加速卡后端）行为是否一致
• 进程数不是 2 的时候（比如 3 卡）是否还成立------很多实现在 2 的幂次上没问题

---

七、最后一点体会

第三方高性能库最大的问题不是"质量差"，而是它的失败模式太安静。

一个手写的 for 循环写错了，你大概率会看到崩或者 NaN。但一个精心优化过的并行 kernel 写错了，它照样能输出形状正确、数值合理的张量------只是那个张量不是你要的。

所以凡是引入第三方的高性能实现，都要在接入的第一天做数值验证，别等到端到端结果不对了再回头查。那时代码已经写了上千行，怀疑的对象太多了。

三十行的验证脚本，能省三天。

---

文中代码为示意，接口按通用约定设计，接入时替换为你实际使用的库。

📎 本文的排查经验，也浓缩进了另一篇更体系化的踩坑总结：《扩散模型推理加速的三个方向，以及移植到国产 NPU 时踩的坑》