run tests in groups

allenai · Apr 2, 2024 · 3aa9ad3 · 3aa9ad3
1 parent 3a786a8
commit 3aa9ad3
Show file tree

Hide file tree

Showing 3 changed files with 15 additions and 2 deletions.
diff --git 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
@@ -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
@@ -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(