Can not compile on cuda:1 on multi-gpus

[alibool]

The following code works well with 'cuda:0', but will raise error when use 'cuda:1'

import torch
from torchvision import models
import torch_tensorrt


dtype = torch.float32
# works fine when cuda:0 but error on cuda:1
device = torch.device("cuda:0")

model = models.resnet50().to(dtype).to(device)
model.eval()
inputs_1 = torch.rand(12, 3, 224, 224).to(device).to(dtype)

optimized_model = torch.compile(model, backend="torch_tensorrt", dynamic=False, fullgraph=True,
                                options={
                                    "precision": dtype,
                                    "device": device
                                })

with torch.no_grad():
    useless = optimized_model(inputs_1)
print(useless)

Two gpus are both A100, pytorch 2.2.1 cu118 torch-tensorrt 2.2.0+cu118 tensorrt 8.6.0

[gs-olive]

Hi - thanks for the report - could you share the error logs when compiling with cuda:1, preferably with "debug": True?

[gs-olive]

As an update on this, I am seeing the following output when running your script:

ERROR: [Torch-TensorRT] - 1: [cudaResources.cpp::~ScopedCudaEvent::24] Error Code 1: Cuda Runtime (an illegal memory access was encountered)
ERROR: [Torch-TensorRT] - 1: [defaultAllocator.cpp::deallocate::61] Error Code 1: Cuda Runtime (an illegal memory access was encountered)

It seems to appear because the active CUDA device in the context is being overwritten or incorrectly set. I am looking into the issue. For a temporary workaround, consider using the prefix argument CUDA_VISIBLE_DEVICES=1 when running the script, to have cuda:1 act as cuda:0 from within the script

[gs-olive]

I have some additional details to share on this, as well as the proposed fix.

In the initial code, the device parameter was set to "cuda:1", but the ambient active CUDA device in the thread was the default ("cuda:0"). In order to have the active thread use "cuda:1", it is necessary to add the line device = torch.device("cuda:1"); torch.cuda.set_device(device) before compilation. This makes the compilation and inference functional on my end.

If the intended usecase was to compile/run on multiple GPUs, I would suggest using separate Python threads, each with a set_device call to ensure the ambient device in that thread is the intended one. Alternatively, you could enable torch_tensorrt.runtime.set_multi_device_safe_mode(True) - see below.

We recently added a feature, Multi-Device Runtime Safety, which gives more verbose messages regarding device contexts and will automatically switch devices if a mismatch is detected, at the cost of runtime device-checking. This can be enabled via torch_tensorrt.runtime.set_multi_device_safe_mode(True). The default is False.

This code works on my end, please let me know if it also resolves your case:

import torch
from torchvision import models
import torch_tensorrt


dtype = torch.float32
device = torch.device("cuda:1")
torch.cuda.set_device(device)

model = models.resnet50().to(dtype).to(device)
model.eval()
inputs_1 = torch.rand(12, 3, 224, 224).to(device).to(dtype)

optimized_model = torch.compile(model, backend="torch_tensorrt", dynamic=False, fullgraph=True,
                                options={
                                    "precision": dtype,
                                    "device": device
                                })

with torch.no_grad():
    useless = optimized_model(inputs_1)
print(useless)

[alibool]

Thanks

I have some additional details to share on this, as well as the proposed fix.

In the initial code, the device parameter was set to "cuda:1", but the ambient active CUDA device in the thread was the default ("cuda:0"). In order to have the active thread use "cuda:1", it is necessary to add the line device = torch.device("cuda:1"); torch.cuda.set_device(device) before compilation. This makes the compilation and inference functional on my end.

If the intended usecase was to compile/run on multiple GPUs, I would suggest using separate Python threads, each with a set_device call to ensure the ambient device in that thread is the intended one. Alternatively, you could enable torch_tensorrt.runtime.set_multi_device_safe_mode(True) - see below.

We recently added a feature, Multi-Device Runtime Safety, which gives more verbose messages regarding device contexts and will automatically switch devices if a mismatch is detected, at the cost of runtime device-checking. This can be enabled via torch_tensorrt.runtime.set_multi_device_safe_mode(True). The default is False.

This code works on my end, please let me know if it also resolves your case:

import torch
from torchvision import models
import torch_tensorrt


dtype = torch.float32
device = torch.device("cuda:1")
torch.cuda.set_device(device)

model = models.resnet50().to(dtype).to(device)
model.eval()
inputs_1 = torch.rand(12, 3, 224, 224).to(device).to(dtype)

optimized_model = torch.compile(model, backend="torch_tensorrt", dynamic=False, fullgraph=True,
                                options={
                                    "precision": dtype,
                                    "device": device
                                })

with torch.no_grad():
    useless = optimized_model(inputs_1)
print(useless)

Thanks, when will be this nice feature released and should I close this issue?

[gs-olive]

Hi - this feature should already be in the version you are using. Once you have verified the resolution works for you, feel free to close this bug and open a new one if there are other issues.