t-vi
commented
3 weeks ago More minimal repro:
def fn(a):
return a.type('torch.DoubleTensor')
a = torch.randn(3)
jfn = thunder.jit(fn, executors=(thunder.executors.torch_compile.torch_compile_ex,))
jfn(a)The trace before transform for execution is
# Constructed by Dead Code Elimination (took 0 milliseconds)
import thunder
import thunder.torch as ltorch
import torch
from thunder.executors.torchex import no_autocast
@torch.no_grad()
@no_autocast
def computation(a):
# a: "cpu f32[2, 2]"
# <ipython-input-18-c7e8b0e10321>:2: return a.type('torch.DoubleTensor')
t0 = ltorch.type(a, 'torch.DoubleTensor', False) # t0: "cpu f64[2, 2]"
# t0 = ltorch.to(a, devices.Device("cpu"), torch.float64, device=None, dtype=None, copy=False, memory_format=None) # t0: "cpu f64[2, 2]"
# t0 = prims.convert_element_type(a, dtypes.float64) # t0: "cpu f64[2, 2]"
return t0but the torchex executor knows to resolve the type to the ltorch.to subsymbol whereas the torch_compile_ex tries to use the type symbol directly:
After transform to execution with the torchex:
# Constructed by Transform for execution (took 0 milliseconds)
from torch import Tensor
import torch
from thunder.executors.torchex import no_autocast
@torch.no_grad()
@no_autocast
def computation(a):
# a: "cpu f32[2, 2]"
t0 = Tensor.to(a, copy=False, device=torch.device("cpu"), dtype=torch.float64) # t0: "cpu f64[2, 2]"
# t0 = ltorch.to(a, None, None, device=torch.device("cpu"), dtype=torch.float64, copy=False, memory_format=None) # t0: "cpu f64[2, 2]"
# t0 = prims.convert_element_type(a, dtypes.float64) # t0: "cpu f64[2, 2]"
# <ipython-input-18-c7e8b0e10321>:2: return a.type('torch.DoubleTensor')
return t0So back in the day, @IvanYashchuk (?) added a comment:
Now I seem to get good results with just calling transform_for_execution with only the torchex as executors: PR #1041 does seem to work.
🐛 Bug
We appear to be missing a key in a map internal to the torch.compile executor:
To Reproduce
See bug #343 for NeVA setup information.
Code sample
Workaround
There's a workaround to just remove the executor:
Environment
Additional context
cc @apaz-cli @tfogal