mark fa2

huggingface · Apr 15, 2024 · ad2f79d · ad2f79d
1 parent af413a1
commit ad2f79d
Show file tree

Hide file tree

Showing 3 changed files with 3 additions and 0 deletions.
diff --git a/tests/test_base_model.py b/tests/test_base_model.py
@@ -10,6 +10,7 @@
 
 
 @pytest.mark.parametrize("tp,dp,pp", [(1, 1, 1), (2, 2, 2)])
+@pytest.mark.fa2
 @rerun_if_address_is_in_use()
 def test_get_named_modules_in_pp_rank(tp: int, dp: int, pp: int):
     model_args = ModelArgs(init_method=RandomInit(std=1.0), model_config=TINY_LLAMA_CONFIG)

diff --git a/tests/test_optimizer.py b/tests/test_optimizer.py
@@ -14,6 +14,7 @@
 @pytest.mark.parametrize(
     "parametrization_method", [ParametrizationMethod.STANDARD, ParametrizationMethod.SPECTRAL_MUP]
 )
+@pytest.mark.fa2
 @rerun_if_address_is_in_use()
 def test_get_custom_lr(tp: int, dp: int, pp: int, parametrization_method: ParametrizationMethod):
     LR = 1e-3

diff --git a/tests/test_parametrization.py b/tests/test_parametrization.py
@@ -12,6 +12,7 @@
 
 @pytest.mark.parametrize("tp,dp,pp", [(2, 1, 1)])
 @pytest.mark.parametrize("parametrization_method", [ParametrizationMethod.SPECTRAL_MUP])
+@pytest.mark.fa2
 @rerun_if_address_is_in_use()
 def test_parametrization(tp: int, dp: int, pp: int, parametrization_method: ParametrizationMethod):
     if parametrization_method == ParametrizationMethod.STANDARD: