[Distributed] Limit world_size to 8 for FSDP Unit tests (#103412)

lcskrishna · pytorchmergebot · commit 78b84655659d · 2023-11-08T17:21:38.000Z
There are few unit tests in FSDP that can support upto 8 GPUs. In this case, for example test_fsdp_uneven has an input size of [8,3]. For each process/rank we pass the data as input[self.rank] as below. So when we use 16 GPUs for our tests, these tests throw an index/key error. So basically to avoid such corner cases, I would like to add this change to use 8GPUs if there are more than 8 GPUs. This is applicable to both ROCm and CUDA builds as well. https://github.com/pytorch/pytorch/blob/main/test/distributed/fsdp/test_fsdp_uneven.py#L44 https://github.com/pytorch/pytorch/blob/main/test/distributed/fsdp/test_fsdp_uneven.py#L55 Pull Request resolved: #103412 Approved by: https://github.com/jithunnair-amd, https://github.com/pruthvistony, https://github.com/malfet
diff --git a/torch/testing/_internal/common_fsdp.py b/torch/testing/_internal/common_fsdp.py
@@ -873,7 +873,7 @@ def setUp(self):
 
     @property
     def world_size(self):
-        return torch.cuda.device_count() if torch.cuda.is_available() else 4
+        return min(torch.cuda.device_count(), 8) if torch.cuda.is_available() else 4
 
     @property
     def process_group(self):
