[Tensor Subclasses] Trace transfom to interpret __torch_dispatch__

[crcrpar]

• add a generic proxy class to represent Tensor Wrapper Subclasses that call torch.Tensor._make_wrapper_subclass in their __new__ and define their own __torch_dispatch__
• add a trace transform that evaluate BoundSymbols of a trace one by one so that we could make a trace free from actual tensor subclass objects as possible and write out the actual behavior that is defined by __torch_dispatch__ in a trace

[t-vi]

To my mind, we would get better error handling (eg stack traces) if we resolved the __torch_dispatch__ earlier by checking for subclasses in the torchsymbol handling logic. Also, I don't think we have much information about the outputs - e.g. are they subclasses again or the original class, shapes etc. - without doing so, so we cannot evaluate even metadata in control flow.

[mruberry]

There's a lot going on with this PR, and it's pretty complicated. Maybe we should schedule an online sync, @crcrpar and @IvanYashchuk, to see if we can make it more incremental?

[t-vi]

I would still prefer to do the flattening during interpretation for the benefit of getting error messages with backtraces.

[crcrpar]

There's a lot going on with this PR, and it's pretty complicated. Maybe we should schedule an online sync, @crcrpar and @IvanYashchuk, to see if we can make it more incremental?

We don't have a nice time slots that work all of us.
What would have you find PRs incremental?
Would it look better to you to split this into two, one to add a proxy and the other to add trace transforms?

EDIT: #1583, #1584, and #1585 are a sequence of PRs that cover this one and #1415.

I would still prefer to do the flattening during interpretation for the benefit of getting error messages with backtraces.

I embarrassingly am not familiar with interpretation implementation at all. Could you give me some pointers to add flattening so that it happens during interpretation?

EDIT: __torch_dispatch__ overrides the behavior of ops used in backward and the customized behavior could use torch ops that do not have their backward definition. Thus I do think we'd need to put the flattening AFTER forward-backward split as in #1415

[mruberry]

There's a lot going on with this PR, and it's pretty complicated. Maybe we should schedule an online sync, @crcrpar and @IvanYashchuk, to see if we can make it more incremental?

We don't have a nice time slots that work all of us. What would have you find PRs incremental? Would it look better to you to split this into two, one to add a proxy and the other to add trace transforms?

OK; we can try to work asynchronously. For a first incremental PR, would you create a PR adding support for aten operators? In particular, if someone were to call something like torch.ops.aten.add then how would that work? Can that operator be added to the torch operations (and be in the torch language context), or should it be added to its own aten language context that has a separate file (or set of files)?

[crcrpar]

if someone were to call something like torch.ops.aten.add then how would that work?

So far in my implementation there's an optimistic mapping from core aten ops to ltorch ops:

lightning-thunder/thunder/transforms/tensor_wrapper_subclass.py

Lines 339 to 348 in 515d425

	for node in list_of_function_call_node:
	if not hasattr(ltorch, node.target._opname):
	msg = (
	f"`thunder.torch` does not have corresponding op for {node.target._opname}. "
	"Think about adding it to thunder/torch/default_torch_ops.py"
	f"\nThe op is found while flattening the following BoundSymbol:\n{bsym}"
	f"\ntorch.fx graph:\n{fx_graph.print_readable(print_output=False)}"
	)
	raise RuntimeError(msg)
	ltorch_ops_for_node_of_ops.append(getattr(ltorch, node.target._opname))

.

By the way, if we're to cover core aten ops, then I'd say it'd be worth thinking of using thunder as a custom backend after AOTAutograd.

Can that operator be added to the torch operations (and be in the torch language context), or should it be added to its own aten language context that has a separate file (or set of files)?

Currently torchsymbol has some for core aten ops, apparently

lightning-thunder/thunder/torch/__init__.py

Lines 172 to 173 in 9d79b8d

	elif hasattr(torch.ops.aten, name):
	id = f"torch.ops.aten.{name}"

lightning-thunder/thunder/torch/__init__.py

Line 4700 in 9d79b8d

