Add 3d+fp8 to test runner

ghstack-source-id: 8009c4635e277fccb8059efccc510ce89c985936 Pull Request resolved: #563
pytorch · Aug 29, 2024 · b8c34fe · b8c34fe
1 parent 8c497b7
commit b8c34fe
Show file tree

Hide file tree

Showing 2 changed files with 25 additions and 0 deletions.
diff --git a/.github/workflows/integration_test_8gpu.yaml b/.github/workflows/integration_test_8gpu.yaml
@@ -37,5 +37,6 @@ jobs:
         pip config --user set global.progress_bar off
 
         python -m pip install --force-reinstall --pre torch --index-url https://download.pytorch.org/whl/nightly/cu121
+        USE_CPP=0 python -m pip install git+https://github.com/pytorch/ao.git
         mkdir artifacts-to-be-uploaded
         python ./test_runner.py artifacts-to-be-uploaded --ngpu 8
diff --git a/test_runner.py b/test_runner.py
@@ -244,6 +244,30 @@ def build_test_list():
             requires_seed_checkpoint=True,
             ngpu=8,
         ),
+        OverrideDefinitions(
+            [
+                # [
+                #     "--experimental.pipeline_parallel_degree 2",
+                #     "--experimental.pipeline_parallel_split_points layers.4",
+                #     "--training.data_parallel_degree 2",
+                #     "--training.tensor_parallel_degree 2",
+                #     "--training.compile"
+                # ],
+                [
+                    "--experimental.pipeline_parallel_degree 2",
+                    "--experimental.pipeline_parallel_split_points layers.4",
+                    "--training.data_parallel_degree 2",
+                    "--training.tensor_parallel_degree 2",
+                    "--float8.enable_float8_linear",
+                    "--float8.enable_fsdp_float8_all_gather",
+                    "--float8.precompute_float8_dynamic_scale_for_fsdp",
+                ],
+            ],
+            "PP+DP+TP 3D test with fp8 and torch.compile",
+            "pp_dp_tp_composability",
+            requires_seed_checkpoint=True,
+            ngpu=8,
+        ),
         OverrideDefinitions(
             [
                 [