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

[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, dtypetorch.float64) A M M.transpose(-1, -2) torch.eye(D, dtypetorch.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, dtypetorch.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, atol1e-10, rtol1e-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放到CPUfloat64 在 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-4ROCm 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, chunk16): # 把大 batch 拆小规避 gfx942 多 batch batched 内核缺陷 out [] for i in range(0, A.size(0), chunk): Ai, Bi A[i:ichunk], B[i:ichunk] 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, dim0) # 用 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(atol1e-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(atol1e-10, rtol1e-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 float64MKL/LAPACK 可靠GPU 数据搬过去算再搬回降 float32业务允许时降精度float32 容差宽松可「通过」牺牲精度分批规避大 batch 拆小 chunk避开 batched 内核的缺陷路径升级ROCm / rocSOLVER / PyTorch 修复后gfx942 上多 batch float64 恢复准确。 记住float64 容差极严~1e-12任何内核的微小数值偏差都会在 float64 下暴露ROCm 新架构gfx942的 batched LAPACK 内核常有此类未修复偏差关键计算用 CPU 参考兜底最稳。