[Bug已解决] cholesky_solve 批处理多 batch CUDA float64 在 ROCm gfx942 上数值不接近解决方案

一键难忘

542人浏览 · 2026-07-17 12:09:40

一键难忘 · 2026-07-17 12:09:40 发布

[Bug已解决] cholesky_solve 批处理多 batch 在 ROCm gfx942 CUDA float64 上数值不接近（Tensor-likes are not close）解决方案

一、现象长什么样

你在 AMD ROCm 平台（gfx942，如 Instinct MI300 系列） 上用 torch.linalg.cholesky_solve 解对称正定线性方程组，并且输入是很多 batch 的批处理、float64（双精度）：

import torch
# A: (N, D, D) 多个 batch 的对称正定矩阵，float64
X = torch.linalg.cholesky_solve(B, A_cholesky)   # 解 A X = B

结果测试/验证报：

test_cholesky_solve_batched_many_batches_cuda_float64 fails on ROCm gfx942
(Tensor-likes are not close)

即 pytorch/pytorch#188553。含义：在 ROCm gfx942 上，`cholesky_solve` 的「批处理 + 很多 batch + float64」这一组合，数值结果与参考（CPU / 其他架构）不接近，测试失败。不是崩溃，而是算出来答案误差超 tolerance。本文聚焦：cholesky_solve 是什么、为什么 ROCm gfx942 的多 batch float64 会数值偏差大、怎么绕过（用 CPU 参考 / 换算法 / 升级）。

二、背景：cholesky_solve 与批处理

`torch.linalg.cholesky_solve(B, A_cholesky)` 求解 A X = B，其中 A 是对称正定矩阵，用 Cholesky 分解 `A = L L^T` 高效求解（比直接求逆稳）。 `A_cholesky` 是 `torch.linalg.cholesky(A)` 的结果（下三角 L）。`B` 是右端项，`X` 是解。批处理（batched）：当 A 是 `(batch, D, D)`、B 是 `(batch, D, K)`，算子对每个 batch 独立求解。这依赖底层 batched BLAS / LAPACK（如 cuSOLVER / hipSOLVER 的 potrsBatched）。 float64：双精度，对数值误差敏感（tolerance 更严——float64 参考解的容差通常是 1e-12 量级，而 float32 是 1e-4 量级）。「Tensor-likes are not close」说明 gfx942 上该组合算出的 X 与参考差超出了 float64 的严容差。

三、为什么 ROCm gfx942 多 batch float64 会数值不接近

ROCm 上 cholesky_solve 走 hipSOLVER / rocSOLVER 的 batched potrs。问题来源：

batched 内核在 gfx942 上的实现缺陷：「很多 batch」触发了某个 batched 代码路径（可能是把大 batch 拆块、或用某种打包格式），在 gfx942 上该路径的数值累积误差偏大，或某分支算错；
float64 严容差放大了偏差：若是 float32，误差在 1e-4 内就过了；但 float64 参考要求 1e-12，于是「微小偏差」在 float64 下被判定为「不接近」；
gfx942 是新架构：rocSOLVER 的 batched float64 内核在 gfx942 上可能还没充分调优/验证，存在未修复的数值 bug；
多 batch 触发边界：batch 数很大时（测试中 "many batches"），可能踩中批处理调度的边界条件（如 stride 计算、batch 分块），导致部分 batch 解错。本质：ROCm gfx942 的 batched float64 cholesky_solve 内核有数值/实现缺陷，在严容差下暴露。

四、最小可运行复现（带守卫，跨设备对比）

下面演示跨设备对比，确认「ROCm gfx942 与 CPU 参考不一致」（实际触发需 gfx942；用守卫说明）：

import torch
def make_spdbatch(batch, D):
    # 构造对称正定矩阵：A = M M^T + I
    M = torch.randn(batch, D, D, dtype=torch.float64)
    A = M @ M.transpose(-1, -2) + torch.eye(D, dtype=torch.float64).expand(batch, D, D)
    return A
def solve_and_compare():
    batch, D = 64, 8
    A = make_spdbatch(batch, D)
    B = torch.randn(batch, D, 3, dtype=torch.float64)
    # CPU 参考（最可信）
    L_cpu = torch.linalg.cholesky(A)
    X_cpu = torch.linalg.cholesky_solve(B, L_cpu)
    # GPU 解（ROCm gfx942 上可能偏差大）
    if torch.cuda.is_available():
        A_g = A.cuda()
        B_g = B.cuda()
        L_g = torch.linalg.cholesky(A_g)
        X_g = torch.linalg.cholesky_solve(B_g, L_g)
        close = torch.allclose(X_g.cpu(), X_cpu, atol=1e-10, rtol=1e-8)
        print("GPU 与 CPU 参考接近:", close)
        # 不接近 → 命中 gfx942 多 batch float64 偏差
    else:
        print("[skip] 无 GPU，仅说明对比方法")
if __name__ == "__main__":
    solve_and_compare()

