Unflatten traced module (#954)

## Description
- Move tracer from `_export_to_torch_ir` to the official
`torch.export.export`
- Add unflatten utils (from torch/export/unflatten.py) to unflatten each
stage module

Purpose of this PR is to:
- be composable with FSDP and TP, which requires structured FQNs like
`a.b.c` to submodules to specify their policies.
- be nice to DCP which would not like to see change of FQNs compared to
original model.
- retire use of `_export_to_torch_ir` per Export Team's plan.

## Test
Added `test_transformer.py`.
```
class TransformerLike(torch.nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.layers = torch.nn.Sequential(
            *[
                MLPModule(d_hid)
                for _ in range(n_layers)
            ]
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.layers(x)
```

We split the model into two stages. Each stages would preserve the
`layers.<i>` structure as in the original model.
```
Stage 0: 
 GraphModule(
  (layers): InterpreterModule(
    (0): InterpreterModule(
      (net1): InterpreterModule()
      (relu): InterpreterModule()
      (net2): InterpreterModule()
    )
    (1): InterpreterModule(
      (net1): InterpreterModule()
      (relu): InterpreterModule()
      (net2): InterpreterModule()
    )
    (2): InterpreterModule(
      (net1): InterpreterModule()
      (relu): InterpreterModule()
      (net2): InterpreterModule()
    )
    (3): InterpreterModule(
      (net1): InterpreterModule()
      (relu): InterpreterModule()
      (net2): InterpreterModule()
    )
  )
)
 ```
```
Stage 1: 
 GraphModule(
  (layers): InterpreterModule(
    (4): InterpreterModule(
      (net1): InterpreterModule()
      (relu): InterpreterModule()
      (net2): InterpreterModule()
    )
    (5): InterpreterModule(
      (net1): InterpreterModule()
      (relu): InterpreterModule()
      (net2): InterpreterModule()
    )
    (6): InterpreterModule(
      (net1): InterpreterModule()
      (relu): InterpreterModule()
      (net2): InterpreterModule()
    )
    (7): InterpreterModule(
      (net1): InterpreterModule()
      (relu): InterpreterModule()
      (net2): InterpreterModule()
    )
  )
)
```

Caveat:
I temporarily disabled multi-use parameter support (aka. shared paramters or tied parameters). So some real examples may break. Will add the support back in next PR.

Author: kwen2501

.flake8

@@ -28,5 +28,5 @@ exclude =
     ./torch/include,
     ./torch/lib,
     ./venv,
-    ./pippy/fx,
+    ./pippy/unflatten.py,
     *.pyi

.github/workflows/gpu_tests.yaml

@@ -42,9 +42,9 @@ jobs:
       - name: Activate conda env
         run: conda activate test
       - name: Install dependencies
-        run:
+        run: |
           pip install numpy
-          pip install -r requirements.txt --find-links https://download.pytorch.org/whl/nightly/cu121
+          pip install --pre -r requirements.txt --find-links https://download.pytorch.org/whl/nightly/cu121/torch_nightly.html
       - name: Install pippy
         run: python setup.py install
       - name: Run forward-only integration test
@@ -61,10 +61,12 @@ jobs:
         run: pip install transformers
       - name: Run GPT2
         run: torchrun --nproc-per-node 4 examples/huggingface/pippy_gpt2.py
-      - name: Run T5
-        run: torchrun --nproc-per-node 2 examples/huggingface/pippy_t5.py
-      - name: Run BERT
-        run: torchrun --nproc-per-node 4 examples/huggingface/pippy_bert.py
+      - name: Test CPU init + GPU run
+        run: torchrun --nproc-per-node 4 examples/cpu_init/gpt2_cpu_init.py
+      # - name: Run T5
+      #   run: torchrun --nproc-per-node 2 examples/huggingface/pippy_t5.py
+      # - name: Run BERT
+      #   run: torchrun --nproc-per-node 4 examples/huggingface/pippy_bert.py
 
   backward_tests_4gpu:
     runs-on: linux.g5.12xlarge.nvidia.gpu
@@ -87,7 +89,7 @@ jobs:
       - name: Install dependencies
         run:
           pip install numpy
-          pip install -r requirements.txt --find-links https://download.pytorch.org/whl/nightly/cu121
+          pip install -r requirements.txt --find-links https://download.pytorch.org/whl/nightly/cu121/torch_nightly.html
       - name: Install pippy
         run: python setup.py install
       - name: Run forward-backward test

.github/workflows/pippy_tests.yaml

@@ -89,10 +89,15 @@ jobs:
           if [ -f requirements.txt ]; then pip install -r requirements.txt --find-links https://download.pytorch.org/whl/nightly/cpu/torch_nightly.html; fi
       - name: Install pippy
         run: "python setup.py install"
+      # IR tests
       - name: Test forward pipe generation
         run: python test/test_pipe.py
       - name: Test backward pipe generation
         run: python test/test_pipe_bwd.py
+      - name: Test unflatten
+        run: python test/test_unflatten.py
+      - name: Test Transformer
+        run: python test/test_transformer.py
       - name: Test pipeline schedule
         run: python test/test_pipeline_schedule.py
       # - name: Run null_coalesce_accumulate integration test

check.sh

@@ -49,7 +49,7 @@ RETVAL=0
 
 if (( SKIP_FORMAT == 0 )); then
   echo; echo "Running format check ..."
-  ufmt diff pippy/*.py pippy/hf/*.py test/*.py
+  ufmt diff pippy/*.py test/*.py
   (( RETVAL |= $? ))
 fi
 

format.sh

@@ -7,7 +7,7 @@
 DEFAULT_TARGETS=()
 for f in $(git ls-files | grep '\.py$'); do
   case "$f" in
-    'pippy/fx/'*)
+    'pippy/unflatten.py')
       # ignore
       ;;