vision: ops.deform_conv2d causes CUDA illegal memory access

The following script produces an error but not the exact error from above (although it appears related)

torchvision_bug.py

from typing import Any

import torch
from torch.nn import Conv2d, LeakyReLU
from torch.nn.modules.utils import _pair
from torchvision.ops import DeformConv2d, deform_conv2d


class ModulatedDeformConvPack(DeformConv2d):
    def __init__(self, *args, offset_channels, deformable_groups: int = 1, post_activation: Any = None, **kwargs):
        super(ModulatedDeformConvPack, self).__init__(*args, **kwargs)

        self.deformable_groups = deformable_groups

        self.conv_offset = Conv2d(
            offset_channels,
            self.deformable_groups * 3 * self.kernel_size[0] * self.kernel_size[1],
            kernel_size=self.kernel_size,
            stride=_pair(self.stride),
            padding=_pair(self.padding),
            dilation=_pair(self.dilation),
            bias=True,
        )

        self.init_offset()

        self.post_activation = post_activation()

    def init_offset(self):
        self.conv_offset.weight.data.zero_()
        self.conv_offset.bias.data.zero_()

    def forward(self, x, m):
        out = self.conv_offset(m)
        o1, o2, mask = torch.chunk(out, 3, dim=1)
        offset = torch.cat((o1, o2), dim=1)
        mask = torch.sigmoid(mask)
        res = deform_conv2d(input=x, offset=offset, weight=self.weight, stride=_pair(self.stride), padding=_pair(self.padding), dilation=_pair(self.dilation), mask=mask)

        if self.post_activation is not None:
            res = self.post_activation(res)

        return res


device = torch.device("cuda")

bs = 23

test_in = torch.randn(bs, 160, 256, 256, device=device)
test_offset = torch.randn(bs, 64, 256, 256, device=device)
dfc = ModulatedDeformConvPack(160, 64, 3, 1, 1, offset_channels=64, deformable_groups=8, post_activation=LeakyReLU).to(device)

print(f"Total memory used before DFC call: {(torch.cuda.max_memory_allocated() / torch.cuda.get_device_properties(device).total_memory) * 100}%")

test_out = dfc(test_in, test_offset)

print(f"Total memory used after DFC call: {(torch.cuda.max_memory_allocated() / torch.cuda.get_device_properties(device).total_memory) * 100}%")

print(test_out.cpu()[0, 0, 0, 0])

To reproduce run:

CUDA_LAUNCH_BLOCKING=1 python torchvision_bug.py

error message this time is:

error in deformable_im2col: an illegal memory access was encountered
Traceback (most recent call last):
  File "torchvision_bug.py", line 56, in <module>
    test_out = dfc(test_in, test_offset)
  File "/home/mehrlich/.cache/pypoetry/virtualenvs/fvr-_yTNPG6U-py3.8/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1051, in _call_impl
    return forward_call(*input, **kwargs)
  File "torchvision_bug.py", line 38, in forward
    res = deform_conv2d(input=x, offset=offset, weight=self.weight, stride=_pair(self.stride), padding=_pair(self.padding), dilation=_pair(self.dilation), mask=mask)
  File "/home/mehrlich/.cache/pypoetry/virtualenvs/fvr-_yTNPG6U-py3.8/lib/python3.8/site-packages/torchvision/ops/deform_conv.py", line 89, in deform_conv2d
    return torch.ops.torchvision.deform_conv2d(
RuntimeError: CUDA error: CUBLAS_STATUS_NOT_INITIALIZED when calling `cublasCreate(handle)`

Note the variable bs. On my test machine I can reproduce at bs=23 or higher, i.e., bs=22 works fine, bs=23 reproduces the error.

I am using a Tesla V100 for testing, this may be related to how much memory the card has, which is why I’m printing the max allocated memory as reported by pytorch. Note that in the working case (22) it’s only ~38% on my card so there should be plenty of headroom, although I dont know for sure how accurate that is.

pytorch version: 1.9.0 torchvision version: 0.10.0 cuda version (as reported by nvidia-smi let me know if thats not right): 11.2

Why is kMaxGridDim set to 65K? Maximum gridDim.x is 2**31, it can also be queried from device properties rather than hardcoded.

int grid_x = std::min<int>(
                at::cuda::getCurrentDeviceProperties()->maxGridSize[0],
                cuda::ATenCeilDiv(safe_downcast<int, int64_t>(outputWidth), block_y*BLOCK_STRIDE));