dcnv4 can't run on cuda1!!!!!!

[alioth-k]

I installed it through pip install dcnv4
here is my code

import os
os.environ["CUDA_LAUNCH_BLOCKING"]='1'
os.environ['TORCH_USE_CUDA_DSA']="1"
import torch
from torch import nn
from DCNv4.modules.dcnv4 import DCNv4
class DeformableConv(nn.Module):
    default_act = nn.GELU
    def __init__(self, c1, c2, k=3, s=1, p=1, g=4, d=1, act=True):
        super().__init__()
        self.conv = nn.Conv2d(c1, c2, 1) if c1!=c2 else nn.Identity()
        self.dcn = DCNv4(c2, kernel_size=3, stride=s,pad=1, group=g, dilation=d)
        self.bn = nn.BatchNorm2d(c2)
        self.act = self.default_act() if act else nn.Identity()
    def forward(self, x):
        x = self.conv(x)
        b, c, h, w = x.shape
        x = x.permute(0, 2, 3, 1).view(b, h * w, c)
        x = self.dcn(x, shape=(h, w))
        x = x.view(b, h, w, c).permute(0, 3, 1, 2) # dcnv4
        return self.act(self.bn(x))


def cuda_test(cuda_idx):
    x = torch.rand(1,64,160,160).cuda(cuda_idx)
    dc = DeformableConv(64,64, k=3, g=4).cuda(cuda_idx)
    print(dc(x).shape)
    del x, dc

if __name__ == "__main__":
    cuda_test(0)
    cuda_test(1)

and the output is

(py312) zhouyufei@ubuntu:~/ultralytics$ bash debug.sh seek.py 
torch.Size([1, 64, 160, 160])
error in dcnv4_im2col_cuda: an illegal memory access was encountered
launch arguments: gridDim=(400, 1, 1), blockDim=(2, 4, 64), shm_size=0

Traceback (most recent call last):
  File "<frozen runpy>", line 198, in _run_module_as_main
  File "<frozen runpy>", line 88, in _run_code
  File "/home/zhouyufei/anaconda3/envs/py312/lib/python3.12/site-packages/debugpy/__main__.py", line 71, in <module>
    cli.main()
  File "/home/zhouyufei/anaconda3/envs/py312/lib/python3.12/site-packages/debugpy/server/cli.py", line 501, in main
    run()
  File "/home/zhouyufei/anaconda3/envs/py312/lib/python3.12/site-packages/debugpy/server/cli.py", line 351, in run_file
    runpy.run_path(target, run_name="__main__")
  File "/home/zhouyufei/anaconda3/envs/py312/lib/python3.12/site-packages/debugpy/_vendored/pydevd/_pydevd_bundle/pydevd_runpy.py", line 310, in run_path
    return _run_module_code(code, init_globals, run_name, pkg_name=pkg_name, script_name=fname)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/zhouyufei/anaconda3/envs/py312/lib/python3.12/site-packages/debugpy/_vendored/pydevd/_pydevd_bundle/pydevd_runpy.py", line 127, in _run_module_code
    _run_code(code, mod_globals, init_globals, mod_name, mod_spec, pkg_name, script_name)
  File "/home/zhouyufei/anaconda3/envs/py312/lib/python3.12/site-packages/debugpy/_vendored/pydevd/_pydevd_bundle/pydevd_runpy.py", line 118, in _run_code
    exec(code, run_globals)
  File "seek.py", line 32, in <module>
    cuda_test(1)
  File "seek.py", line 27, in cuda_test
    print(dc(x).shape)
          ^^^^^
  File "/home/zhouyufei/anaconda3/envs/py312/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1511, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/zhouyufei/anaconda3/envs/py312/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1520, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "seek.py", line 19, in forward
    x = self.dcn(x, shape=(h, w))
        ^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/zhouyufei/anaconda3/envs/py312/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1511, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/zhouyufei/anaconda3/envs/py312/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1520, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/zhouyufei/anaconda3/envs/py312/lib/python3.12/site-packages/DCNv4/modules/dcnv4.py", line 130, in forward
    x = DCNv4Function.apply(
        ^^^^^^^^^^^^^^^^^^^^
  File "/home/zhouyufei/anaconda3/envs/py312/lib/python3.12/site-packages/torch/autograd/function.py", line 553, in apply
    return super().apply(*args, **kwargs)  # type: ignore[misc]
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/zhouyufei/anaconda3/envs/py312/lib/python3.12/site-packages/torch/cuda/amp/autocast_mode.py", line 115, in decorate_fwd
    return fwd(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^
  File "/home/zhouyufei/anaconda3/envs/py312/lib/python3.12/site-packages/DCNv4/functions/dcnv4_func.py", line 103, in forward
    output = ext.dcnv4_forward(*args)
             ^^^^^^^^^^^^^^^^^^^^^^^^
RuntimeError: false INTERNAL ASSERT FAILED at "/tmp/pip-install-ac16vv1o/dcnv4_169f40bf1e344ced97039510708b25ee/src/cuda/dcnv4_im2col_cuda.cuh":323, please report a bug to PyTorch. kernel launch error

[libra0037]

same issue here, try torch.cuda.set_device(1)

[gtebbutt]

Wrapping the DCNv4 forward call in with torch.cuda.device(...) seems to work cleanly, and doesn't modify global state (the torch docs recommend against using set_device).

I'm not sure exactly how the two differ, but with torch.device(...) does not work for me, even if passed a CUDA device. It has to explicitly be the torch.cuda version.

In this example, the following should do the job:

def cuda_test(cuda_idx):
    x = torch.rand(1,64,160,160).cuda(cuda_idx)
    dc = DeformableConv(64,64, k=3, g=4).cuda(cuda_idx)
    with torch.cuda.device(x.device.index):
        print(dc(x).shape)
    del x, dc

Wrapping the calls in a device context manager like that is definitely just a workaround, though, it doesn't address whatever's causing the overflow.

If anyone has the time to look at fixing the root cause and updating the CUDA code, it looks like there was a very similar issue in Deformable DETR: open-mmlab/mmdetection#7186 that was fixed in open-mmlab/mmcv#1823