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.14k stars 73 forks source link

duplicate symbols from autoregistration (torch.foo vs. torch.Tensor.foo) #1140

Open t-vi opened 1 week ago

t-vi commented 1 week ago

Autoregistration registers duplicate symbols under the same name, which is not ideal, and we should go the same way as for manually registered symbols (see below).

Repro:

import thunder, torch
def fn(a):
     return torch.positive(a.positive())
jf = thunder.jit(fn)
jf(torch.randn(1))
lt = thunder.last_traces(jf)[0]
print(lt)
s1 = lt.bound_symbols[1].sym
s2 = lt.bound_symbols[2].sym
assert s1 == s2, f"{s1} != {s2}"
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[1]"

  # <ipython-input-2-d90f2ca6b8df>:2:       return torch.positive(a.positive())
  t0 = ltorch.positive(a)  # t0: "cpu f32[1]"
  t1 = ltorch.positive(t0)  # t1: "cpu f32[1]"
  return t1

AssertionError                            Traceback (most recent call last)
----> 1 assert s1 == s2, f"{s1} != {s2}"

AssertionError: [Symbol name=positive] != [Symbol name=positive]

By contrast, manually registered functions register both with a single symbol: https://github.com/Lightning-AI/lightning-thunder/blob/3770bf02240cdb6ec86a6d5224807bc6ad39d2a9/thunder/torch/__init__.py#L251-L256

(credit for discovery and analysis: @k223kim , thank you)

tfogal commented 1 week ago

@kiya00 is this something that you could take a look at?