Repository: pytorch
Issue: [export] Errors out when unflattening TorchTitan
Labels: oncall: export

### 🐛 Describe the bug

## Background
TorchTitan (mimic of LLaMA) has a forward function that is similar to the mini repro below:
```
import torch

# Simple module for demonstration
class layer(torch.nn.Module):
    def forward(self, x, const) -> torch.Tensor:
        return x + const

class M(torch.nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.register_buffer("const", torch.ones(64, 256))
        self.layers = torch.nn.ModuleList([layer() for _ in range(2)])

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        for layer in self.layers:
            x = layer(x, self.const)
        return x
```
The above forward function is a mimic of TorchTitan's [forward](https://github.com/pytorch/torchtitan/blob/54b5fa259d8dd6c55c044a536b949363e31ee9f9/torchtitan/models/llama/model.py#L431-L432):
```
for layer in self.layers.values():
    h = layer(h, self.freqs_cis)
```
That is, a buffer is registered at root level, and passed to each layer's forward as an input argument.

## Repro
Here are simple steps to export and unflatten it:
```
# Export
mod = M()
x = torch.randn(64, 256)
ep = torch.export.export(mod, (x,))
unflattened = torch.export.unflatten(ep)

# Run
unflattened(x)
```

## Error 1:
The export step works fine, but the unflatten step hits the following error:
```
$ python titan_like.py 
Traceback (most recent call last):
  File "/data/users/kw2501/export-playground/titan_like.py", line 24, in <module>
    unflattened = torch.export.unflatten(ep)
                  ^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/data/users/kw2501/pytorch/torch/export/unflatten.py", line 468, in unflatten
    return UnflattenedModule(module, flat_args_adapter)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/data/users/kw2501/pytorch/torch/export/unflatten.py", line 348, in __init__
    node.name not in inputs_to_state
AssertionError: b_const was not sunk into the module layers.0 which has the graph: graph():
    %b_const : [num_users=1] = placeholder[target=b_const]
    %x : [num_users=1] = placeholder[target=x]
    %add : [num_users=1] = call_function[target=torch.ops.aten.add.Tensor](args = (%x, %b_const), kwargs = {})
    return add
```

## Error 2:
If we comment out the assert from torch/export/unflatten.py (which is a self check), and continue with a run, we will hit a runtime error:
```
  File "/data/users/kw2501/pytorch/torch/fx/interpreter.py", line 235, 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 x, but one was not passed in!

While executing %x : [num_users=1] = placeholder[target=x]
Original traceback:
None

While executing %layers_0 : [num_users=1] = call_module[target=layers.0](args = (%x,), kwargs = {})
Original traceback:
None
```

This is expected because `layers.0` is expecting two inputs, but the parent module is feeding only one (i.e. the `x`).
```
: graph():
    %x : [num_users=1] = placeholder[target=x]
    %layers_0 : [num_users=1] = call_module[target=layers.0](args = (%x,), kwargs = {})
    %layers_1 : [num_users=1] = call_module[target=layers.1](args = (%layers_0,), kwargs = {})
    return (layers_1,)
layers.0: graph():
    %b_const : [num_users=1] = placeholder[target=b_const]
    %x : [num_users=1] = placeholder[target=x]
    %add : [num_users=1] = call_function[target=torch.ops.aten.add.Tensor](args = (%x, %b_const), kwargs = {})
    return add
```

### Versions

pytorch nightly

Cc @pianpwk @zhxchen17 @tugsbayasgalan @avikchaudhuri 
Cc @wconstab @H-Huang @wanchaol 

cc @avikchaudhuri @gmagogsfm @zhxchen17 @tugsbayasgalan @angelayi @suo @ydwu4