Repository: pytorch
Issue: Treating SymNode compute more fungibly for the purpose of export flattener (or not?)
Labels: module: dynamic shapes,export-triaged,oncall: export

### 🐛 Describe the bug

When working on https://github.com/pytorch/pytorch/pull/132421 I often ran into "RuntimeError: Unable to find nn_module_stack for node sym_size_int" errors from export's unflattener. Here's an example situation where this occurs.

I'm exporting:

```
class M(Module):
  def forward(self, x):
    return x / x.size(0)
```

This will turn into a graph like

```
def f(x):
  # from M.forward
  x_size = torch.ops.aten.sym_size.int(x, 0)
  r = torch.ops.aten.div.Tensor(x, x_size)
  return r
```

Now, the `torch.ops.aten.sym_size.int(x, 0)` doesn't /really/ belong to M.forward, in some sense. Specifically, the way make_fx works, we will only ever trace the actually sym_size.int call once, and then just reuse x_size as necessary downstream in the model.

Here is a full example:

```
    def test_unflatten_nested_warble(self):
        class Child1(torch.nn.Module):
            def forward(self, x):
                return x / x.size(0)
            
        class Child2(torch.nn.Module):
            def forward(self, x):
                return x / x.size(0)
                
        class MyModule(torch.nn.Module):
            def __init__(self) -> None:
                super().__init__()
                self.foo = Child1()
                self.bar = Child2()
                
            def forward(self, x):
                y1 = self.foo(x)
                y2 = self.bar(x)
                return y1 + y2
                
        orig_eager = MyModule()
        d = torch.export.Dim("foo")
        export_module = torch.export.export(orig_eager, (torch.rand(2, 3),), dynamic_shapes=((d,),))
        print(export_module)
        unflattened = unflatten(export_module)
        print(unflattened.foo.graph)
        print(unflattened.bar.graph)
        print(unflattened.graph)
        
        inputs = (torch.rand(2, 3),)
        
        # Compare the root modules and all submodules
        self.compare_outputs(orig_eager, unflattened, inputs)
        #self.compare_outputs(orig_eager.foo, unflattened.foo, inputs)
        #self.compare_outputs(orig_eager.bar, unflattened.bar, inputs)
```

and here are the generated subgraphs:

```
graph():        
    %sym_size_int_1 : [num_users=1] = placeholder[target=sym_size_int_1]
    %x : [num_users=1] = placeholder[target=x]
    %div : [num_users=1] = call_function[target=torch.ops.aten.div.Tensor](args = (%x, %sym_size_int_1), kwargs = {})
    return div
graph():
    %sym_size_int_1 : [num_users=1] = placeholder[target=sym_size_int_1]
    %x : [num_users=1] = placeholder[target=x]
    %div_1 : [num_users=1] = call_function[target=torch.ops.aten.div.Tensor](args = (%x, %sym_size_int_1), kwargs = {})
    return div_1
graph():
    %x : [num_users=3] = placeholder[target=x]
    %sym_size_int_1 : [num_users=2] = call_function[target=torch.ops.aten.sym_size.int](args = (%x, 0), kwargs = {})
    %foo : [num_users=1] = call_module[target=foo](args = (%sym_size_int_1, %x), kwargs = {})
    %bar : [num_users=1] = call_module[target=bar](args = (%sym_size_int_1, %x), kwargs = {})
    %add_6 : [num_users=1] = call_function[target=torch.ops.aten.add.Tensor](args = (%foo, %bar), kwargs = {})
    return (add_6,)
```

You can see that only one sym_size.int eventually got traced into the graph, and it is now being passed as an argument to the submodules. In some sense this is "wrong", because the sym_size.int call actually happened inside the submodules, but if you don't care about the call signatures of the submodules, it doesn't really matter.

Except, @angelayi tells me, we SHOULD care about the signatures of the submodules, because you might want to do module swapping on them! So it feels like there is some dynamic shapes / unflattener disaster waiting to happen.

Some of the simple stuff works. But here is a torture case that doesn't:

```
    def test_unflatten_nested_warble(self):
        class Child1(torch.nn.Module):
            def forward(self, x2):
                return x2 / x2.size(0)
        
        class Child2(torch.nn.Module):
            def forward(self, x):
                return x2 / x2.size(0)
        
        class MyModule(torch.nn.Module):
            def __init__(self) -> None:
                super().__init__()
                self.foo = Child1()
                self.bar = Child2()
            
            def forward(self, x):
                x2 = torch.cat([x, x])
                y1 = self.foo(x2)
                y2 = self.bar(x2)
                return y1 + y2
        
        orig_eager = MyModule()
        d = torch.export.Dim("foo")
        dyn = ((d,),)
        # dyn = None
        export_module = torch.export.export(orig_eager, (torch.rand(2, 3),), dynamic_shapes=dyn, preserve_module_call_signature=("foo", "bar"))
        print(export_module)
        unflattened = unflatten(export_module)
        print(unflattened.foo.graph)
        print(unflattened.bar.graph)
        print(unflattened.graph)
        
        inputs = (torch.rand(2, 3),)
        
        # Compare the root modules and all submodules
        self.compare_outputs(orig_eager, unflattened, inputs)
        
        inputs = (torch.rand(4, 3),)
        
        self.compare_outputs(orig_eager.foo, unflattened.foo, inputs)
        self.compare_outputs(orig_eager.bar, unflattened.bar, inputs)
```

