BERT文本分类推理时CUDA环境下torch.matmul报CUBLAS_STATUS_INVALID_VALUE错误

问题详情 #

我在做BERT文本分类任务，模型微调完成后要在多台机器上做推理，大部分环境都正常，但某特定环境一直报CUDA error: CUBLAS_STATUS_INVALID_VALUE错误。所有环境的Python、torch、numpy、pytorch-transformers版本完全一致：

• Python 3.9.16
• torch 1.13.0
• pytorch-transformers 1.2.0

更奇怪的是，CPU环境下一切正常，只有CUDA环境会报错，排查后发现错误出在torch.matmul(tensor1, tensor2)操作：

CPU环境执行正常 #

输出torch.Size([3])：

device = 'cpu'
tensor1 = torch.randn(3, 4).to(device)
tensor2 = torch.randn(4).to(device)
torch.matmul(tensor1, tensor2).size()

CUDA环境执行报错 #

RuntimeError: CUDA error: CUBLAS_STATUS_INVALID_VALUE when calling cublasSgemv(handle, op, m, n, &alpha, a, lda, x, incx, &beta, y, incy)

device = 'cuda:0'
tensor1 = torch.randn(3, 4).to(device)
tensor2 = torch.randn(4).to(device)
torch.matmul(tensor1, tensor2).size()

排查与解决思路 #

• 检查CUDA驱动兼容性：torch 1.13.0对CUDA驱动有版本要求（推荐>=470.57.02对应CUDA11.6，或>=510.39.01对应CUDA11.7），确认该机器的驱动版本是否达标，不匹配则升级/降级驱动。
• 清理CUDA缓存并验证设备状态：先执行torch.cuda.empty_cache()清理缓存，再用torch.cuda.is_available()、torch.cuda.device_count()确认CUDA设备是否被正常识别，排除硬件或驱动初始化异常。
• 统一张量数据类型：强制将参与运算的张量转为同一数据类型（比如float32），避免混合精度导致的CUBLAS异常：

  tensor1 = torch.randn(3, 4).float().to(device)
  tensor2 = torch.randn(4).float().to(device)
  

• 禁用CUBLAS优化选项：部分老显卡或特殊环境下，Tensor Core优化可能触发错误，尝试添加环境变量或代码设置：
  • 终端设置：export CUBLAS_WORKSPACE_CONFIG=:4096:8
  • 代码开头添加：

    torch.backends.cudnn.benchmark = False
    torch.backends.cuda.matmul.allow_tf32 = False
    

• 重新编译PyTorch（极端情况）：如果上述方法都无效，可能是该机器的PyTorch二进制包存在编译问题，尝试从源码编译对应版本的PyTorch。

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