allow multi gpu tests to be run (#765)

Summary:
Pull Request resolved: #765

Tests were always skipped since their name does not follow the `test_*.py` or `*_test.py` format.

Instead, skip tests in CPU/single-GPU environments using the `unittest.skipIf` decorator.

Reviewed By: HuanyuZhang

Differential Revision: D76141869

fbshipit-source-id: b8009714e02e702f865012d5606c2de05265f27c

Author: iden-kalemaj

opacus/tests/multigpu_adaptive_clipping.py renamed to opacus/tests/multigpu_adaptive_clipping_test.py

@@ -67,7 +67,7 @@ def forward(self, x):
 
 def demo_basic(rank, weight, world_size, dp):
     torch.manual_seed(world_size)
-    batch_size = 32
+    batch_size = 2
     setup(rank, world_size)
 
     # create model and move it to GPU with id rank
@@ -129,7 +129,8 @@ def run_demo(demo_fn, weight, world_size, dp):
 
 
 class GradientComputationTestAdaptiveClipping(unittest.TestCase):
-    def test_gradient_correct_adaptive(self) -> None:
+    @unittest.skipIf(torch.cuda.device_count() < 2, "Need at least 2 GPUs")
+    def test_adaptive_gradient_correct(self) -> None:
 
         # Tests that gradient is the same with DP or without DP in the distributed setting
         n_gpus = torch.cuda.device_count()

opacus/tests/multigpu_gradcheck.py renamed to opacus/tests/multigpu_gradcheck_test.py

@@ -112,7 +112,7 @@ def run_ghost_clipping_test(
 
 def demo_basic(rank, weight, world_size, dp, clipping, grad_sample_mode):
     torch.manual_seed(world_size)
-    batch_size = 32
+    batch_size = 2
     setup(rank, world_size)
 
     # create model and move it to GPU with id rank
@@ -188,6 +188,7 @@ def run_demo(demo_fn, weight, world_size, dp, clipping, grad_sample_mode):
 
 
 class GradientComputationTest(unittest.TestCase):
+    @unittest.skipIf(torch.cuda.device_count() < 2, "Need at least 2 GPUs")
     def test_gradient_correct(self) -> None:
         # Tests that gradient is the same with DP or without DDP
         n_gpus = torch.cuda.device_count()