Add ssm

FlagOpen · Dec 17, 2024 · 5b7fa1c · 5b7fa1c
1 parent faa7b7f
commit 5b7fa1c
Show file tree

Hide file tree

Showing 3 changed files with 8 additions and 0 deletions.
diff --git a/.github/workflows/all-tests.yml b/.github/workflows/all-tests.yml
@@ -32,6 +32,7 @@ jobs:
           - inference
           - models
           - pipeline_parallel
+          - ssm
           - tensor_parallel
           - transformer/moe
           - transformer

diff --git a/tests/scripts/unit_tests/config.yml b/tests/scripts/unit_tests/config.yml
@@ -17,6 +17,12 @@ megatron:
         - test_t5_model.py::TestT5Model::test_post_process_forward
         - test_t5_model.py::TestT5Model::test_forward_output_encoder_hidden_only
         - test_t5_model.py::TestT5Model::test_forward_with_encoder_hidden_states
+    ssm:
+      ignore:
+        - test_mamba_block.py
+        - test_mamba_hybrid_layer_allocation.py
+        - test_mamba_layer.py
+        - test_mamba_mixer.py
     transformer/moe:
       type: single 
       ignore: 

diff --git a/tests/scripts/unit_tests/test_all.sh b/tests/scripts/unit_tests/test_all.sh
@@ -16,6 +16,7 @@ commands=(
     "tests/scripts/unit_tests/test_subset.sh --backend megatron --subset inference"
     "tests/scripts/unit_tests/test_subset.sh --backend megatron --subset models"
     "tests/scripts/unit_tests/test_subset.sh --backend megatron --subset pipeline_parallel"
+    "tests/scripts/unit_tests/test_subset.sh --backend megatron --subset ssm"
     "tests/scripts/unit_tests/test_subset.sh --backend megatron --subset tensor_parallel"
     "tests/scripts/unit_tests/test_subset.sh --backend megatron --subset transformer/moe"
     "tests/scripts/unit_tests/test_subset.sh --backend megatron --subset transformer"