CUBLAS_STATUS_EXECUTION_FAILED when calling cublasSgemm 解决方法

问题描述

在运行pytorch项目的时候遇到了这个问题
RuntimeError: CUDA error: CUBLAS_STATUS_EXECUTION_FAILED when calling cublasSgemm( handle, opa, opb, m, n, k, &alpha, a, lda, b, ldb, &beta, c, ldc)
简单描述一下我遇到的情况：
在运行Bert的SelfAttention操作时出错，定位在一个nn.Linear操作上。我在CPU上可以成功运行，但是在GPU上运行就会产生这个错误。所以怀疑是版本的问题。
直到我看见这个网页，https://discuss.pytorch.org/t/cublas-status-execution-failed-when-calling-cublassgemm-handle-opa-opb-m-n-k-alpha-a-lda-b-ldb-beta-c-ldc/116740/9

然后我发现这个是一个普遍的错误。
我的版本配置是python=3.8，torch=1.8，cuda=11.1，这个是cuda=11.1时会出现的一个错误。
假如你还无法确定自己是不是跟我有同样的问题，可以运行下面的实例代码，如果报错，说明问题是一样的：

import os
os.environ['CUDA_VISIBLE_DEVICES']='0'
import torch

class Classifier(torch.nn.Module):
    def __init__(self):
        super().__init__()
        #classifier
        self.linear_1 = torch.nn.Linear(32*8*4*8,32*8*4)
        self.linear_2 = torch.nn.Linear(32*8*4,32*8)
        self.batch_norm_1 = torch.nn.BatchNorm1d(32*8)

        self.linear_3 = torch.nn.Linear(32*8,32)
        self.linear_4 = torch.nn.Linear(32,1)

        self.activation = torch.nn.ReLU()


    def forward(self,x):
        x = x.reshape(x.shape[0], -1)
        x = self.linear_1(x)
        x = self.activation(x)
        x = self.linear_2(x)
        x = self.activation(x)
        x = self.batch_norm_1(x)

        x = self.linear_3(x)
        x = self.activation(x)
        x = self.linear_4(x)
        return x
    
net = Classifier().cuda()
inp = torch.rand(2,32,8,4,8).cuda()
gt = torch.rand(2,1).cuda()
outp = net(inp)
loss = torch.nn.BCEWithLogitsLoss()(outp, gt)
loss.backward()

解决方案

在我把cuda版本降为10.2之后就能成功运行代码了。
cuda安装指令：

conda install pytorch torchvision torchaudio cudatoolkit=10.2 -c pytorch-lts