PyTorch同一矩阵Cholesky分解GPU失败、CPU正常的问题求助

问题重现 #

以下代码在CPU环境下可正常执行Cholesky分解：

import torch
device = torch.device('cpu')
torch.manual_seed(1)
size = 4096
F = torch.rand(int(size / 2), int(size / 2)).to(device)
F = torch.matmul(F, F.T)
torch.linalg.cholesky(F)

但将设备切换为GPU后，执行相同逻辑的代码会抛出正定检查失败的异常：

import torch
device = torch.device('cuda')
torch.manual_seed(1)
size = 4096
F = torch.rand(int(size / 2), int(size / 2)).to(device)
F = torch.matmul(F, F.T)
torch.linalg.cholesky(F)

报错信息：

torch.linalg.cholesky(F)
torch._C._LinAlgError: linalg.cholesky: The factorization could not be completed because the input is not positive-definite (the leading minor of order 2044 is not positive-definite).

环境信息：PyTorch 2.3.0，CUDA 12.1

原因分析 #

理论上，F = F @ F.T构造的矩阵是半正定矩阵，而Cholesky分解要求输入为严格正定矩阵（或半正定且无零特征值）。CPU与GPU执行结果的差异源于两点：

1. 浮点数计算精度差异：GPU默认单精度（float32）运算的数值误差比CPU单精度运算更大；且CPU与GPU的浮点运算实现（舍入规则、矩阵乘法优化逻辑）存在差异，导致GPU上构造的矩阵出现微小负特征值或非正定主子式，触发正定检查失败。
2. 设备随机数生成差异：即使设置了全局随机种子，CPU与GPU的随机数生成器实现不同，生成的初始矩阵存在细微差异，进一步放大了后续矩阵乘法的数值误差。

解决方法 #

1. 使用更高精度的数据类型 #

将矩阵转换为float64（双精度）类型，大幅降低数值误差，使矩阵更接近严格正定：

import torch
device = torch.device('cuda')
torch.manual_seed(1)
size = 4096
# 生成双精度随机矩阵
F = torch.rand(int(size / 2), int(size / 2), dtype=torch.float64).to(device)
F = torch.matmul(F, F.T)
torch.linalg.cholesky(F)

2. 添加对角扰动确保正定 #

给矩阵对角线添加极小正数扰动，强制矩阵成为严格正定矩阵：

import torch
device = torch.device('cuda')
torch.manual_seed(1)
size = 4096
F = torch.rand(int(size / 2), int(size / 2)).to(device)
F = torch.matmul(F, F.T)
# 添加对角扰动，epsilon可根据矩阵规模调整
epsilon = 1e-6
F = F + epsilon * torch.eye(int(size / 2), device=device)
torch.linalg.cholesky(F)

3. 使用容错性更强的分解API #

使用torch.linalg.cholesky_ex，该API返回分解结果和错误状态，允许处理非正定情况：

import torch
device = torch.device('cuda')
torch.manual_seed(1)
size = 4096
F = torch.rand(int(size / 2), int(size / 2)).to(device)
F = torch.matmul(F, F.T)
# 获取分解结果和错误码
L, info = torch.linalg.cholesky_ex(F)
if info.item() == 0:
    print("Cholesky分解成功")
else:
    # 分解失败时添加扰动重试
    F = F + 1e-6 * torch.eye(int(size / 2), device=device)
    L = torch.linalg.cholesky(F)

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