PyTorch矩阵乘法切片不一致问题：Transformer批量输入差异排查

在处理Transformer模型长输入的批量运算时，发现批量计算和逐段非批量计算的结果存在差异。通过隔离测试得到如下代码：

import torch

n = 20

vec = torch.rand(n, 20)
a = torch.rand(30, 20)

for i in range(1, n+1):
    print(i, torch.equal(
        torch.nn.functional.linear(vec, a)[:i],
        torch.nn.functional.linear(vec[:i], a)))

运行输出：

1 False
2 False
3 False
4 True
5 True
6 True
7 False
8 False
9 False
10 True
11 True
12 True
13 False
14 False
15 False
16 True
17 True
18 True
19 True
20 True

原因分析 #

这种差异源于浮点数运算的精度特性以及PyTorch内部的矩阵乘法优化策略：

• 浮点数（如FP32）本身存在精度限制，批量计算是对整个vec矩阵与a做乘法，非批量计算则是对vec[:i]子集运算，两者的计算顺序、中间舍入步骤不同，最终累积出微小偏差。
• PyTorch会根据输入张量的形状自动选择最优矩阵乘法实现（如不同CUDA核或CPU优化指令），批量与非批量输入形状差异会触发不同的优化逻辑，进一步放大精度差异。
• 输出呈现的周期性规律（如3个False后3个True）和张量内存对齐、分块计算策略有关，不同大小的输入会触发不同分块处理方式，导致误差表现出周期性。

解决方案 #

针对这类问题，可从以下角度处理：

• 放宽结果校验精度阈值：放弃使用torch.equal（要求完全精确匹配），改用torch.allclose并设置合理容差，例如torch.allclose(result1, result2, rtol=1e-5, atol=1e-8)，阈值可根据业务需求调整。
• 统一计算路径：若需要严格一致的结果，强制让批量与非批量计算使用相同逻辑。比如非批量计算时补零对齐批量输入维度，或手动指定矩阵乘法实现方式（如torch.matmul结合precision参数，需匹配PyTorch版本与硬件支持）。
• 使用更高精度浮点数：将张量类型从float32改为float64（double），更高精度能显著降低舍入误差累积，缩小批量与非批量结果的差异。
• 禁用自动优化：通过torch.backends.cudnn.benchmark = False关闭PyTorch自动选择最优核的功能，强制使用固定计算实现，不过会牺牲一定性能。

内容的提问来源于stack exchange，提问作者Sasha