2025-02-13 nightly release (fd45bdc)

Author: pytorchbot

torchrec/distributed/comm.py

@@ -226,9 +226,14 @@ def intra_and_cross_node_pg_2D(
 
     if _INTRA_PG_2D is None:
         for group_rank in range(step):
-            sharding_pg_peers = [
-                step * r + group_rank for r in range(sharding_group_size)
-            ]
+            if env.use_inter_host_allreduce:
+                # for inter host all reduce, we change the sharding group calculation to be continuous
+                ranks = group_rank * sharding_group_size
+                sharding_pg_peers = list(range(ranks, ranks + sharding_group_size))
+            else:
+                sharding_pg_peers = [
+                    step * r + group_rank for r in range(sharding_group_size)
+                ]
             for group in range(len(sharding_pg_peers) // devices_per_node):
                 intra_pg_peers = sharding_pg_peers[
                     group * devices_per_node : (group + 1) * devices_per_node

torchrec/distributed/embedding.py

@@ -32,6 +32,7 @@
 from torch.distributed._shard.sharding_spec import EnumerableShardingSpec
 from torch.distributed._tensor import DTensor
 from torch.nn.parallel import DistributedDataParallel
+from torchrec.distributed.comm import get_local_size
 from torchrec.distributed.embedding_sharding import (
     EmbeddingSharding,
     EmbeddingShardingInfo,
@@ -69,13 +70,16 @@
     QuantizedCommCodecs,
     ShardedTensor,
     ShardingEnv,
+    ShardingEnv2D,
     ShardMetadata,
 )
 from torchrec.distributed.utils import (
     add_params_from_parameter_sharding,
     convert_to_fbgemm_types,
+    create_global_tensor_shape_stride_from_metadata,
     maybe_annotate_embedding_event,
     merge_fused_params,
+    none_throws,
     optimizer_type_to_emb_opt_type,
 )
 from torchrec.modules.embedding_configs import (
@@ -534,12 +538,9 @@ def __init__(
                 if table_name in self._table_names
             },
         )
-        # output parameters as DTensor in state dict
-        self._output_dtensor: bool = (
-            fused_params.get("output_dtensor", False) if fused_params else False
-        )
-
         self._env = env
+        # output parameters as DTensor in state dict
+        self._output_dtensor: bool = env.output_dtensor
         # TODO get rid of get_ec_index_dedup global flag
         self._use_index_dedup: bool = use_index_dedup or get_ec_index_dedup()
         sharding_type_to_sharding_infos = create_sharding_infos_by_sharding(
@@ -842,6 +843,14 @@ def _initialize_torch_state(self) -> None:  # noqa
                         )
                     )
                 else:
+                    shape, stride = create_global_tensor_shape_stride_from_metadata(
+                        none_throws(self.module_sharding_plan[table_name]),
+                        (
+                            self._env.node_group_size
+                            if isinstance(self._env, ShardingEnv2D)
+                            else get_local_size(self._env.world_size)
+                        ),
+                    )
                     # empty shard case
                     self._model_parallel_name_to_dtensor[table_name] = (
                         DTensor.from_local(
@@ -851,6 +860,8 @@ def _initialize_torch_state(self) -> None:  # noqa
                             ),
                             device_mesh=self._env.device_mesh,
                             run_check=False,
+                            shape=shape,
+                            stride=stride,
                         )
                     )
             else:
@@ -861,7 +872,11 @@ def _initialize_torch_state(self) -> None:  # noqa
                     ShardedTensor._init_from_local_shards(
                         local_shards,
                         self._name_to_table_size[table_name],
-                        process_group=self._env.process_group,
+                        process_group=(
+                            self._env.sharding_pg
+                            if isinstance(self._env, ShardingEnv2D)
+                            else self._env.process_group
+                        ),
                     )
                 )
 

torchrec/distributed/model_parallel.py

@@ -690,6 +690,7 @@ def __init__(
         init_data_parallel: bool = True,
         init_parameters: bool = True,
         data_parallel_wrapper: Optional[DataParallelWrapper] = None,
+        use_inter_host_allreduce: bool = False,
     ) -> None:
         assert device.type == "cuda", "DMPCollection only supports CUDA"
         self._device = device
