From 7c5e20fef54f9d4199e14fc7b92dfff5bcc19578 Mon Sep 17 00:00:00 2001
From: epwalsh <petew@allenai.org>
Date: Tue, 2 Apr 2024 15:00:49 -0700
Subject: [PATCH] more groups

---
 .github/workflows/main.yml             | 48 +++++++++++---------------
 pyproject.toml                         |  3 +-
 src/test/distributed/fsdp/fsdp_test.py | 14 ++++----
 3 files changed, 30 insertions(+), 35 deletions(-)

diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml
index c478ed1e..d87933a1 100644
--- a/.github/workflows/main.yml
+++ b/.github/workflows/main.yml
@@ -23,46 +23,40 @@ env:
 
 jobs:
   checks:
-    name: ${{ matrix.task.name }} (py ${{ matrix.python }})
+    name: ${{ matrix.task.name }}
     runs-on: [ubuntu-latest]
     timeout-minutes: 15
     strategy:
       fail-fast: false
       matrix:
-        python: ['3.8', '3.10']
+        python: ['3.10']
         task:
           - name: Lint
-            run: |
-              make lint-check
+            run: make lint-check
 
-        include:
-          - python: '3.10'
-            task:
-              name: Test (main group)
-              run: pytest -v --color=yes --durations=5 src/test/ -m 'not group1'
+          - name: Test (default)
+            run: pytest -v --color=yes --durations=5 src/test/ -m 'not fsdp1 and not fsdp2'
 
-          - python: '3.10'
-            task:
-              name: Test (group 1)
-              run: pytest -v --color=yes --durations=5 src/test/ -m 'group1'
+          - name: Test FSDP (1)
+            run: pytest -v --color=yes --durations=5 src/test/ -m 'fsdp1'
 
-          - python: '3.10'
-            task:
-              name: Type check
-              run: |
-                make type-check
+          - name: Test FSDP (2)
+            run: pytest -v --color=yes --durations=5 src/test/ -m 'fsdp2'
 
-          - python: '3.10'
-            task:
-              name: Build
-              run: |
-                make build
+          - name: Type check
+            run: make type-check
+
+          - name: Build
+            run: make build
 
-          - python: '3.10'
+          - name: Style
+            run: make style-check
+
+        include:
+          - python: '3.8'
             task:
-              name: Style
-              run: |
-                make style-check
+              name: Lint (min Python)
+              run: make lint-check
 
     steps:
       - uses: actions/checkout@v3
diff --git a/pyproject.toml b/pyproject.toml
index fd14ad69..5e49e213 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -151,7 +151,8 @@ log_cli = false
 log_cli_level = "DEBUG"
 markers = [
     "gpu",
-    "group1",
+    "fsdp1",
+    "fsdp2",
 ]
 filterwarnings = [
     'ignore::FutureWarning:huggingface_hub\.file_download',
diff --git a/src/test/distributed/fsdp/fsdp_test.py b/src/test/distributed/fsdp/fsdp_test.py
index 0f355a26..10af397d 100644
--- a/src/test/distributed/fsdp/fsdp_test.py
+++ b/src/test/distributed/fsdp/fsdp_test.py
@@ -85,7 +85,7 @@ def run_fsdp_against_non_distributed_model(model_factory, model_data_factory):
             )
 
 
-@pytest.mark.group1
+@pytest.mark.fsdp1
 @pytest.mark.parametrize("backend", BACKENDS)
 def test_fsdp_against_non_distributed_model(backend, tiny_model_factory, tiny_model_data_factory):
     run_distributed_test(
@@ -159,7 +159,7 @@ def run_fsdp_against_ddp(model_factory, model_data_factory):
     optim.step()
 
 
-@pytest.mark.group1
+@pytest.mark.fsdp1
 @pytest.mark.parametrize("backend", BACKENDS)
 def test_fsdp_against_ddp(backend, tiny_model_factory, tiny_model_data_factory):
     run_distributed_test(
@@ -217,7 +217,7 @@ def run_fsdp_with_gradient_accumulation(model_factory, model_data_factory):
             )
 
 
-@pytest.mark.group1
+@pytest.mark.fsdp1
 @pytest.mark.parametrize("backend", BACKENDS)
 def test_fsdp_with_gradient_accumulation(backend, tiny_model_factory, tiny_model_data_factory):
     run_distributed_test(
@@ -350,7 +350,7 @@ def forward(self, x):
     loss.backward()
 
 
-@pytest.mark.group1
+@pytest.mark.fsdp1
 @pytest.mark.parametrize("backend", BACKENDS)
 def test_nested_fsdp_api(backend, tiny_model_factory, tiny_model_data_factory):
     run_distributed_test(
@@ -387,7 +387,7 @@ def run_fsdp_with_mixed_precision(model_factory, model_data_factory, precision):
             assert param.grad.dtype == param.dtype
 
 
-@pytest.mark.group1
+@pytest.mark.fsdp2
 @pytest.mark.parametrize("backend", BACKENDS)
 @pytest.mark.parametrize("precision", FSDP_MIXED_PRECISION)
 def test_fsdp_with_mixed_precision(backend, tiny_model_factory, tiny_model_data_factory, precision):
@@ -436,7 +436,7 @@ def __init__(self):
     assert fsdp.module.fc3.out_proj.max_prefetch_count == 3
 
 
-@pytest.mark.group1
+@pytest.mark.fsdp1
 @pytest.mark.parametrize("backend", BACKENDS)
 def test_auto_wrap(backend):
     run_distributed_test(
@@ -497,7 +497,7 @@ def initialize_and_check(m: nn.Module):
         assert (param.data.detach() == 1.1).all()
 
 
-@pytest.mark.group1
+@pytest.mark.fsdp1
 @pytest.mark.parametrize("backend", BACKENDS)
 def test_apply(backend):
     run_distributed_test(