Triton Error [CUDA]: invalid device context when autograd.backward a triton kernel

[ppwwyyxx]

🐛 Describe the bug

import torch
import triton
import triton.language as tl

class Rua(torch.autograd.Function):

    @staticmethod
    def rua(hidden):
        output = torch.empty_like(hidden)
        _rua_kernel[(1,)](hidden)
        return output

    @staticmethod
    def forward(ctx, hidden):
        return Rua.rua(hidden)

    @staticmethod
    def backward(ctx, output_grad):
        return Rua.rua(output_grad)

@triton.jit
def _rua_kernel(hidden_sh_ptr):
    return

if __name__ == "__main__":

    input_data = torch.empty(1, device='cuda', requires_grad=True)
    output_grad = torch.empty(1, device='cuda')

    output = Rua.apply(input_data)
    # torch.empty(1, device='cuda', requires_grad=True).backward()
    torch.autograd.backward(output, output_grad)
    torch.cuda.synchronize()

Prints:

Traceback (most recent call last):
  File "/home/X/a.py", line 33, in <module>
    torch.autograd.backward(output, output_grad)
  File "/home/X/.local/lib/python3.11/site-packages/torch/autograd/__init__.py", line 266, in backward
    Variable._execution_engine.run_backward(  # Calls into the C++ engine to run the backward pass
  File "/home/X/.local/lib/python3.11/site-packages/torch/autograd/function.py", line 289, in apply
    return user_fn(self, *args)
           ^^^^^^^^^^^^^^^^^^^^
  File "/home/X/a.py", line 20, in backward
    return Rua.rua(output_grad)
           ^^^^^^^^^^^^^^^^^^^^
  File "/home/X/a.py", line 11, in rua
    _rua_kernel[(1,)](hidden)
  File "/home/X/.local/lib/python3.11/site-packages/triton/runtime/jit.py", line 550, in run
    bin.c_wrapper(
RuntimeError: Triton Error [CUDA]: invalid device context

If I uncomment the commented line in the code, it runs correctly.

Versions

PyTorch version: 2.2.2+cu121
Is debug build: False
CUDA used to build PyTorch: 12.1
ROCM used to build PyTorch: N/A

OS: Arch Linux (x86_64)
GCC version: (GCC) 13.2.1 20230801
Clang version: 17.0.6
CMake version: version 3.29.2
Libc version: glibc-2.39

Python version: 3.11.8 (main, Feb 12 2024, 14:50:05) [GCC 13.2.1 20230801] (64-bit runtime)
Python platform: Linux-6.6.27-1-lts-x86_64-with-glibc2.39
Is CUDA available: True
CUDA runtime version: 12.4.131
CUDA_MODULE_LOADING set to: LAZY
GPU models and configuration: GPU 0: NVIDIA GeForce RTX 3080
Nvidia driver version: 550.67

cc @ezyang @albanD @gqchen @pearu @nikitaved @soulitzer @Varal7 @ptrblck @msaroufim @bertmaher @int3 @davidberard98 @nmacchioni @lezcano

[albanD]

Does triton have thread local context that they need preserved when running from different threads?

[bertmaher]

Weird, if I add a print(output) right after output = torch.empty_like(hidden), the program runs to completion :-o

[albanD]

See the triton issue I just opened.
Anything we do that properly initialize cuda in the autograd thread (printing or running another backward) will fix this issue indeed.

I would suggest that keeping "torch.empty(1, device='cuda', requires_grad=True).backward()" at the top of your program would be a cheap-enough workaround while we work on a fix with triton.
Or if we need someone more urgent, I'm happy to add that directly in the autograd engine if we think this is a big issue.

[titu1994]

Can reproduce with some trivial op compilation too - remove the 2.0 and it will work without the hack, add the 2.0 and it will crash the grad kernel with the above error (hack fixes it perfectly). This is with Pytorch 2.3 on Ubuntu 20.04, CUDA 12.1 (conda install)

import torch

# Hack for a triton bug - https://github.com/pytorch/pytorch/issues/124565
# torch.empty(1, device='cuda', requires_grad=True).backward()

def add(a, b):
    return 2.0 * a + b

def forward_backward(x, y):
    c = add(x, y)
    c.backward(torch.ones_like(c))
    return c

op = add
op_grad = torch.compile(forward_backward, dynamic=True)

x = torch.ones(2, 3, device="cuda", requires_grad=True) * 3
y = torch.ones(2, 3, device="cuda") * 2

z = op(x, y)
_ = op_grad(x, y)

print(z)

System info in case it helps

Collecting environment information...
PyTorch version: 2.3.0
Is debug build: False
CUDA used to build PyTorch: 12.1
ROCM used to build PyTorch: N/A

OS: Ubuntu 20.04.1 LTS (x86_64)
GCC version: (Ubuntu 9.4.0-1ubuntu1~20.04.2) 9.4.0
Clang version: Could not collect
CMake version: Could not collect
Libc version: glibc-2.31

Python version: 3.10.14 (main, May  6 2024, 19:42:50) [GCC 11.2.0] (64-bit runtime)
Python platform: Linux-5.10.102.1-microsoft-standard-WSL2-x86_64-with-glibc2.31
Is CUDA available: True
CUDA runtime version: Could not collect
CUDA_MODULE_LOADING set to: LAZY
GPU models and configuration: GPU 0: NVIDIA GeForce RTX 2080 with Max-Q Design
Nvidia driver version: 546.17
cuDNN version: Could not collect
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

[yiliu30]

It seems like the issue was addressed in triton-lang/triton#3731, but the fix is not included in torch 2.5.0.dev20240812.

[BrokenArrow1404]

thanks for the solution
# torch.empty(1, device='cuda', requires_grad=True).backward()
this works for me