@@ -705,13 +706,16 @@ def __init__(
                 global_rank=self._global_rank,
                 world_size=world_size,
                 local_size=sharding_group_size,
+                use_inter_host_allreduce=use_inter_host_allreduce,
             )
         )
 
         self._remap_sharding_plan(
             plan=plan,
             rank=self._global_rank,
-            num_nodes=world_size // sharding_group_size,
+            step=world_size // sharding_group_size,
+            sharding_group_size=sharding_group_size,
+            use_inter_host_allreduce=use_inter_host_allreduce,
         )
         super().__init__(
             module,
@@ -720,6 +724,7 @@ def __init__(
                 sharding_pg=self._sharding_pg,
                 device_mesh=self._device_mesh,
                 node_group_size=node_group_size,
+                use_inter_host_allreduce=use_inter_host_allreduce,
             ),
             device,
             plan,
@@ -768,7 +773,11 @@ def sync(self, include_optimizer_state: bool = True) -> None:
                 handle.wait()
 
     def _create_process_groups(
-        self, global_rank: int, world_size: int, local_size: int
+        self,
+        global_rank: int,
+        world_size: int,
+        local_size: int,
+        use_inter_host_allreduce: bool = False,
     ) -> Tuple[DeviceMesh, dist.ProcessGroup, dist.ProcessGroup]:
         """
         Creates process groups for sharding and replication, the process groups
@@ -784,17 +793,29 @@ def _create_process_groups(
                 replication process group, and allreduce process group.
         """
         peer_matrix = []
-        num_nodes = world_size // local_size
+        mesh, sharding_pg, replica_pg = None, None, None
 
-        for group_rank in range(world_size // local_size):
-            peers = [num_nodes * r + group_rank for r in range(local_size)]
-            peer_matrix.append(peers)
+        logger.warning(f"[2D] Use inter host all reduce: {use_inter_host_allreduce}")
+
+        if use_inter_host_allreduce:
+            # We shard on continuous set of ranks and nodes. Thereby forcing our all reduce to be inter host.
+            # Under this scheme sharding types such as TWRW and GRID will now take
+            # advantage of intra node comms as a result of the continuous set of ranks.
+            peer_matrix = [
+                list(range(i, i + local_size)) for i in range(0, world_size, local_size)
+            ]
+        else:
+            step = world_size // local_size
+            for group_rank in range(world_size // local_size):
+                peers = [step * r + group_rank for r in range(local_size)]
+                peer_matrix.append(peers)
 
         mesh = DeviceMesh(
             device_type=self._device.type,
             mesh=peer_matrix,
             mesh_dim_names=("replicate", "shard"),
         )
+
         logger.warning(f"[Connection] 2D Device Mesh created: {mesh}")
         sharding_pg = mesh.get_group(mesh_dim="shard")
         logger.warning(
@@ -808,7 +829,12 @@ def _create_process_groups(
         return mesh, sharding_pg, replica_pg
 
     def _remap_sharding_plan(
-        self, plan: ShardingPlan, rank: int, num_nodes: int
+        self,
+        plan: ShardingPlan,
+        rank: int,
+        step: int,
+        sharding_group_size: int,
+        use_inter_host_allreduce: bool = False,
     ) -> None:
         """
         Remaps the sharding plan to the local replica process group ranks
@@ -822,20 +848,32 @@ def _remap_sharding_plan(
             global_rank (int): The global rank of the current process.
             num_nodes (int): The number of nodes.
         """
-
-        group_start = rank % num_nodes
+        group_start = rank % step
         for key in plan.plan:
             # pyre-ignore[16]
             for _, param_sharding in plan.plan[key].items():
                 new_ranks = []
-                for shard_rank in param_sharding.ranks:
-                    new_ranks.append(shard_rank * num_nodes + group_start)
+                if use_inter_host_allreduce:
+                    group = rank // sharding_group_size
+                    new_ranks = [
+                        shard_rank + (group * sharding_group_size)
+                        for shard_rank in param_sharding.ranks
+                    ]
+                else:
+                    for shard_rank in param_sharding.ranks:
+                        new_ranks.append(shard_rank * step + group_start)
                 param_sharding.ranks = new_ranks
+
                 if isinstance(param_sharding.sharding_spec, EnumerableShardingSpec):
                     shards = param_sharding.sharding_spec.shards
                     if shards is not None:
                         for shard in shards:
-                            shard_rank = shard.placement._rank * num_nodes + group_start
+                            if use_inter_host_allreduce:
+                                shard_rank = shard.placement._rank + (
+                                    (rank // sharding_group_size) * sharding_group_size
+                                )
+                            else:
+                                shard_rank = shard.placement._rank * step + group_start
                             shard.placement = _remote_device(
                                 f"rank:{shard_rank}/cuda:{shard_rank % get_local_size()}"
                             )

torchrec/distributed/sharding/cw_sharding.py

@@ -45,6 +45,7 @@
     QuantizedCommCodecs,
     ShardedTensorMetadata,
     ShardingEnv,
+    ShardingType,
     ShardMetadata,
 )
 from torchrec.distributed.utils import none_throws
@@ -191,7 +192,7 @@ def _shard(
             for i, rank in enumerate(info.param_sharding.ranks):
                 # Remap rank by number of replica groups if 2D parallelism is enabled
                 rank = (
-                    rank // self._env.num_sharding_groups()  # pyre-ignore[16]
+                    self._env.remap_rank(rank, ShardingType.COLUMN_WISE)  # pyre-ignore[16]
                     if self._is_2D_parallel
                     else rank
                 )

torchrec/distributed/sharding/grid_sharding.py

@@ -250,7 +250,7 @@ def _shard(
             # pyre-fixme [6]
             for i, rank in enumerate(info.param_sharding.ranks):
                 rank = (
-                    rank // self._env.num_sharding_groups()  # pyre-ignore[16]
+                    self._env.remap_rank(rank, ShardingType.GRID_SHARD)  # pyre-ignore[16]
                     if self._is_2D_parallel
                     else rank
                 )

torchrec/distributed/sharding/tw_sharding.py

@@ -48,6 +48,7 @@
     ShardedTensorMetadata,
     ShardingEnv,
     ShardingEnv2D,
+    ShardingType,
     ShardMetadata,
 )
 from torchrec.distributed.utils import none_throws
@@ -128,7 +129,7 @@ def _shard(
             )
 
             dtensor_metadata = None
-            if info.fused_params.get("output_dtensor", False):  # pyre-ignore[16]
+            if self._env.output_dtensor:
                 dtensor_metadata = DTensorMetadata(
                     mesh=(
                         self._env.device_mesh["replicate"]  # pyre-ignore[16]
@@ -142,12 +143,12 @@ def _shard(
                     ),
                     stride=info.param.stride(),
                 )
-            # to not pass onto TBE
-            info.fused_params.pop("output_dtensor", None)  # pyre-ignore[16]
 
             rank = (
                 # pyre-ignore [16]
-                info.param_sharding.ranks[0] // self._env.num_sharding_groups()
+                self._env.remap_rank(
+                    info.param_sharding.ranks[0], ShardingType.TABLE_WISE  # pyre-ignore[16]
+                )
                 if self._is_2D_parallel
                 else info.param_sharding.ranks[0]
             )

torchrec/distributed/sharding/twrw_sharding.py

@@ -13,7 +13,7 @@
 
 import torch
 import torch.distributed as dist
-from torch.distributed._tensor import Shard
+from torch.distributed._tensor import Replicate, Shard
 from torch.distributed.distributed_c10d import get_process_group_ranks
 from torchrec.distributed.comm import (
     get_local_size,
@@ -165,10 +165,11 @@ def _shard(
 
             dtensor_metadata = None
             if self._env.output_dtensor:
-                placements = (Shard(0),)
                 dtensor_metadata = DTensorMetadata(
                     mesh=self._env.device_mesh,
-                    placements=placements,
+                    placements=(
+                        (Replicate(), Shard(1)) if self._is_2D_parallel else (Shard(1),)
+                    ),
                     size=(
                         info.embedding_config.num_embeddings,
                         info.embedding_config.embedding_dim,

torchrec/distributed/test_utils/test_model_parallel.py

@@ -149,6 +149,7 @@ def _test_sharding(
         global_constant_batch: bool = False,
         pooling: PoolingType = PoolingType.SUM,
         data_type: DataType = DataType.FP32,
+        use_inter_host_allreduce: bool = False,
     ) -> None:
         self._build_tables_and_groups(data_type=data_type)
         self._run_multi_process_test(
@@ -170,6 +171,7 @@ def _test_sharding(
             apply_optimizer_in_backward_config=apply_optimizer_in_backward_config,
             variable_batch_per_feature=variable_batch_per_feature,
             global_constant_batch=global_constant_batch,
+            use_inter_host_allreduce=use_inter_host_allreduce,
         )
 
 

torchrec/distributed/test_utils/test_sharding.py

@@ -315,6 +315,7 @@ def sharding_single_rank_test(
     global_constant_batch: bool = False,
     world_size_2D: Optional[int] = None,
     node_group_size: Optional[int] = None,
+    use_inter_host_allreduce: bool = False,
     input_type: str = "kjt",  # "kjt" or "td"
 ) -> None:
     with MultiProcessContext(rank, world_size, backend, local_size) as ctx:
@@ -432,6 +433,7 @@ def sharding_single_rank_test(
                 plan=plan,
                 sharders=sharders,
                 device=ctx.device,
+                use_inter_host_allreduce=use_inter_host_allreduce,
             )
         else:
             local_model = DistributedModelParallel(