Automatic registration of torch operators using FakeTensor

[kiya00]

What does this PR do?

Design doc: https://docs.google.com/document/d/1BVr_5co6GDj3wC7NrWq0ckfsbgX5l6w1dOSFBi2c-Hc/edit?usp=sharing

This PR did some experiments to use FakeTensor for our meta function to allow automatic fallback to torch operators when the op is not supposed to have a fusion backend.

Currently it can work for the example test cases(the sinc, hsplit, normalize, adaptive_avg_pool3d are ops currently not added in thunder ):

import torch
import thunder
from torch import nn

def func(a, out):
    t1 = torch.nn.functional.adaptive_avg_pool3d(a, output_size=out)
    return torch.nn.functional.normalize(t1)

def func1(a):
    x = torch.sinc(a)
    t = torch.arange(16.0).reshape(4,4)
    return torch.hsplit(t, 2), x


a = torch.randn((1, 64, 8, 9, 10)).cuda()
jfunc = thunder.jit(func)
out = jfunc(a, (5, 7, 9))
res = func(a, (5, 7, 9))
torch.testing.assert_close(out, res)

inp = torch.rand(2,2).cuda()
jfunc1 = thunder.jit(func1)
out = jfunc1(inp)
expected = func1(inp)
torch.testing.assert_close(out, expected)

import torch
import thunder
def func(a):
    x = torch.sinc(a)
    x1,x2= torch.hsplit(x, 2)
    return x1+x2

def func1(a, out):
    t1 = torch.nn.functional.adaptive_avg_pool3d(a, output_size=out)
    return torch.nn.functional.normalize(t1)


inp = torch.rand(4,4).cuda().requires_grad_()
inp_1 = inp.detach().clone().requires_grad_()


jfunc = thunder.jit(func)
out = jfunc(inp)
out.sum().backward()
print(thunder.last_traces(jfunc)[-1])
print(thunder.last_backward_traces(jfunc)[-1])


expected = func(inp_1)
expected.sum().backward()

torch.testing.assert_close(out, expected)
torch.testing.assert_close(inp.grad, inp_1.grad)


a = torch.randn((1, 64, 8, 9, 10)).cuda().requires_grad_()
a1 = a.detach().clone().requires_grad_()

res = func1(a1, (5, 7, 9))
res.sum().backward()

jfunc = thunder.jit(func1)
out = jfunc(a, (5, 7, 9))
out.sum().backward()
print(thunder.last_traces(jfunc)[-1])
print(thunder.last_backward_traces(jfunc)[-1])

torch.testing.assert_close(out, res)
torch.testing.assert_close(a.grad, a1.grad)

Fixes #811

cc: @IvanYashchuk

[riccardofelluga]

Hey this is pretty cool! I've tried it with this example that we had no op implemented and it worked!

import torch
import thunder

@thunder.jit
def func(x, y, z):
    t0 = torch.baddbmm(x, y, z)
    return t0

t0 = torch.randn(10, 3, 5)
b1 = torch.randn(10, 3, 4)
b2 = torch.randn(10, 4, 5)
out = func(t0, b1, b2)
print(thunder.last_traces(func)[-1])
print(out.size())

looking forward for this feature!

[kiya00]

currently output trace is like:

@torch.no_grad()
@no_autocast
def augmented_forward_fn(a):
  # a: "cuda:0 f32[1, 64, 8, 9, 10]"
  t0 = torch.nn.functional.adaptive_avg_pool3d(a, output_size=(5, 7, 9))  # t0: "cuda:0 f32[1, 64, 5, 7, 9]"
  t1 = torch.nn.functional.normalize(t0)  # t1: "cuda:0 f32[1, 64, 5, 7, 9]"
  return {'output': t1, 'flat_args': [a], 'flat_output': (t1,)}, ((a, t0), (5, 7, 9))
# Constructed by Delete Last Used (took 0 milliseconds)
import torch
from thunder.executors.torchex import no_autocast

@torch.no_grad()
@no_autocast
def backward_fn(saved_for_backward, cotangents):
  # saved_for_backward: "Collection"
  # cotangents: "Collection"
  C0, C1, = saved_for_backward
  clear_mutable_collection(saved_for_backward)
  del saved_for_backward
  t2, = cotangents
  clear_mutable_collection(cotangents)
  del cotangents
  a, t0, = C0
  clear_mutable_collection(C0)
  del C0
  i0, i1, i2, = C1
  clear_mutable_collection(C1)
  del C1
  (t14,) = normalize_vjp({'inputs': ((t0,), {}), 'func': _function_0}, t2)
  del t0, t2
  (t15,) = adaptive_avg_pool3d_vjp({'inputs': ((a,), {'output_size': (i0, i1, i2)}), 'func': _function_1}, t14)
  del a, i0, i1, i2, t14
  return (t15,)

[t-vi]

First off, great job, @kiya00 !

One fundamental question I would have about the design is: given that we are talking about a finite set of operators here, would it be more transparent to autogenerate the code for them and keep them a file? That would have us need to keep up / update for new stuff PyTorch introduces, but in general, it would be less magic. When we implement something in a more detailed way, it would be removed from the autogeneration to manual.

I'm imagining that we could/would still use the bulk of your code, just the time of registration would move.

WDYT?

[IvanYashchuk]

Awesome work!

There are only a few cosmetic changes to be done and then I think we should merge this.

[crcrpar]

what an awesome work

[kshitij12345]

The PR looks really great! Just had a couple of suggestions, thank you @kiya00!

[kiya00]

Hi @t-vi , I found the CI failure seems unrelated to this PR, and can only reproduce it locally by running pytest thunder/tests/test_inplace_functionalization.py thunder/tests/test_nvfuser_remat.py::test_find_cut_dropout_nvfuser_cuda_None together (Separate execution succeeds). I tried to fix it by #928

[t-vi]

Thank you @kiya00 @kshitij12345 @crcrpar @IvanYashchuk @riccardofelluga @mruberry