Repository: pytorch
Issue: [dynamo][inlining-inbuilt-nn-modules] decide how dynamo/export should handle parametrizations
Labels: triaged,module: __torch_dispatch__,oncall: pt2,module: dynamo,oncall: export,module: pt2-dispatcher

Parametrizations have surfaced as a potentially-useful way to handle the situation where:

(1) you have a model where it has params/buffers are python tensor subclasses
(2) you want to be able to export it (to a C++ runtime where you cannot use tensor subclasses, like AOTI)

The idea is that you register a parametrization, such that the tensors stored on your state dict are the plain inner tensors. The parametrization then constructs a subclass out of these tensors, and that construction happens inside the graph (and gets traced away as part of subclass desugaring).

One reason this is nice because avoids the problem of what a C++ runtime behind export should do if it is running a model whos parameters are tensor subclasses: instead of having to manually re-write your state dict ahead-of-time to desugar the tensor subclasses yourself, you can:

(a) rely the parametrization to remove subclasses from the state dict (so the export runtime no longer needs to do extra work)
(b) rely on our existing tracing infra (dynamo/AOT) to trace through the tensor subclass construction logic, which will now be in the graph, thanks to the parametrization.

However, this doesn't quite work today for a few reasons.

When you run `aot_compile`, we [first generate a graph](https://github.com/pytorch/pytorch/blob/main/torch/_export/__init__.py#L368) from dynamo that goes to inductor, which thens calls `[aot_export_module](https://github.com/pytorch/pytorch/blob/main/torch/_inductor/compile_fx.py#L1390)`.

`aot_export_module` will then [lift](https://github.com/pytorch/pytorch/blob/main/torch/_functorch/aot_autograd.py#L1047) params/buffers on the module into graph inputs

I made a small repro below that errors with `RuntimeError: aot_export is not currently supported with traceable tensor subclass`.

**What we probably want**: When we export a model that registers a parametrization, the underlying graph that we produce should have graph inputs that from before-the-parametrization (aka the inner tensors, not the subclass

**Open question**: Do we want the same behavior when you `torch.compile` a model with a parametrization? I think the answer here is **yes**: in particular, parametrizations allow autograd to flow from the new parameter to the underlying tensors. In order for this gradient compute to be captured in the graph, AOTAutograd needs to see the underlying tensors as inputs to its `autograd.Function.forward`.

**What happens instead (in the below repro)**

The initial call to generate the torch IR graph ends up invoking the parametrization, and dynamo registers the parametrization output (a tensor subclass) as a **buffer** on the GraphModule that it sends to AOTAutograd.

**How should this work?** I can think of a few options:

(1) AOTAutograd does not see the parametrization

In this version, we would arrange for the GraphModule that dynamo sends to AOTAutograd to have a set of parameters/buffers corresponding to the **before-parametrization** code. In order for this to work, the user-defined parametrization code would need to have been inlined by dynamo into the graph

(2) AOTAutograd sees the parametrization

This would be kind of like allow-in-graph'ing the parametrization: dynamo would plumb the parametrization on each weight directly to the GraphModule that AOTAutograd sees. If this happens, then AOTAutograd will need to ensure that (a) lifts the underlying tensors of the parametrization as inputs instead of the parametrization itself, and (b) invokes the parametrization during its round of tracing.

After writing both options out, (1) seems like the right thing to do for safety reasons (parametrization are defined by the user, so we need dynamo safety. I would be interested in whether anyone has alternative ideas / if I'm missing anything.

**repro code**
```
import torch
from torch.testing._internal.two_tensor import TwoTensor
import torch.nn.utils.parametrize as parametrize

class Mod(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.lin1 = torch.nn.Linear(16, 16)
        self.relu = torch.nn.ReLU()
        self.lin2 = torch.nn.Linear(16, 16)

    def forward(self, x):
        x = self.lin1(x)
        x = self.relu(x)
        x = self.lin2(x)
        return x


class TwoTensorParametrization(torch.nn.Module):
    def __init__(self):
        super().__init__()

    def forward(self, a, b):
        return TwoTensor(a, b)

    def right_inverse(self, subclass_instance):
        return subclass_instance.a, subclass_instance.b


m = Mod()
# Convert both linear's weight params to TwoTensors
m.lin1.weight = torch.nn.Parameter(TwoTensor(m.lin1.weight, m.lin1.weight.clone()))
m.lin2.weight = torch.nn.Parameter(TwoTensor(m.lin2.weight, m.lin2.weight.clone()))
# Register parametrization for both subclasses, so their desugared tensors how up in m.state_dict()
parametrize.register_parametrization(m.lin1, "weight", TwoTensorParametrization())
parametrize.register_parametrization(m.lin2, "weight", TwoTensorParametrization())

x = torch.ones(16, 16)
# API used by aot_compile, before going to aot_export_module
from torch.export._trace import _export_to_torch_ir
gm = _export_to_torch_ir(m, (x,), {}, dynamic_shapes=None, disable_constraint_solver=None, restore_fqn=False)
print([f'fqn={k}, param_type={type(v)}' for k, v in gm.named_parameters()])
print([f'fqn={k}, buffer_type={type(v)}' for k, v in gm.named_buffers()])
```

cc @Chillee @ezyang @zou3519 @albanD @samdow @chauhang @penguinwu @voznesenskym @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @chenyang78 @kadeng @amjames @rec @avikchaudhuri @gmagogsfm @zhxchen17 @tugsbayasgalan @angelayi @suo @ydwu4 @msaroufim @anijain2305