@torchsymbol(torch.ops.aten.embedding_backward)

and

lightning-thunder/thunder/torch/__init__.py

Line 4308 in 9d79b8d

torch.ops.aten._adaptive_avg_pool2d_backward,

are a core aten op. So I think extending thunder/torch/__init__.py would be fair.

[mruberry]

if someone were to call something like torch.ops.aten.add then how would that work?

So far in my implementation there's an optimistic mapping from core aten ops to ltorch ops:

lightning-thunder/thunder/transforms/tensor_wrapper_subclass.py

Lines 339 to 348 in 515d425

	for node in list_of_function_call_node:
	if not hasattr(ltorch, node.target._opname):
	msg = (
	f"`thunder.torch` does not have corresponding op for {node.target._opname}. "
	"Think about adding it to thunder/torch/default_torch_ops.py"
	f"\nThe op is found while flattening the following BoundSymbol:\n{bsym}"
	f"\ntorch.fx graph:\n{fx_graph.print_readable(print_output=False)}"
	)
	raise RuntimeError(msg)
	ltorch_ops_for_node_of_ops.append(getattr(ltorch, node.target._opname))

.
By the way, if we're to cover core aten ops, then I'd say it'd be worth thinking of using thunder as a custom backend after AOTAutograd.

Can that operator be added to the torch operations (and be in the torch language context), or should it be added to its own aten language context that has a separate file (or set of files)?

Currently torchsymbol has some for core aten ops, apparently

lightning-thunder/thunder/torch/__init__.py

Lines 172 to 173 in 9d79b8d

	elif hasattr(torch.ops.aten, name):
	id = f"torch.ops.aten.{name}"

lightning-thunder/thunder/torch/__init__.py

Line 4700 in 9d79b8d

@torchsymbol(torch.ops.aten.embedding_backward)

and

lightning-thunder/thunder/torch/__init__.py

Line 4308 in 9d79b8d

torch.ops.aten._adaptive_avg_pool2d_backward,

are a core aten op. So I think extending thunder/torch/__init__.py would be fair.

OK; expanding thunder/torch/init.py sounds good for now. Let's not "optimistically" try to map ATen operations to torch operations for the moment, but just treat them like different operations.

Would you submit a PR adding torch.ops.aten.add to the torch operations?

EDITED BELOW.

As a follow-up PR to that, what about working with a program like

# Original program
def foo(x):
  return x

# Trace
def computation(x):
  # x: "MyTensorSubclass[cuda:0 f32[12, 12]]" 
  return x

Where the initial trace shows the tensor subclass and its flattened information, and the prologue validates the subclass and its flattening. Then I'd be curious to see addition with that tensor, like this:

# Original program
def foo(x):
  return x + 1

# Trace
def computation(x):
  # x: "MyTensorSubclass[cuda:0 f32[12, 12]]" 
  t0 = MyTensorSubclass.torch.add(x, 1)  # t0: "MyTensorSubclass[cuda:0 f32[12, 12]]"
    # t1 = flatten_tensor_subclass(MyTensorSubclass, x)
    # t2 = torch.ops.aten.add(t1, 1)
      # <decomposition of aten.add into prims would go here> 
    # t0 = unflatten_tensor_subclass(MyTensorSubclass, t2)
  return t0

This can be translated for execution by PyTorch, but I think working through this will be interesting. Then the follow-up question is what the grad transform for it looks like, and how this operation should be translated for execution by nvFuser.

[crcrpar]

IMHO, it'd sound more natural to me to register core aten ops to thunder.torch namespace after merging #1583.

Then registration comes after the aforementioned PR, before #1584 and #1585, followed by some refinement of prologue and how traces with tensor subclasses look accompanied by #1584.

the follow-up question is what the grad transform for it looks like, and how this operation should be translated for execution by nvFuser.

With the experience of #1585, I do think we'd have to let the trace get split into forward and backward before interpreting __torch_dispatch__ partly because the extended behavior of certain ops could be dependent on ops without any backward definitions.