Disable fsdp1 tests on single gpu

Signed-off-by: Keval Morabia <[email protected]>

Author: kevalmorabia97

tests/examples/llm_qat/test_llm_qat.py

@@ -15,6 +15,7 @@
 
 
 import pytest
+import torch
 from _test_utils.examples.run_command import run_example_command
 
 
@@ -37,7 +38,7 @@ def _run_command(extra_cmd_args: list[str]):
     )
 
 @pytest.mark.parametrize("backend", [
-    "fsdp1",
+    pytest.param("fsdp1", marks=pytest.mark.skipif(torch.cuda.device_count() < 2, reason="need 2 GPUs!")),
     "fsdp2",
     "deepspeed",
     "ddp",
@@ -68,7 +69,7 @@ def test_llama_qat_int4w_int8a(tiny_llama_path, tmp_path, backend):
     )
 
 @pytest.mark.parametrize("backend", [
-    "fsdp1",
+    pytest.param("fsdp1", marks=pytest.mark.skipif(torch.cuda.device_count() < 2, reason="need 2 GPUs!")),
     "fsdp2",
     "deepspeed",
     "ddp",

tests/examples/speculative_decoding/test_medusa.py

@@ -67,6 +67,7 @@ def test_llama_medusa_fp8_qat(tiny_llama_path, num_gpus, tiny_daring_anteater_pa
             "--output_dir", tmp_path / "medusa-tinyllama-qat-finetune",
             "--quant_cfg", "FP8_DEFAULT_CFG",
             "--calib_size", "64",
+            "--backend", "fsdp2",
         ],
         "llm_qat",
     )