如何解决 RuntimeError: CUDA error: CUBLAS_STATUS_ALLOC_FAILED when calling cublasCreate(handle) 错误

如何解决 RuntimeError: CUDA error: CUBLAS_STATUS_ALLOC_FAILED when calling cublasCreate(handle) 错误

在深度学习的训练过程中，使用CUDA加速运算是常见的做法。但有时，我们可能会遇到 RuntimeError: CUDA error: CUBLAS_STATUS_ALLOC_FAILED when calling cublasCreate(handle) 这样的错误信息。这个错误通常出现在使用 PyTorch 或 TensorFlow 等深度学习框架时，原因多种多样，涉及到硬件资源、驱动版本、CUDA环境等多个因素。下面我们将详细解析可能导致这个问题的原因，并提供解决方案。

一、错误分析

错误信息 RuntimeError: CUDA error: CUBLAS_STATUS_ALLOC_FAILED when calling cublasCreate(handle) 通常出现在 CUDA 库的初始化过程中，尤其是在 CUBLAS（CUDA Basic Linear Algebra Subprograms）库中。CUBLAS 是一个用于矩阵计算和线性代数运算的高效库，常用于深度学习框架中。

错误的核心原因通常是与 GPU 资源的分配有关，尤其是 GPU 的内存不足或没有正确初始化。通常，这个错误发生在以下几种情况下：

1. GPU 内存不足：当前任务需要的内存超过了 GPU 可用内存。
2. CUDA 驱动问题：CUDA 驱动和库的版本不匹配，或者环境配置错误。
3. 多任务争用 GPU 资源：多个程序或进程占用了 GPU 资源，导致内存不足。
4. 硬件故障或驱动问题：显卡硬件故障或驱动版本不兼容。
5. CUDA 环境配置错误：在配置 CUDA 环境时，某些参数设置不正确。

二、解决方案

针对 CUBLAS_STATUS_ALLOC_FAILED 错误，有几个常见的解决步骤。下面将逐一列出并详细介绍如何进行排查和解决。

1. 检查 GPU 内存使用情况

GPU 内存不足是最常见的原因之一。我们可以使用 nvidia-smi 命令来检查当前 GPU 的内存占用情况。如果有其他进程正在占用 GPU 内存，考虑终止这些进程，或者降低当前任务的内存需求。

• 打开终端，运行：

nvidia-smi

• 输出会显示每个 GPU 的内存占用情况。如果内存已满，建议尝试以下几种方式：
  • 结束其他不必要的进程。
  • 如果是深度学习任务，可以尝试减小 batch size。
  • 使用混合精度训练（FP16），以减少内存占用。

2. 重启机器或显卡

有时，GPU 内存没有正确释放或者被占用，导致后续任务无法正确初始化。尝试重启机器或者单独重启 GPU 驱动，可以解决这种问题。

• 重启机器：

sudo reboot

• 如果只需要重启 GPU 驱动，可以运行以下命令：

sudo systemctl restart nvidia-persistenced

3. 更新 CUDA 和 GPU 驱动

如果出现驱动不兼容或版本问题，可以考虑更新 NVIDIA 驱动程序和 CUDA 工具包。确保 CUDA 驱动版本与深度学习框架支持的版本匹配。

• 查看当前 CUDA 和驱动版本：

nvcc --version  # 查看 CUDA 版本
nvidia-smi      # 查看 GPU 驱动版本

• 更新 NVIDIA 驱动和 CUDA（以 Ubuntu 为例）：

# 更新 NVIDIA 驱动
sudo apt-get update
sudo apt-get install nvidia-driver-XXX  # 根据硬件选择正确的版本

# 更新 CUDA
sudo apt-get install cuda-toolkit-XXX  # 选择合适的版本

4. 清理 CUDA 缓存

PyTorch 和 TensorFlow 等深度学习框架会缓存某些中间数据，如果缓存过多，可能会导致内存问题。你可以尝试清理这些缓存。

• 清理 PyTorch 的缓存：

import torch
torch.cuda.empty_cache()

• 清理 TensorFlow 的缓存：

import tensorflow as tf
tf.keras.backend.clear_session()

5. 调整 PyTorch 配置（如果使用 PyTorch）

如果你使用的是 PyTorch，可能需要调整一些 CUDA 配置以避免该错误。例如，使用 torch.backends.cudnn.enabled 来关闭 cuDNN 的优化，可能有助于解决内存分配问题。

• 禁用 cuDNN 优化：

import torch
torch.backends.cudnn.enabled = False

6. 检查硬件和驱动兼容性

如果以上方法都没有解决问题，可能需要检查显卡的硬件问题或驱动兼容性。尝试将显卡插入不同的插槽，或更换显卡驱动版本。此外，确保你的硬件支持当前版本的 CUDA 和驱动。

7. 设置 CUDA_VISIBLE_DEVICES

有时，多个程序或进程同时使用 GPU，导致 GPU 资源争用。可以通过设置 CUDA_VISIBLE_DEVICES 环境变量来指定某个 GPU 进行任务调度。

• 设置 CUDA_VISIBLE_DEVICES 环境变量（指定 GPU 0）：

export CUDA_VISIBLE_DEVICES=0

这将使得应用只使用 GPU 0，而忽略其他 GPU。

8. 选择适当的 PyTorch 和 CUDA 版本

PyTorch 与 CUDA 版本之间存在一定的兼容性要求。可以通过 PyTorch 官网查看与 CUDA 对应的兼容版本，并确保使用正确的版本。

• 安装兼容版本的 PyTorch 和 CUDA：

# 安装与 CUDA 10.1 兼容的 PyTorch
pip install torch==1.7.0+cu101

三、总结

RuntimeError: CUDA error: CUBLAS_STATUS_ALLOC_FAILED when calling cublasCreate(handle) 错误通常与 GPU 资源分配和 CUDA 环境配置有关。通过检查 GPU 内存使用情况、更新驱动和 CUDA 工具包、清理缓存、禁用 cuDNN 优化以及调整 PyTorch 配置等手段，可以有效解决该问题。如果以上步骤都不能解决，考虑检查硬件或更换 CUDA 和驱动版本。