Tensor.item should be disallowed in the CUDA Graphs region - Githubissues

Lightning-AI / lightning-thunder

Make PyTorch models up to 40% faster! Thunder is a source to source compiler for PyTorch. It enables using different hardware executors at once; across one or thousands of GPUs.

Apache License 2.0

1.15k stars 77 forks source link

Tensor.item should be disallowed in the CUDA Graphs region #727

Open IvanYashchuk opened 3 months ago

IvanYashchuk commented 3 months ago

Note: If you have a model or program that is not supported yet but should be, please use the program coverage template.

🐛 Bug

import thunder
import torch

def func(x):
    return x.sum().item()

a = torch.randn(3, 3, device="cuda")

jfunc = thunder.jit(func)
jfunc(a) # works!

jfunc = thunder.jit(func, use_cudagraphs=True)
jfunc(a) # fails!

File thunder.CUDAGraph0_fn_3:10, in CUDAGraph0_fn(***failed resolving arguments***)
      9 del x
---> 10 f1 = Tensor.item(t0)  # f1: "float ?"
     11 del t0

RuntimeError: CUDA error: operation not permitted when stream is capturing

tfogal commented 3 months ago

related to #728