Attempting to run cuBLAS, but there was no current CUDA context!

Hi,

i want to calculate a jacobian of an linear layer using torch.func.jacrev(). Unfortunately when doing so i get the warning shown in my minimal example below. Why does this warning occur? When I use it in my main code the runtime of different code blocks changes.

import torch

dummy = torch.tensor(1.0).cuda() # dummy to init cuda
print(torch.cuda.is_initialized())

f = torch.nn.Sequential(
        torch.nn.Linear(in_features=16, out_features=16, bias=False),
    ).cuda()

X_k = torch.rand(16).cuda()
H = torch.func.jacrev(f)(X_k)

print(H.shape)

Output Warning:

python3.10/site-packages/torch/autograd/graph.py:769: UserWarning: Attempting to run cuBLAS, but there was no current CUDA context! Attempting to set the primary context... (Triggered internally at ../aten/src/ATen/cuda/CublasHandlePool.cpp:135.)
  return Variable._execution_engine.run_backward(  # Calls into the C++ engine to run the backward pass

@janssenn I have the same issue. Any idea how to fix this?

This warning was introduced in this PR setting the primary context to the current one for multi-threaded use cases which could otherwise fail.

1 Like

@ptrblck, could you clarify how the warning relates to multi-threading?

If I understand correctly, the warning appears even when running the “vanilla” code with a single GPU and a single thread. Do you have any suggestions on adjustments to prevent the warning from appearing?

1 Like

PyTorch uses a separate thread in the backward pass. If no other CUDA operation was used before the cuBLAS call, the warning will be raised as we will set the primary context to the current one.

Hi, I encountered a similar problem when I was running the WGAN-GP code. I have already set the primary context to the current one for multi-threaded use cases, but the problem still exists.(There is no error when I remove the gradient penalty part in the program.)The following is the content of the error:D:\anaconda3\envs\py312\Lib\site-packages\torch\autograd\graph.py:825: UserWarning: Attempting to run cuBLAS, but there was no current CUDA context! Attempting to set the primary context… (Triggered internally at C:\cb\pytorch_1000000000000\work\aten\src\ATen\cuda\CublasHandlePool.cpp:135.)
return Variable._execution_engine.run_backward( # Calls into the C++ engine to run the backward pass Could you give me some advice, thanks!!!

It’s not an error but an expected warning raised once if the primary context in a new thread wasn’t set already.

1 Like