RuntimeError: CUDA error: invalid argument

Andymulb: <p>Hi, I’m currently testing various SLAM systems for autonomous driving on embedded hardware (an aarch64-based NVIDIA Jetson AGX Orin with 64 GB of unified memory). I cannot run some of the systems due to a PyTorch multiprocessing error. Here’s the error message I received when testing DROID-Splat:</p>
<p>[Main]: Load pretrained checkpoint from ./pretrained/droid.pth!<br>
/home/DROID-Splat/src/slam.py:260: FutureWarning: You are using <code>torch.load</code> with <code>weights_only=False</code> (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See <a href="https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models" class="inline-onebox" rel="noopener nofollow ugc">pytorch/SECURITY.md at main · pytorch/pytorch · GitHub</a> for more details). In a future release, the default value for <code>weights_only</code> will be flipped to <code>True</code>. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via <code>torch.serialization.add_safe_globals</code>. We recommend you start setting <code>weights_only=True</code> for any use case where you don’t have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.<br>
state_dict = OrderedDict([(k.replace(“module.”, “”), v) for (k, v) in torch.load(pretrained).items()])<br>
[2024-12-09 08:08:06,685][<strong>main</strong>][INFO] - Running on 2000 frames<br>
/home/anaconda3/envs/DROID-Splat/lib/python3.10/site-packages/torchvision/io/image.py:13: UserWarning: Failed to load image Python extension:<br>
warn(f"Failed to load image Python extension: {e}")<br>
Traceback (most recent call last):<br>
File “”, line 1, in<br>
File “/home/anaconda3/envs/DROID-Splat/lib/python3.10/multiprocessing/spawn.py”, line 116, in spawn_main<br>
exitcode = _main(fd, parent_sentinel)<br>
File “/home/anaconda3/envs/DROID-Splat/lib/python3.10/multiprocessing/spawn.py”, line 126, in _main<br>
self = reduction.pickle.load(from_parent)<br>
File “/home/anaconda3/envs/DROID-Splat/lib/python3.10/site-packages/torch/multiprocessing/reductions.py”, line 149, in rebuild_cuda_tensor<br>
storage = storage_cls._new_shared_cuda(<br>
File “/home/anaconda3/envs/DROID-Splat/lib/python3.10/site-packages/torch/storage.py”, line 1420, in _new_shared_cuda<br>
return torch.UntypedStorage._new_shared_cuda(*args, **kwargs)<br>
RuntimeError: CUDA error: invalid argument<br>
Compile with <code>TORCH_USE_CUDA_DSA</code> to enable device-side assertions.</p>
<p>I also just tested DROID-Splat on a different x86-based system with an RTX 4080 GPU (CUDA 12.4), and everything works fine apart from the memory boundaries. However, the Jetson AGX Orin has 64 GB has sufficient memory. Do you have an idea what might cause the problem?</p>
<p>On the Jetson PyTorch 2.5.0a0+872d972e41.nv24.08 is installed, which is the only version available that works with CUDA 12.6</p>

NailikLN: <p>Hi ! i have the same problem using torch multiprocessing and spawn method. I pass as an argument torch models (yolo models) and i have more than enough GPU memory,.i also tested it on a an RTX 4080 and it worked, however it was not on arm like the jetson are (i don’t know if it is the problem).<br>
My test were run on docker and my jetson orin NX run on jetpack 6.</p>
<p>Did you manage to fix the issue and if not does someone have an idea ?</p>

Andymulb: <p>Hi <a class="mention" href="/u/nailikln">@NailikLN</a>! I didn’t solve the issue yet. I’m testing other SLAM systems as well and until now, I encountered a similar issue when testing DROID-SLAM, and exactly the same issue when testing NeRF-LOAM. Do you have JetPack 6.1 or JetPack 6.0 and CUDA 12.6.10 or CUDA 12.2.1 installed? If I don’t find another solution I will try to downgrade JetPack to use a earlier CUDA version than 12.6.</p>

NailikLN: <p>I also have jetpack 6.0 and cuda 12.2, let me know if jetpack 5 resolve the issue ! I also tested some others ways but none of them worked (different docker image etc) so i do not see what else it could be…</p>

Doruk_Sonmez: <p>Hi, I’m also having the same problem on NVIDIA Jetson AGX Orin 64GB while using PyTorch 2.5.0a0+872d972e41.nv24.08 and JetPack 6.0 (but inside a JetPack 6.1-based docker container). My workload is related to some LLM/VLM application but the error is pretty much the same:</p>
<pre><code class="lang-auto">Process EmbeddingProcess-2:
Traceback (most recent call last):
  File "/usr/lib/python3.10/multiprocessing/process.py", line 314, in _bootstrap
    self.run()
  File "/opt/nvidia/myapp/process_base.py", line 188, in run
    item = self._queue.get()
  File "/usr/lib/python3.10/multiprocessing/queues.py", line 122, in get
    return _ForkingPickler.loads(res)
  File "/usr/local/lib/python3.10/dist-packages/torch/multiprocessing/reductions.py", line 149, in rebuild_cuda_tensor
    storage = storage_cls._new_shared_cuda(
  File "/usr/local/lib/python3.10/dist-packages/torch/storage.py", line 1420, in _new_shared_cuda
    return torch.UntypedStorage._new_shared_cuda(*args, **kwargs)
RuntimeError: CUDA error: invalid argument
Compile with `TORCH_USE_CUDA_DSA` to enable device-side assertions.
</code></pre>

Andymulb: <p>I finally found the solution and it’s very simple. I tried to install HI-SLAM2 on the Jetson AGX Orin with Jetpack 6.1 and CUDA 12.6 when running the demo, a similar torch multiprocessing error occurred. The solution was to change the compile args of the project’s setup.py and the setup.py files of the submodules as well due to the newer architecture of the Jetson AGX Orin.</p>
<p>Here is an example in HI-SLAM2:</p>
<pre><code class="lang-auto">setup(
    name='droid_backends',
    ext_modules=[
        CUDAExtension('droid_backends',
            include_dirs=[osp.join(ROOT, 'thirdparty/eigen')],
            sources=[
                'src/droid.cpp', 
                'src/droid_kernels.cu',
                'src/correlation_kernels.cu',
                'src/altcorr_kernel.cu',
            ],
            extra_compile_args={
                'cxx': ['-O3'],
                'nvcc': ['-O3',
                    # '-gencode=arch=compute_60,code=sm_60',
                    # '-gencode=arch=compute_61,code=sm_61',
                    # '-gencode=arch=compute_70,code=sm_70',
                    # '-gencode=arch=compute_75,code=sm_75',
                    # '-gencode=arch=compute_80,code=sm_80',
                    # '-gencode=arch=compute_86,code=sm_86',
                    '-gencode=arch=compute_87,code=sm_87', # this is the important change
                ]
            }),
    ],
    cmdclass={ 'build_ext' : BuildExtension }
)
</code></pre>

vixp24: <p>Hi All,</p>
<p>I’ve got the same issue building MASt3R-SLAM on the Jetson Orin AGX running JP6.0 with CUDA 12.2, and using a docker container with Pytorch 2.4.0a0+07cecf4168.nv24.05.</p>
<p><a class="mention" href="/u/andymulb">@Andymulb</a> , I’ve tried your suggestion with only specifying nvcc args applicable for the Jetson only, but I ended up with the same issue! Curious if you had changed anything else? Or if anyone has come across another solution.</p>
<p>I’m currently upgrading to JP6.1, CUDA 12.6 and Pytorch 2.5, to see if that resolves it for me.</p>

vixp24: <p>Hi <a class="mention" href="/u/andymulb">@Andymulb</a>,</p>
<p>I tried again on JP6.1, CUDA 12.6 and Pytorch 2.5, and added the arch code changes in all relevant packages, and I keep getting the same runtime issue.</p>
<p>I came across <a href="https://forums.developer.nvidia.com/t/multiprocessing-pytorch-inference-with-tensorrt-on-jetson-orin-nx-devices/289892" rel="noopener nofollow ugc">this post</a> and some others that mention that <code>torch.multiprocessing</code> is effectively not supported on the Jetson Orin AGX.</p>
<p>I don’t understand how you were able to get past this issue. Can you please share which SLAM system you were able to run successfully? Did it use <code>torch.multiprocessing</code>?</p>
<p>I am hitting the same runtime issue with DroidSLAM and MASt3R-SLAM, unless I get rid of <code>torch.multiprocessing</code></p>
<p>Thanks,</p>

J4Q8: <p>Hi <a class="mention" href="/u/andymulb">@Andymulb</a> and <a class="mention" href="/u/vixp24">@vixp24</a>,</p>
<p>Sadly, Jetson devices do not support CUDA IPC (inter-process communication), which is used by torch.multiprocessing for sharing cuda tensors across processes. I had the same problem in a project of mine, where I wrote custom queue and buffer that are able to share cuda tensors across processes by levaraging low level cuda API (cuMemExportToShareableHandle).</p>
<p>Queue: <a href="https://github.com/nasa-jpl/visual-perception-engine/blob/7f7c56e8c325729c8742def91183f22a0ff089ff/src/vp_engine/cuda_queue.py#L41" class="inline-onebox" rel="noopener nofollow ugc">visual-perception-engine/src/vp_engine/cuda_queue.py at 7f7c56e8c325729c8742def91183f22a0ff089ff · nasa-jpl/visual-perception-engine · GitHub</a><br>
Buffer: <a href="https://github.com/nasa-jpl/visual-perception-engine/blob/7f7c56e8c325729c8742def91183f22a0ff089ff/src/vp_engine/cuda_buffer.py#L36" class="inline-onebox" rel="noopener nofollow ugc">visual-perception-engine/src/vp_engine/cuda_buffer.py at 7f7c56e8c325729c8742def91183f22a0ff089ff · nasa-jpl/visual-perception-engine · GitHub</a></p>
<p>How does it work under the hood (if you are interested): <a href="https://github.com/nasa-jpl/visual-perception-engine/blob/7f7c56e8c325729c8742def91183f22a0ff089ff/src/vp_engine/cuda_utils.py#L429" class="inline-onebox" rel="noopener nofollow ugc">visual-perception-engine/src/vp_engine/cuda_utils.py at 7f7c56e8c325729c8742def91183f22a0ff089ff · nasa-jpl/visual-perception-engine · GitHub</a></p>
<p>I hope that you will find those useful!</p>
<p>Best,</p>