This fails with 

```
======================================================================
ERROR: test_unflatten_nested_warble (__main__.TestUnflatten)
----------------------------------------------------------------------
Traceback (most recent call last):
  File "/data/users/ezyang/a/pytorch/torch/fx/interpreter.py", line 232, in placeholder
    return next(self.args_iter)
StopIteration
                
The above exception was the direct cause of the following exception:
        
Traceback (most recent call last):
  File "/data/users/ezyang/a/pytorch/torch/testing/_internal/common_utils.py", line 2918, in wrapper
    method(*args, **kwargs)
  File "/data/users/ezyang/a/pytorch/test/export/test_unflatten.py", line 101, in test_unflatten_nested_warble
    self.compare_outputs(orig_eager.foo, unflattened.foo, inputs)
  File "/data/users/ezyang/a/pytorch/test/export/test_unflatten.py", line 59, in compare_outputs
    unflattened_output = unflattened(*args)
  File "/data/users/ezyang/a/pytorch/torch/nn/modules/module.py", line 1736, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/data/users/ezyang/a/pytorch/torch/nn/modules/module.py", line 1747, in _call_impl
    return forward_call(*args, **kwargs)
  File "/data/users/ezyang/a/pytorch/torch/export/unflatten.py", line 115, in forward
    return torch.fx.Interpreter(self, graph=self.graph).run(
  File "/data/users/ezyang/a/pytorch/torch/fx/interpreter.py", line 147, in run
    self.env[node] = self.run_node(node)
  File "/data/users/ezyang/a/pytorch/torch/fx/interpreter.py", line 204, in run_node
    return getattr(self, n.op)(n.target, args, kwargs)
  File "/data/users/ezyang/a/pytorch/torch/fx/interpreter.py", line 237, in placeholder
    raise RuntimeError(f'Expected positional argument for parameter {target}, but one was not passed in!') from si
RuntimeError: Expected positional argument for parameter _positional_arg_0, but one was not passed in!
        
While executing %_positional_arg_0 : [num_users=1] = placeholder[target=_positional_arg_0]
Original traceback:
None    
        
To execute this test, run the following from the base repo dir:
    python test/export/test_unflatten.py TestUnflatten.test_unflatten_nested_warble
        
This message can be suppressed by setting PYTORCH_PRINT_REPRO_ON_FAILURE=0
```

It is clear that the call signature of the module has not been preserved. If you turn off dynamic shapes, the call signature matches and the test passes.

Here is one of the submodule graphs:

```
graph():            
    %x : [num_users=1] = placeholder[target=x]                                     
    %_positional_arg_0 : [num_users=1] = placeholder[target=_positional_arg_0]                                                                                        
    %_spec_0 : [num_users=1] = get_attr[target=_spec_0]
    %tree_flatten_spec : [num_users=1] = call_function[target=torch.fx._pytree.tree_flatten_spec](args = (((%_positional_arg_0,), {}), %_spec_0), kwargs = {})
    %cat : [num_users=1] = call_function[target=operator.getitem](args = (%tree_flatten_spec, 0), kwargs = {})
    %sym_size_int : [num_users=1] = call_function[target=torch.ops.aten.sym_size.int](args = (%x, 0), kwargs = {})
    %mul : [num_users=1] = call_function[target=operator.mul](args = (2, %sym_size_int), kwargs = {})
    %div : [num_users=1] = call_function[target=torch.ops.aten.div.Tensor](args = (%cat, %mul), kwargs = {})
    %_spec_1 : [num_users=1] = get_attr[target=_spec_1]
    %tree_unflatten : [num_users=1] = call_function[target=torch.utils._pytree.tree_unflatten](args = ((%div,), %_spec_1), kwargs = {})
    return tree_unflatten
```

In a last ditch effort to get the quantity of x2.size(0), we have somehow decided we must pass x in as an argument, so that we can call x.size(0) on it and then multiply it by two (replaying the cat logic). This breaks signature preservation.

It seems to me that the unflattener really wants size compute to NOT be fungible. Instead of assuming that s0 can be served from anywhere, we want to accurately trace the dependencies inside the model, so we don't generate extra dependencies that can interfere with making module signatures. This is pretty different from where the deferred runtime asserts pass was going, where we instead said "hey, there will be a single, canonical FX node for any given symbol, and you can use it everywhere".

If you want to maintain fungibility, maybe you want to just be willing to regenerate sizes inside modules. This can be kind of tricky, but it should always be in principle be possible. But in that case, the unflattener really shouldn't complain if there's a `x_size = torch.ops.aten.sym_size.int(x, 0)` hanging out with no NN module stack associated with it. It doesn't matter! Use it, don't use it, I don't care.

cc @penguinwu @avikchaudhuri @gmagogsfm @zhxchen17 @tugsbayasgalan @angelayi @suo @ydwu4 @pianpwk @eellison 

### Versions

main