Repository: pytorch
Issue: [Pipelining] Problem using `torch.distributed.pipelining` on `Gemma2ForCausalLM`
Labels: oncall: distributed

Hi there!

I have problem testing torch pipelining against gemma2 model. Here is the code snippet:

```python
config = AutoConfig.from_pretrained("google/gemma-2-2b-it")
model = Gemma2ForCausalLM(config)
pipeline(model, mb_args=(torch.LongTensor([[1,2,3],[4,5,6]]),), split_spec={"model.layers.13": SplitPoint.BEGINNING})
```
Output:
```
/home/sbarikbin/.venv/lib/python3.12/site-packages/torch/export/_unlift.py:60: UserWarning: Attempted to insert a get_attr Node with no underlying reference in the owning GraphModule! Call GraphModule.add_submodule to add the necessary submodule, GraphModule.add_parameter to add the necessary Parameter, or nn.Module.register_buffer to add the necessary buffer
  getattr_node = gm.graph.get_attr(lifted_node)
/home/sbarikbin/.venv/lib/python3.12/site-packages/torch/fx/graph.py:1586: UserWarning: Node model_lifted_tensor_0 target model.lifted_tensor_0 lifted_tensor_0 of model does not reference an nn.Module, nn.Parameter, or buffer, which is what 'get_attr' Nodes typically target
  warnings.warn(f'Node {node} target {node.target} {atom} of {seen_qualname} does '
Traceback (most recent call last):
  File "/home/sbarikbin/test_pp.py", line 42, in <module>
    test1()
  File "/home/sbarikbin/test_pp.py", line 18, in test1
    p = pipeline(model, mb_args=(torch.LongTensor([[1,2,3],[4,5,6]]),), split_spec={"model.layers.13": SplitPoint.BEGINNING})
        ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/sbarikbin/.venv/lib/python3.12/site-packages/torch/distributed/pipelining/_IR.py", line 1231, in pipeline
    return Pipe.from_tracing(
           ^^^^^^^^^^^^^^^^^^
  File "/home/sbarikbin/.venv/lib/python3.12/site-packages/torch/distributed/pipelining/_IR.py", line 1051, in from_tracing
    pipe = Pipe._from_traced(
           ^^^^^^^^^^^^^^^^^^
  File "/home/sbarikbin/.venv/lib/python3.12/site-packages/torch/distributed/pipelining/_IR.py", line 750, in _from_traced
    new_submod = _outline_submodules(submodule.graph)
                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/sbarikbin/.venv/lib/python3.12/site-packages/torch/distributed/pipelining/_unflatten.py", line 24, in _outline_submodules
    ).run_outer()
      ^^^^^^^^^^^
  File "/home/sbarikbin/.venv/lib/python3.12/site-packages/torch/export/unflatten.py", line 1014, in run_outer
    self.run_from(node_idx)
  File "/home/sbarikbin/.venv/lib/python3.12/site-packages/torch/export/unflatten.py", line 1094, in run_from
    ).run_from(node_idx)
      ^^^^^^^^^^^^^^^^^^
  File "/home/sbarikbin/.venv/lib/python3.12/site-packages/torch/export/unflatten.py", line 1094, in run_from
    ).run_from(node_idx)
      ^^^^^^^^^^^^^^^^^^
  File "/home/sbarikbin/.venv/lib/python3.12/site-packages/torch/export/unflatten.py", line 1094, in run_from
    ).run_from(node_idx)
      ^^^^^^^^^^^^^^^^^^
  File "/home/sbarikbin/.venv/lib/python3.12/site-packages/torch/export/unflatten.py", line 1071, in run_from
    self.finalize_outputs()
  File "/home/sbarikbin/.venv/lib/python3.12/site-packages/torch/export/unflatten.py", line 993, in finalize_outputs
    _verify_graph_equivalence(self.cached_graph_module, self.module)
  File "/home/sbarikbin/.venv/lib/python3.12/site-packages/torch/export/unflatten.py", line 655, in _verify_graph_equivalence
    assert graph_dump(x.graph) == graph_dump(y.graph)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
AssertionError
```

I would be grateful if you help me on that. For example I couldn't find info about what is `torch.export.unflatten._ModuleFrame` (which `_outline_submodules` uses) for.

@H-Huang , @kwen2501 


cc @H-Huang @awgu @kwen2501 @wanchaol @fegin @fduwjj @wz337 @wconstab @d4l3k @c-p-i-o