要点：用 CPU 参考解作基准对比，GPU 结果偏差超 float64 容差即命中。

五、解决方案一：用 CPU 计算 cholesky_solve（最稳）

若你的场景可接受，把 cholesky_solve 放到 CPU（float64 在 CPU 上用 MKL/LAPACK，数值最稳）：

import torch
def solve_on_cpu(A, B):
    # A, B 在 CPU float64，结果可靠
    L = torch.linalg.cholesky(A)
    X = torch.linalg.cholesky_solve(B, L)
    return X
# 用法：数据在 GPU 时先搬 CPU 算，再搬回
if torch.cuda.is_available():
    A_g, B_g = A.cuda(), B.cuda()
    X = solve_on_cpu(A_g.cpu(), B_g.cpu()).cuda()

代价：多一次 CPU/GPU 搬运，但数值绝对可靠，规避 ROCm gfx942 的 batched float64 偏差。若 batch 不大、D 不大，CPU 上 LAPACK 也很快。

六、解决方案二：降低精度到 float32 放宽容差

若业务允许，用 float32——float32 的参考容差宽松（1e-4），ROCm gfx942 的偏差通常在 float32 容差内「通过」：

import torch
A = make_spdbatch(batch, D).float()   # float32
B = torch.randn(batch, D, 3).float()
if torch.cuda.is_available():
    A, B = A.cuda(), B.cuda()
L = torch.linalg.cholesky(A)
X = torch.linalg.cholesky_solve(B, L)

但注意：float32 牺牲精度，只适合对数值误差不敏感的场景（如推理、非科学计算）。科学计算要 float64 就别用这招。

七、解决方案三：避免「很多 batch」触发缺陷路径（分批）

若是「many batches」触发了 batched 内核的坏路径，可把大 batch 拆成小块，每块用非 batched / 小 batched 路径：

import torch
def solve_in_chunks(A, B, chunk=16):
    # 把大 batch 拆小，规避 gfx942 多 batch batched 内核缺陷
    out = []
    for i in range(0, A.size(0), chunk):
        Ai, Bi = A[i:i+chunk], B[i:i+chunk]
        if torch.cuda.is_available():
            Ai, Bi = Ai.cuda(), Bi.cuda()
        Li = torch.linalg.cholesky(Ai)
        Xi = torch.linalg.cholesky_solve(Bi, Li)
        out.append(Xi.cpu())
    return torch.cat(out, dim=0)
# 用 chunk 较小的批处理，可能绕开「many batches」坏路径

原理：减小单次 batched 的 batch 数，避开触发缺陷的批处理调度边界。

八、解决方案四：升级 ROCm / PyTorch 到修复版本

#188553 是 ROCm gfx942 的内核缺陷，官方会修（rocSOLVER 更新）。升级：

# 升级 PyTorch（其绑定更新的 ROCm / rocSOLVER）
pip install --upgrade torch --index-url https://download.pytorch.org/whl/rocm6.3
# 或升级系统 rocSOLVER / hipSOLVER

判断修复：同样「many batches float64」`cholesky_solve` 在 gfx942 上与 CPU 参考 `allclose(atol=1e-10)` 通过。修复前，用 CPU 计算 / 降 float32 / 分批绕过。

九、排查清单

ROCm gfx942 + 多 batch + float64 的 cholesky_solve 数值偏差 → 确认 #188553。
用 CPU 参考 对比：GPU 结果超 float64 容差（1e-12 级）即命中。
绕过：改 CPU 算 cholesky_solve（最稳）；或降 float32（放宽容差，牺牲精度）；
分批：把 many batches 拆小 chunk，避开 batched 坏路径。
升级 ROCm / PyTorch / rocSOLVER 到修复版。
验证：修复后与 CPU 参考 allclose(atol=1e-10, rtol=1e-8)。

十、小结

test_cholesky_solve_batched_many_batches_cuda_float64 fails on ROCm gfx942 (Tensor-likes are not close)（#188553）的本质是：在 AMD ROCm gfx942 上，torch.linalg.cholesky_solve 的「批处理 + 很多 batch + float64」组合，依赖的 rocSOLVER/hipSOLVER batched 内核有数值/实现缺陷，算出的解与 CPU 参考差超出 float64 的严容差（1e-12 级），测试失败。float64 的严容差把「微小偏差」放大成了「不接近」。应对：

CPU 计算最稳：把 cholesky_solve 放 CPU float64（MKL/LAPACK 可靠），GPU 数据搬过去算再搬回；
降 float32：业务允许时降精度，float32 容差宽松可「通过」（牺牲精度）；
分批规避：大 batch 拆小 chunk，避开 batched 内核的缺陷路径；
升级：ROCm / rocSOLVER / PyTorch 修复后，gfx942 上多 batch float64 恢复准确。记住：float64 容差极严（~1e-12），任何内核的微小数值偏差都会在 float64 下暴露；ROCm 新架构（gfx942）的 batched LAPACK 内核常有此类未修复偏差，关键计算用 CPU 参考兜底最稳。