Numba并行化jit函数结果异常问题排查求助

我尝试并行化一个经Numba JIT编译的最小二乘函数（用于处理含NaN的因变量），原本以为这是个易并行的问题，但并行版本与串行版本输出结果不一致，出现异常情况。请问可能的原因是什么，以及有什么解决办法？

可复现示例 #

import numba
import numpy as np
from sklearn.datasets import make_regression


@numba.jit(nopython=True)
def nanlstsq(X, y):
    """Return the least-squares solution to a linear matrix equation

    Analog to ``numpy.linalg.lstsq`` for dependant variable containing ``Nan``

    Args:
        X ((M, N) np.ndarray): Matrix of independant variables
        y ({(M,), (M, K)} np.ndarray): Matrix of dependant variables

    Returns:
        np.ndarray: Least-squares solution, ignoring ``Nan``
    """
    beta = np.zeros((X.shape[1], y.shape[1]), dtype=np.float64)
    for idx in range(y.shape[1]):
        # subset y and X
        isna = np.isnan(y[:,idx])
        X_sub = X[~isna]
        y_sub = y[~isna,idx]
        # Compute beta on data subset
        XTX = np.linalg.inv(np.dot(X_sub.T, X_sub))
        XTY = np.dot(X_sub.T, y_sub)
        beta[:,idx] = np.dot(XTX, XTY)
    return beta


@numba.jit(nopython=True, parallel=True)
def nanlstsq_parallel(X, y):
    beta = np.zeros((X.shape[1], y.shape[1]), dtype=np.float64)
    for idx in numba.prange(y.shape[1]):
        # subset y and X
        isna = np.isnan(y[:,idx])
        X_sub = X[~isna]
        y_sub = y[~isna,idx]
        # Compute beta on data subset
        XTX = np.linalg.inv(np.dot(X_sub.T, X_sub))
        XTY = np.dot(X_sub.T, y_sub)
        beta[:,idx] = np.dot(XTX, XTY)
    return beta


# Generate random data
n_targets = 10000
n_features = 3
X, y = make_regression(n_samples=200, n_features=n_features,
                       n_targets=n_targets)
# Add random nan to y array
y.ravel()[np.random.choice(y.size, 5*n_targets, replace=False)] = np.nan
# Run the regression
beta = nanlstsq(X, y)
beta_parallel = nanlstsq_parallel(X, y)

np.testing.assert_allclose(beta, beta_parallel)

可能的原因 #

• 数值稳定性差异：直接通过np.linalg.inv计算逆矩阵求解最小二乘的方式本身不稳定，当X_sub的列接近线性相关或样本数不足时，XTX会接近奇异矩阵，求逆操作会放大数值误差。串行与并行环境下Numba对矩阵运算的优化策略不同，导致误差放大程度不一致，最终结果差异显著。
• 临时数组内存竞争：在prange循环内创建X_sub、y_sub等临时数组时，Numba的内存管理器在多线程环境下可能出现隐式内存竞争，导致部分迭代的临时数据被意外覆盖，计算出错误结果。
• 并行模式下的函数行为差异：Numba对np.isnan的JIT编译在并行模式下可能存在实现差异，导致部分NaN值未被正确识别，进而选取错误的数据集子集计算。

解决办法 #

1. 改用更稳定的求解方式
   放弃直接求逆的逻辑，使用np.linalg.lstsq专门求解最小二乘问题，数值稳定性更高：

   # 替换原求逆计算代码
   beta[:, idx] = np.linalg.lstsq(X_sub, y_sub, rcond=None)[0]
   

2. 避免临时数组，直接用掩码计算
   通过掩码直接计算XTX和XTY，减少临时数组创建，规避内存竞争：

   mask = ~np.isnan(y[:, idx])
   # 先检查样本数是否足够
   if mask.sum() >= X.shape[1]:
       XTX = np.dot(X[mask].T, X[mask])
       XTY = np.dot(X[mask].T, y[mask, idx])
       # 条件数判断避免奇异矩阵
       if np.linalg.cond(XTX) < 1e10:
           beta[:, idx] = np.dot(np.linalg.inv(XTX), XTY)
       else:
           beta[:, idx] = np.linalg.lstsq(X[mask], y[mask, idx], rcond=None)[0]
   else:
       # 样本不足时的 fallback 逻辑，比如填充NaN
       beta[:, idx] = np.nan
   

3. 显式隔离线程私有变量
   在prange循环内显式声明所有临时变量（如isna、XTX），确保每个线程拥有独立的变量副本，避免共享内存冲突。

内容的提问来源于stack exchange，提问作者Loïc Dutrieux