From 3aa9ad366ed2f24cbd26a23b9bfcb6f49705b851 Mon Sep 17 00:00:00 2001
From: epwalsh <petew@allenai.org>
Date: Tue, 2 Apr 2024 14:46:11 -0700
Subject: [PATCH] run tests in groups

---
 .github/workflows/main.yml             | 9 +++++++--
 pyproject.toml                         | 1 +
 src/test/distributed/fsdp/fsdp_test.py | 7 +++++++
 3 files changed, 15 insertions(+), 2 deletions(-)

diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml
index d32f1e19..c478ed1e 100644
--- a/.github/workflows/main.yml
+++ b/.github/workflows/main.yml
@@ -38,8 +38,13 @@ jobs:
         include:
           - python: '3.10'
             task:
-              name: Test
-              run: pytest -v --color=yes --durations=5 src/test/
+              name: Test (main group)
+              run: pytest -v --color=yes --durations=5 src/test/ -m 'not group1'
+
+          - python: '3.10'
+            task:
+              name: Test (group 1)
+              run: pytest -v --color=yes --durations=5 src/test/ -m 'group1'
 
           - python: '3.10'
             task:
diff --git a/pyproject.toml b/pyproject.toml
index 2de42ddb..fd14ad69 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -151,6 +151,7 @@ log_cli = false
 log_cli_level = "DEBUG"
 markers = [
     "gpu",
+    "group1",
 ]
 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 ec0b7736..0f355a26 100644
--- a/src/test/distributed/fsdp/fsdp_test.py
+++ b/src/test/distributed/fsdp/fsdp_test.py
@@ -85,6 +85,7 @@ def run_fsdp_against_non_distributed_model(model_factory, model_data_factory):
             )
 
 
+@pytest.mark.group1
 @pytest.mark.parametrize("backend", BACKENDS)
 def test_fsdp_against_non_distributed_model(backend, tiny_model_factory, tiny_model_data_factory):
     run_distributed_test(
@@ -158,6 +159,7 @@ def run_fsdp_against_ddp(model_factory, model_data_factory):
     optim.step()
 
 
+@pytest.mark.group1
 @pytest.mark.parametrize("backend", BACKENDS)
 def test_fsdp_against_ddp(backend, tiny_model_factory, tiny_model_data_factory):
     run_distributed_test(
@@ -215,6 +217,7 @@ def run_fsdp_with_gradient_accumulation(model_factory, model_data_factory):
             )
 
 
+@pytest.mark.group1
 @pytest.mark.parametrize("backend", BACKENDS)
 def test_fsdp_with_gradient_accumulation(backend, tiny_model_factory, tiny_model_data_factory):
     run_distributed_test(
@@ -347,6 +350,7 @@ def forward(self, x):
     loss.backward()
 
 
+@pytest.mark.group1
 @pytest.mark.parametrize("backend", BACKENDS)
 def test_nested_fsdp_api(backend, tiny_model_factory, tiny_model_data_factory):
     run_distributed_test(
@@ -383,6 +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.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):
@@ -431,6 +436,7 @@ def __init__(self):
     assert fsdp.module.fc3.out_proj.max_prefetch_count == 3
 
 
+@pytest.mark.group1
 @pytest.mark.parametrize("backend", BACKENDS)
 def test_auto_wrap(backend):
     run_distributed_test(
@@ -491,6 +497,7 @@ def initialize_and_check(m: nn.Module):
         assert (param.data.detach() == 1.1).all()
 
 
+@pytest.mark.group1
 @pytest.mark.parametrize("backend", BACKENDS)
 def test_apply(backend):
     run_distributed_test(