Unflatten traced module

.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
       ;;