PyTorch训练GeoFormer时出现CUBLAS_STATUS_EXECUTION_FAILED错误求助

环境配置 #

• Pytorch: 1.9.0
• GPU: RTX3090
• CUDA: 11.1
• Python: 3.8
• Linux版本: Ubuntu 16.04

错误信息 #

运行GeoFormer训练代码时触发CUDA错误：CUDA error: CUBLAS_STATUS_EXECUTION_FAILED when calling cublasGemmEx(...)

错误栈：

Traceback (most recent call last):
  File "train_fs.py", line 222, in <module>
    train_one_epoch(epoch, train_loader, model, criterion, optimizer)
  File "train_fs.py", line 65, in train_one_epoch
    outputs = model(support_dict, query_dict, remember=False, training=True)
  File "/home/zpf/anaconda3/envs/py38/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1051, in _call_impl
    return forward_call(*input, **kwargs)
  File "/home/zpf/work/test_hwj/GeoFormer/model/geoformer/geoformer_fs.py", line 546, in forward
    dec_outputs = self.forward_decoder(
  File "/home/zpf/work/test_hwj/GeoFormer/model/geoformer/geoformer_fs.py", line 719, in forward_decoder
    dec_outputs = self.decoder(
  File "/home/zpf/anaconda3/envs/py38/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1051, in _call_impl
    return forward_call(*input, **kwargs)
  File "/home/zpf/work/test_hwj/GeoFormer/model/transformer_detr.py", line 140, in forward
    output, attn = layer(
  File "/home/zpf/anaconda3/envs/py38/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1051, in _call_impl
    return forward_call(*input, **kwargs)
  File "/home/zpf/work/test_hwj/GeoFormer/model/transformer_detr.py", line 484, in forward
    return self.forward_pre_rel(
  File "/home/zpf/work/test_hwj/GeoFormer/model/transformer_detr.py", line 460, in forward_pre_rel
    tgt = self.out_mlp(tgt)
  File "/home/zpf/anaconda3/envs/py38/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1051, in _call_impl
    return forward_call(*input, **kwargs)
  File "/home/zpf/anaconda3/envs/py38/lib/python3.8/site-packages/torch/nn/modules/container.py", line 139, in forward
    input = module(input)
  File "/home/zpf/anaconda3/envs/py38/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1051, in _call_impl
    return forward_call(*input, **kwargs)
  File "/home/zpf/anaconda3/envs/py38/lib/python3.8/site-packages/torch/nn/modules/linear.py", line 96, in forward
    return F.linear(input, self.weight, self.bias)
  File "/home/zpf/anaconda3/envs/py38/lib/python3.8/site-packages/torch/nn/functional.py", line 1847, in linear
    return torch._C._nn.linear(input, weight, bias)
RuntimeError: CUDA error: CUBLAS_STATUS_EXECUTION_FAILED when calling `cublasGemmEx( handle, opa, opb, m, n, k, &falpha, a, CUDA_R_16F, lda, b, CUDA_R_16F, ldb, &fbeta, c, CUDA_R_16F, ldc, CUDA_R_32F, CUBLAS_GEMM_DFALT_TENSOR_OP)`

相关网络代码 #

class TransformerDecoderLayer(nn.Module):
    def __init__(
        self,
        d_model,
        nhead=4,
        dim_feedforward=256,
        dropout=0.1,
        dropout_attn=None,
        activation="relu",
        normalize_before=True,
        use_rel=False,
        norm_fn_name="ln",
    ):
        super().__init__()
        if dropout_attn is None:
            dropout_attn = dropout
        self.self_attn = nn.MultiheadAttention(d_model, nhead, dropout=dropout)

        self.norm1 = NORM_DICT[norm_fn_name](d_model)
        self.norm2 = NORM_DICT[norm_fn_name](d_model)

        self.norm3 = NORM_DICT[norm_fn_name](d_model)
        self.dropout1 = nn.Dropout(dropout, inplace=True)
        self.dropout2 = nn.Dropout(dropout, inplace=True)
        self.dropout3 = nn.Dropout(dropout, inplace=True)

        # Implementation of Feedforward model
        self.linear1 = nn.Linear(d_model, dim_feedforward)
        self.dropout = nn.Dropout(dropout, inplace=True)
        self.linear2 = nn.Linear(dim_feedforward, d_model)

        self.activation = ACTIVATION_DICT[activation]()
        self.normalize_before = normalize_before

        self.use_rel = use_rel
        self.nhead = nhead
        # print(d_model)
        if self.use_rel:
            self.attn_mlp = nn.Sequential(
                nn.Linear(d_model, d_model),
                nn.ReLU(),
                nn.Linear(d_model, d_model),
            )
            self.v_mlp = nn.Sequential(
                nn.Linear(d_model, d_model),
            )
            self.out_mlp = nn.Sequential(
                # nn.Flatten(),
                nn.Linear(d_model, d_model),
                nn.ReLU(),
            )
        else:
            self.multihead_attn = nn.MultiheadAttention(d_model, nhead, dropout=dropout)

    def with_pos_embed(self, tensor, pos: Optional[Tensor]):
        return tensor if pos is None else tensor + pos

    def _reshape_to_batches(self, x):
        batch_size, seq_len, in_feature = x.size()
        sub_dim = in_feature // self.nhead
        return (
            x.reshape(batch_size, seq_len, self.nhead, sub_dim)
            .permute(0, 2, 1, 3)
            .reshape(batch_size * self.nhead, seq_len, sub_dim)
        )

    def _reshape_from_batches(self, x):
        batch_size, seq_len, in_feature = x.size()
        batch_size //= self.head_num
        out_dim = in_feature * self.nhead
        return (
            x.reshape(batch_size, self.nhead, seq_len, in_feature)
            .permute(0, 2, 1, 3)
            .reshape(batch_size, seq_len, out_dim)
        )

    def forward_pre_rel(
        self,
        tgt,
        memory,
        tgt_mask: Optional[Tensor] = None,
        memory_mask: Optional[Tensor] = None,
        tgt_key_padding_mask: Optional[Tensor] = None,
        memory_key_padding_mask: Optional[Tensor] = None,
        pos: Optional[Tensor] = None,
        query_pos: Optional[Tensor] = None,
        relative_pos: Optional[Tensor] = None,
        return_attn_weights: Optional[bool] = False,
    ):

        # NOTE self attn between queries themself: use absolute euclid pos
        tgt2 = self.norm1(tgt)
        q = k = self.with_pos_embed(tgt2, query_pos)
        tgt2 = self.self_attn(q, k, value=tgt2, attn_mask=tgt_mask, key_padding_mask=tgt_key_padding_mask)[0]
        tgt = tgt + self.dropout1(tgt2)
        tgt2 = self.norm2(tgt)

        # NOTE cross attn between queries and contexst: use relative pos
        n_queries, n_context, batch, channel = relative_pos.shape
        tgt2_expand = tgt2[:, None, :, :].repeat(1, n_context, 1, 1)

        memory_expand = memory[None, :, :, :].repeat(n_queries, 1, 1, 1)

        sim = self.attn_mlp(tgt2_expand - memory_expand + relative_pos)
        attn = F.softmax(sim / np.sqrt(sim.shape[-1]), dim=1)

        v2 = self.v_mlp(memory_expand + relative_pos)
        tgt = torch.einsum("qcbf,qcbf->qbf", attn, v2)  # n_queries, batch, channel
        tgt = self.out_mlp(tgt)
        ###########################################################################################

        tgt = tgt + self.dropout2(tgt2)
        tgt2 = self.norm3(tgt)
        tgt2 = self.linear2(self.dropout(self.activation(self.linear1(tgt2))))
        tgt = tgt + self.dropout3(tgt2)
        if return_attn_weights:
            return tgt, attn
        return tgt, None

    def forward(
        self,
        tgt,
        memory,
        tgt_mask: Optional[Tensor] = None,
        memory_mask: Optional[Tensor] = None,
        tgt_key_padding_mask: Optional[Tensor] = None,
        memory_key_padding_mask: Optional[Tensor] = None,
        pos: Optional[Tensor] = None,
        query_pos: Optional[Tensor] = None,
        relative_pos: Optional[Tensor] = None,
        return_attn_weights: Optional[bool] = False,
    ):
        return self.forward_pre_rel(
            tgt,
            memory,
            tgt_mask,
            memory_mask,
            tgt_key_padding_mask,
            memory_key_padding_mask,
            pos,
            query_pos,
            relative_pos,
            return_attn_weights,
        )

已尝试的解决方法 #

• 将PyTorch版本从1.10.0降级到1.9.0
• 确认输入张量尺寸为[256,1,64]，对应线性层[64,64]
• 尝试用Flatten处理3D张量，但问题未解决

疑问 #

这是CUDA 11.1的问题吗？

解答 #

这个错误不一定是CUDA 11.1本身的问题，可从以下维度排查：

1. 张量精度与Tensor Core兼容性 #

错误日志显示使用了CUBLAS_GEMM_DFALT_TENSOR_OP，RTX3090的Tensor Core对FP16操作的形状、精度要求严格：

• 若代码启用了torch.cuda.amp自动混合精度，暂时关闭测试
• 在out_mlp前强制转换张量为FP32：tgt = tgt.float()

2. 张量形状隐性问题 #

• 打印einsum操作后的tgt.shape，确认是[256,1,64]再传入out_mlp
• 检查relative_pos的n_queries、n_context是否与tgt、memory的序列长度匹配，广播不匹配可能导致异常维度

3. 版本兼容性优化 #

PyTorch 1.9.0兼容CUDA 11.1，但Ampere架构的RTX3090在旧CUDA版本中可能存在Tensor Core优化bug：

• 升级CUDA到11.3/11.4（PyTorch 1.9.0支持这些版本）
• 重新升级PyTorch到1.10.0+，新版本对Ampere架构兼容性更好

4. 硬件驱动检查 #

• 用nvidia-smi确认驱动版本，CUDA 11.1要求驱动≥455.23，版本过低会导致Ampere GPU操作异常，需升级驱动

5. 代码笔误修复 #

_reshape_from_batches方法中存在笔误：batch_size //= self.head_num，类内定义的是self.nhead，应改为self.nhead，该错误可能导致张量形状异常，引发后续CUBLAS错误。

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