KrylovKit eigsolve在CPU与GPU上结果不一致的问题咨询

问题描述 #

我尝试对比KrylovKit库中eigsolve函数在CPU和GPU上的运行结果，编写了如下Julia代码：

using CUDA
using CUDA.CUSPARSE
using LinearAlgebra
using LinearMaps
using SparseArrays
using Random
using KrylovKit

function eign_cpu(M::Union{AbstractMatrix,LinearMap{T}}, v0::Vector{T}) where {T<:Number}
    tol = 1e-10
    (E, V, info) = eigsolve(M, v0, 1, :SR; ishermitian=true, tol=tol)
    return (E[1], V[1], info)
end

function eign_gpu(M::Union{CuArray,LinearMap{T}}, v0::CuVector{T}) where {T<:Number}
    tol = 1e-10
    (E, V, info) = eigsolve(M, v0, 1, :SR; ishermitian=true, tol=tol)
    return (E[1], V[1], info)
end

Random.seed!(0)

n = 4
m = 4
Es = [1, 1, 2, 3]
Ed = [2, 3, 4, 4]

w = ones(4)
sw = sqrt.(w)
sw_d = CuArray(sw)

C = LinearMap{Float64}((out::AbstractArray, y::AbstractArray) -> (out .= sw; out .*= -dot(sw, y)), n, ismutating=true, issymmetric=true)
C_d = LinearMap{Float64}((out::CuArray, y::CuArray) -> (out .= sw_d; out .*= -dot(sw_d, y)), n, ismutating=true, issymmetric=true)

M = sparse(Es, Ed, ones(Float64, m), n, n)
M_d = CuSparseMatrixCSC(M)

adj_cpu = p::AbstractArray -> (M.nzval .= p ./ 2; Symmetric(M))
adj_gpu = p::CuArray -> (M_d.nzVal .= p ./ 2; Symmetric(M_d))

p = rand(m, 1)
p_d = CuArray(p)

aux = Vector{Float64}(undef, n)
eig_grad = randn(Float64, n)
v0 = ((randn!(aux) .*= 1e-4) .+= eig_grad)
v0_d = CuArray(v0)

if isapprox(v0, Array(v0_d))
    println("v and v0 are approximate")
end
if isapprox(sw, Array(sw_d))
    println("C and C_d are approximate")
end
if isapprox(M, SparseMatrixCSC(M_d))
    println("M and M_d are approximate")
end
S = C + adj_cpu(p)
(E_cpu, V_cpu, info_cpu) = eign_cpu(S, v0)
println("info cpu = ", info_cpu)
S_d = C_d + adj_gpu(p_d)
(E_gpu, V_gpu, info_gpu) = eign_gpu(S_d, v0_d)
println("info gpu = ", info_gpu)

if !isapprox(V_cpu, Array(V_gpu))
    println("WARN: diverging eigenvalues between cpu and gpu:
",
        "E_cpu = ", E_cpu, " V_cpu = ", V_cpu, "
E_gpu = ", E_gpu, " V_gpu = ", V_gpu)
    readline()
end

输入数据经isapprox校验均返回true，但运行输出如下：

info cpu = ConvergenceInfo: 4 converged values after 1 iterations and 4 applications of the linear map;
norms of residuals are given by (1.3683491030815672e-36, 6.545806401840613e-35, 6.344902430327517e-34, 8.827885896605726e-33).

info gpu = ConvergenceInfo: 4 converged values after 1 iterations and 4 applications of the linear map;
norms of residuals are given by (7.776628566574466e-20, 1.087879714604506e-17, 4.976482876540383e-17, 3.382716406082749e-17).

WARN: different eigenvalues between cpu and gpu:
E_cpu = -3.6566857184706016 V_cpu = [-0.5154937434198816, -0.46021338805521755, -0.5355038728938648, -0.485495046386023]
E_gpu = -3.7700145739756743 V_gpu = [0.48275625184457943, 0.47793647243109605, 0.52265534023327, 0.5151257370300364]

原本期望CPU与GPU的结果近似相等，请问本次实验是否存在疏漏或错误假设？

问题排查与分析 #

1. 线性映射的对称性偏差 #

你标记了C和C_d为对称映射，但GPU上dot(sw_d, y)的浮点数计算精度与CPU存在差异，会导致映射的对称性出现微小偏差。KrylovKit依赖ishermitian=true的假设做优化求解，这种偏差会被放大，最终得到不同的特征对。

2. 稀疏矩阵原地修改的隐患 #

adj_cpu和adj_gpu直接修改全局的M和M_d非零值：

• 你只校验了初始状态的M和M_d，但未验证修改后的对称矩阵是否与CPU端一致。p ./ 2在CPU和GPU上的浮点数计算可能存在微小差异，导致最终的算子S和S_d数学上不等价。
• 这种全局变量原地修改的设计存在风险，后续代码若复用变量会引发意外错误。

3. 特征向量比较逻辑缺陷 #

特征向量的符号翻转是正常数学性质（特征向量乘以任意非零标量仍为特征向量），但你的isapprox直接比较元素，未考虑符号差异，会误判为不相等。不过本次特征值也存在差异，说明还有其他核心问题。

4. 迭代收敛精度差异 #

CPU和GPU的残差范数相差16个数量级，GPU端的残差精度远低于CPU。虽然设置了tol=1e-10，但GPU浮点数运算的特性可能导致迭代提前终止，未收敛到与CPU一致的特征对。

5. 线性算子相加的兼容性问题 #

LinearMap与Symmetric稀疏矩阵相加时，CPU和GPU的底层实现逻辑可能存在差异。需要验证S * x与Array(S_d * CuArray(x))是否近似相等，确认两个算子的数学一致性。

修正建议 #

• 避免全局变量原地修改：adj_cpu和adj_gpu应创建新矩阵，比如adj_cpu(p) = Symmetric(sparse(Es, Ed, p ./ 2, n, n))，确保每次生成的矩阵独立且一致。
• 验证算子等价性：取同一输入向量x，分别计算S * x和Array(S_d * CuArray(x))，检查是否近似相等。
• 调整特征向量比较逻辑：比较时考虑符号差异，比如isapprox(V_cpu, Array(V_gpu)) || isapprox(V_cpu, -Array(V_gpu))。
• 提高迭代精度：调小tol（如1e-12），增加maxiter（如100），让GPU端求解更充分收敛。
• 验证映射对称性：手动检查GPU映射的对称性，比如对任意GPU向量x、y，确认dot(x, C_d * y)与dot(C_d * x, y)近似相等。

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