背景
我最近在做一件很具体的事：把一个扩散模型（图像生成模型）的推理框架，从 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 时踩的坑》
