Check device count before running dist tests

tests/unit/common.py

@@ -108,6 +108,10 @@ def _get_fixture_kwargs(self, request, func):
         return fixture_kwargs
 
     def _launch_procs(self, num_procs):
+        if torch.cuda.is_available() and torch.cuda.device_count() < num_procs:
+            pytest.skip(
+                f"Skipping test because not enough GPUs are available: {num_procs} required, {torch.cuda.device_count()} available"
+            )
         mp.set_start_method('forkserver', force=True)
         skip_msg = mp.Queue()  # Allows forked processes to share pytest.skip reason
         processes = []