diff --git a/tests/models/utils/test_tp_strategy.py b/tests/models/utils/test_tp_strategy.py index 37c1a2941d..cf3107cc1c 100644 --- a/tests/models/utils/test_tp_strategy.py +++ b/tests/models/utils/test_tp_strategy.py @@ -11,11 +11,11 @@ from llmfoundry.models.mpt.modeling_mpt import ComposerMPTCausalLM from llmfoundry.utils.builders import build_tp_strategy - +@pytest.mark.gpu def test_ffn_tp_strategy_layer_plan(): # Actual layer plan from tp_strategy=fnn tp_config = { - 'strategy': 'ffn', + 'strategy': 'ffn' } model_cfg = {