[release/2.14] Keep FSDP/HSDP/TP tests within the RCCL /dev/shm budget (#200632)

test/distributed/_composable/fsdp/test_fully_shard_training.py

@@ -59,7 +59,9 @@
 from torch.testing._internal.common_utils import (
     device_sleep,
     get_cycles_per_ms,
+    instantiate_parametrized_tests,
     MI200_ARCH,
+    parametrize,
     run_tests,
     skipIfRocm,
     skipIfTorchInductor,
@@ -2051,15 +2053,18 @@ def _test_2d_mlp_with_nd_mesh(
             self.assertEqual(p.device_mesh.mesh_dim_names, ("dp", "tp"))
 
     @skip_if_lt_x_gpu(8)
-    def test_shard_placement_fn_tp_ep(self):
+    @parametrize("tp_degree", [1, 2])
+    @parametrize("dp_replicate", [1, 2])
+    def test_shard_placement_fn_tp_ep(self, tp_degree: int, dp_replicate: int):
         # Every new mesh adds NCCL communicators that hold /dev/shm until the
         # processes exit, and the meshes depend only on tp_degree and
-        # dp_replicate, so build them once rather than once per subtest.
+        # dp_replicate, so build them once rather than once per subtest, and
+        # give each configuration its own test (and so its own processes).
         self._parallel_meshes = {}
         self.run_subtests(
             {
-                "tp_degree": [1, 2],
-                "dp_replicate": [1, 2],
+                "tp_degree": [tp_degree],
+                "dp_replicate": [dp_replicate],
                 "reshard_non_layer_modules": [False, True, 2],
             },
             self._test_shard_placement_fn_tp_ep,
@@ -2744,5 +2749,7 @@ def test_two_layer_fully_shard_cudagraph(self):
                 model.zero_grad(set_to_none=True)
 
 
+instantiate_parametrized_tests(TestFullyShardNDTraining)
+
 if __name__ == "__main__":
     run_tests()

test/distributed/fsdp/test_fsdp_hybrid_shard.py

@@ -343,6 +343,12 @@ def patched_collective(orig_collective, counter, *args, **kwargs):
             self.assertEqual(num_hsdp_flat_params, cntr[orig_ar])
             self.assertEqual(num_flat_params, cntr[orig_rs])
 
+        if not use_device_mesh:
+            # FSDP created these groups for this model only, and each one holds
+            # /dev/shm until the workers exit, so free them with the model.
+            for pg in intra_node_pgs | inter_node_pgs:
+                dist.destroy_process_group(pg)
+
     @skip_if_lt_x_gpu(4)
     def test_fsdp_hybrid_shard_parity(self):
         # Build the HSDP groups once rather than once per subtest; see

test/distributed/fsdp/test_fsdp_tp_integration.py

@@ -2,6 +2,7 @@
 import copy
 import sys
 from collections import OrderedDict
+from typing import ClassVar
 
 import torch
 from torch import distributed as dist
@@ -85,6 +86,19 @@ def prepare_output_fn(mod, outputs, device_mesh):
 
 
 class TestTPFSDPIntegration(FSDPTestContinuous):
+    # The workers are reused for every test in this class, and each new mesh adds
+    # NCCL communicators that hold /dev/shm until the workers exit, so build the
+    # (dp, tp) mesh once and share it across tests.
+    _dp_tp_mesh: ClassVar[DeviceMesh | None] = None
+
+    @classmethod
+    def _get_dp_tp_mesh(cls) -> DeviceMesh:
+        if cls._dp_tp_mesh is None:
+            cls._dp_tp_mesh = init_device_mesh(
+                device_type, (cls.world_size // 2, 2), mesh_dim_names=("dp", "tp")
+            )
+        return cls._dp_tp_mesh
+
     def _get_params_and_sharding_info(
         self,
         model: SimpleModel,
@@ -282,11 +296,7 @@ def _test_fsdp_tp_integration(
             sharding_strategy=sharding_strategy,
             use_orig_params=use_orig_params,
         )
-        mesh_2d = init_device_mesh(
-            device_type,
-            (self.world_size // tensor_parallel_size, tensor_parallel_size),
-            mesh_dim_names=["dp", "tp"],
-        )
+        mesh_2d = self._get_dp_tp_mesh()
         # Shard with TP and then wrap with FSDP
         sequence_parallelize_plan = {
             "net1": ColwiseParallel(input_layouts=Shard(0)),
@@ -365,9 +375,7 @@ def test_fsdp_tp_extension_grad(self):
         """
         Tests TP + FSDP extension with correct gradient (i.e. no ACT)
         """
-        mesh_2d = init_device_mesh(
-            device_type, (self.world_size // 2, 2), mesh_dim_names=["dp", "tp"]
-        )
+        mesh_2d = self._get_dp_tp_mesh()
 
         class TestModel(torch.nn.Module):
             def __init__(self) -> None:
@@ -417,9 +425,7 @@ def forward(self, x):
 
     @skip_if_lt_x_gpu(4)
     def test_fsdp_tp_sync_module_state(self):
-        mesh_2d = init_device_mesh(
-            device_type, (self.world_size // 2, 2), mesh_dim_names=["dp", "tp"]
-        )
+        mesh_2d = self._get_dp_tp_mesh()
         tp_mesh = mesh_2d["tp"]
         dp_mesh = mesh_2d["dp"]
 

test/distributed/fsdp/test_hsdp_dtensor_state_dict.py

@@ -2,12 +2,13 @@
 
 import io
 from copy import deepcopy
+from typing import ClassVar
 
 import torch
 import torch.distributed as dist
 import torch.nn as nn
 from torch.distributed._shard.sharded_tensor import ShardedTensor
-from torch.distributed.device_mesh import init_device_mesh
+from torch.distributed.device_mesh import DeviceMesh, init_device_mesh
 from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
 from torch.distributed.fsdp.api import (
     ShardedOptimStateDictConfig,
@@ -49,6 +50,17 @@ def get_input(self, device):
 
 # TODO: Consolidate DeviceMesh based FSDP and HSDP test cases.
 class TestHSDPWithDeviceMeshAndDTensor(DTensorContinuousTestBase):
+    # The workers are reused for every test in this class, and each new mesh adds
+    # NCCL communicators that hold /dev/shm until the workers exit, so build the
+    # 2D mesh once and share it across tests.
+    _mesh_2d: ClassVar[DeviceMesh | None] = None
+
+    def _get_mesh_2d(self) -> DeviceMesh:
+        cls = type(self)
+        if cls._mesh_2d is None:
+            cls._mesh_2d = init_device_mesh(self.device_type, (2, self.world_size // 2))
+        return cls._mesh_2d
+
     def _create_model(self, device, device_mesh=None):
         device_id = device_type
         if device_mesh:
@@ -58,7 +70,7 @@ def _create_model(self, device, device_mesh=None):
                 sharding_strategy=ShardingStrategy.HYBRID_SHARD,
             )
         else:
-            mesh_2d = init_device_mesh(self.device_type, (2, self.world_size // 2))
+            mesh_2d = self._get_mesh_2d()
             intra_node_pg = mesh_2d.get_group(mesh_dim=1)
             inter_node_pg = mesh_2d.get_group(mesh_dim=0)
             model = FSDP(
@@ -77,7 +89,7 @@ def _create_model(self, device, device_mesh=None):
     @skip_if_lt_x_gpu(4)
     def test_hsdp_init_with_device_mesh(self, device):
         device_id = device_type
-        mesh_2d = init_device_mesh(self.device_type, (2, self.world_size // 2))
+        mesh_2d = self._get_mesh_2d()
         model, optim = self._create_model(device_id, mesh_2d)
 
         FSDP.set_state_dict_type(
@@ -112,7 +124,7 @@ def test_hsdp_init_with_device_mesh(self, device):
     @parametrize("offload_to_cpu", [True, False])
     def test_dtensor_sharded_tensor_state_dict_identical(self, device, offload_to_cpu):
         device_id = device_type
-        mesh_2d = init_device_mesh(self.device_type, (2, self.world_size // 2))
+        mesh_2d = self._get_mesh_2d()
 
         model, optim = self._create_model(device_id, mesh_2d)
         FSDP.set_state_dict_type(
@@ -182,7 +194,7 @@ def test_dtensor_sharded_tensor_state_dict_identical(self, device, offload_to_cp
     @parametrize("offload_to_cpu", [True, False])
     def test_dtensor_sharded_optim_load_state_dict(self, device, offload_to_cpu):
         device_id = device_type
-        mesh_2d = init_device_mesh(self.device_type, (2, self.world_size // 2))
+        mesh_2d = self._get_mesh_2d()
         model, optim = self._create_model(device_id, mesh_2d)
 
         FSDP.set_state_dict_type(
@@ -237,7 +249,7 @@ def test_dtensor_sharded_optim_load_state_dict(self, device, offload_to_cpu):
     @parametrize("offload_to_cpu", [True, False])
     def test_dtensor_sharded_model_load_state_dict(self, device, offload_to_cpu):
         device_id = device_type
-        mesh_2d = init_device_mesh(self.device_type, (2, self.world_size // 2))
+        mesh_2d = self._get_mesh_2d()
         model, optim = self._create_model(device_id, mesh_2d)
 
         FSDP.set_state_dict_type(
@@ -299,7 +311,7 @@ def forward(self, x):
                 dist.all_reduce(sparse)
                 return self.dense(sparse)
 
-        mesh_2d = init_device_mesh(self.device_type, (2, self.world_size // 2))
+        mesh_2d = self._get_mesh_2d()
         model = FakeMPModel(device_mesh=mesh_2d).to(device_id)
         optim = torch.optim.Adam(model.parameters(), lr=1e-2)
 

torch/testing/_internal/common_fsdp.py

@@ -746,6 +746,11 @@ def forward(self, *args, **kwargs):
         return self.module(*args, **kwargs)
 
 
+# Each new group adds NCCL communicators that hold /dev/shm until the processes
+# exit, so share one size-1 expert group per parent group across model builds.
+_MOE_EXPERT_GROUPS: dict[dist.ProcessGroup, dist.ProcessGroup] = {}
+
+
 class MixtureOfExperts(NestedWrappedModule):
     def __init__(
         self,
@@ -786,10 +791,11 @@ def __init__(
 
         if wrap_fsdp:
             # we create a process group of size 1 for the expert params
-            expert_group = torch.distributed.new_group(
-                [group.rank()]
-            )  # world size 1 means no shard
-            expert = FSDP(expert, expert_group, **fsdp_kwargs)  # type: ignore[assignment]
+            if group not in _MOE_EXPERT_GROUPS:
+                _MOE_EXPERT_GROUPS[group] = torch.distributed.new_group(
+                    [group.rank()]
+                )  # world size 1 means no shard
+            expert = FSDP(expert, _MOE_EXPERT_GROUPS[group], **fsdp_kwargs)  # type: ignore[assignment]
             shared = FSDP(shared, group, **fsdp_kwargs)  # type: ignore[assignment]
 
         self.module = nn.Sequential(