Fully Sharded 2D Parallelism (#3558)

Summary:

**This diff introduces Fully Sharded 2D Parallelism in TorchRec. It brings forth significant memory (50%+) savings by sharding embedding tables when they are not in use.**

After the embedding lookup, the embedding table is further sharded across the data parallel dimension until it is needed in the backward pass. This allows model layers after the embedding lookup to have more memory headroom. Enabling further scaling of the dense architecture. **Practically speaking, this saves 50%+ embedding memory per GPU which account for upwards of 10GB of memory saving on large models.**

The peak memory during this step becomes, ```O(shard + shard/num_replication)```, which then leads to an embedding memory of ```O(shard/num_replication)``` after the lookup step.

The memory free and collective communications are done in a overhead free manner by maximizing computation and communication collectives through asynchronous handling on multiple streams.

With Fully Sharded 2D, the embedding weight synchronization has to happen every step or trained batches are lost across ranks. We use an asynchronous reduce scatter after the embedding lookup step. We are able to fully overlap this collective with compute to expose no additional overhead.

A new awaitable is introduced, ```ReduceScatterResizeAwaitable``` under the Fully Sharded path that is called with SDD output_dist all to all. This awaitable ```wait()```s on the async reduce scatter and calls the ```resize()``` operation on the embedding memory ensuring no race conditions.

Users can enable fully sharded 2D through, a new arg `ShardingStrategy`
```
DMPCollection(..., sharding_strategy=ShardingStrategy.FULLY_SHARDED)
```

This is part of our work to create an overhead free 2D parallel which will allow us to use it for every model.

Remaining work from this diff is to launch an async all gather in the backward pass, making planner aware of such memory savings, and integrate this work with per module 2D

Differential Revision: D82253387

Author: iamzainhuda

torchrec/distributed/batched_embedding_kernel.py

@@ -343,6 +343,7 @@ def __init__(
         module_fqn: Optional[str] = None,
         sharding_types: Optional[List[str]] = None,
         use_gather_select: bool = False,
+        resize_awaitables: Optional[List[Awaitable[torch.Tensor]]] = None,
     ) -> None:
         super().__init__()
         self._awaitables_per_sharding = awaitables_per_sharding
@@ -354,6 +355,7 @@ def __init__(
         self._module_fqn = module_fqn
         self._sharding_types = sharding_types
         self._use_gather_select = use_gather_select
+        self._resize_awaitables = resize_awaitables
 
     def _wait_impl(self) -> Dict[str, JaggedTensor]:
         jt_dict: Dict[str, JaggedTensor] = {}
@@ -398,6 +400,12 @@ def _wait_impl(self) -> Dict[str, JaggedTensor]:
                     use_gather_select=self._use_gather_select,
                 )
             )
+
+        # free memory and resize
+        # pyre-ignore[16]
+        for awaitable in self._resize_awaitables:
+            awaitable.wait()
+
         return jt_dict
 
 
@@ -1588,6 +1596,8 @@ def compute_and_output_dist(
     ) -> LazyAwaitable[Dict[str, JaggedTensor]]:
         awaitables_per_sharding: List[Awaitable[torch.Tensor]] = []
         features_before_all2all_per_sharding: List[KeyedJaggedTensor] = []
+        resize_awaitables = []
+
         for lookup, odist, features, sharding_ctx, sharding_type in zip(
             self._lookups,
             self._output_dists,
@@ -1604,6 +1614,9 @@ def compute_and_output_dist(
                 EmbeddingEvent.LOOKUP, self._module_fqn, sharding_type
             ):
                 embs = lookup(features)
+                if hasattr(lookup, "get_resize_awaitables"):
+                    # pyre-ignore[29]
+                    resize_awaitables.extend(lookup.get_resize_awaitables())
                 if self.post_lookup_tracker_fn is not None:
                     self.post_lookup_tracker_fn(features, embs, self, None)
 
@@ -1631,6 +1644,7 @@ def compute_and_output_dist(
             module_fqn=self._module_fqn,
             sharding_types=list(self._sharding_type_to_sharding.keys()),
             use_gather_select=self._use_gather_select,
+            resize_awaitables=resize_awaitables,
         )
 
     def _embedding_dim_for_sharding_type(self, sharding_type: str) -> int:

torchrec/distributed/embedding.py

@@ -39,6 +39,8 @@
     BatchedFusedEmbeddingBag,
     KeyValueEmbedding,
     KeyValueEmbeddingBag,
+    ShardedBatchedFusedEmbedding,
+    ShardedBatchedFusedEmbeddingBag,
     ZeroCollisionEmbeddingCache,
     ZeroCollisionKeyValueEmbedding,
     ZeroCollisionKeyValueEmbeddingBag,
@@ -65,7 +67,15 @@
     QuantBatchedEmbedding,
     QuantBatchedEmbeddingBag,
 )
-from torchrec.distributed.types import rank_device, ShardedTensor, ShardingType
+from torchrec.distributed.types import (
+    LazyAwaitable,
+    rank_device,
+    ShardedTensor,
+    ShardingEnv,
+    ShardingEnv2D,
+    ShardingStrategy,
+    ShardingType,
+)
 from torchrec.sparse.jagged_tensor import KeyedJaggedTensor
 
 logger: logging.Logger = logging.getLogger(__name__)
@@ -185,12 +195,15 @@ def __init__(
         grouped_configs: List[GroupedEmbeddingConfig],
         pg: Optional[dist.ProcessGroup] = None,
         device: Optional[torch.device] = None,
+        env: Optional[ShardingEnv] = None,
     ) -> None:
         super().__init__()
         self._emb_modules: nn.ModuleList = nn.ModuleList()
         self._need_prefetch: bool = False
         for config in grouped_configs:
-            self._emb_modules.append(self._create_embedding_kernel(config, pg, device))
+            self._emb_modules.append(
+                self._create_embedding_kernel(config, pg, device, env)
+            )
 
         self._feature_splits: List[int] = []
         for config in grouped_configs:
@@ -218,6 +231,7 @@ def _create_embedding_kernel(
         config: GroupedEmbeddingConfig,
         pg: Optional[dist.ProcessGroup],
         device: Optional[torch.device],
+        env: Optional[ShardingEnv] = None,
     ) -> BaseEmbedding:
         for table in config.embedding_tables:
             if (
@@ -234,11 +248,20 @@ def _create_embedding_kernel(
                 device=device,
             )
         elif config.compute_kernel == EmbeddingComputeKernel.FUSED:
-            return BatchedFusedEmbedding(
-                config=config,
-                pg=pg,
-                device=device,
-            )
+            if (
+                env
+                and isinstance(env, ShardingEnv2D)
+                and env.sharding_strategy == ShardingStrategy.FULLY_SHARDED
+            ):
+                return ShardedBatchedFusedEmbedding(
+                    config=config, pg=pg, device=device, env=env
+                )
+            else:
+                return BatchedFusedEmbedding(
+                    config=config,
+                    pg=pg,
+                    device=device,
+                )
         elif config.compute_kernel == EmbeddingComputeKernel.KEY_VALUE:
             return KeyValueEmbedding(
                 config=config,
@@ -329,6 +352,14 @@ def forward(
 
         return embeddings_cat_empty_rank_handle(embeddings, self._dummy_embs_tensor)
 
+    def get_resize_awaitables(self) -> List[LazyAwaitable[torch.Tensor]]:
+        # TODO - we can probably do some smart grouping to make this more efficient
+        return [
+            emb_module.get_rs_awaitable()  # pyre-ignore[29]
+            for emb_module in self._emb_modules
+            if hasattr(emb_module, "get_rs_awaitable")
+        ]
+
     # pyre-fixme[14]: `state_dict` overrides method defined in `Module` inconsistently.
     def state_dict(
         self,
@@ -512,12 +543,14 @@ def __init__(
         feature_processor: Optional[BaseGroupedFeatureProcessor] = None,
         scale_weight_gradients: bool = True,
         sharding_type: Optional[ShardingType] = None,
+        env: Optional[ShardingEnv] = None,
     ) -> None:
         super().__init__()
+        self._env = env
         self._emb_modules: nn.ModuleList = nn.ModuleList()
         for config in grouped_configs:
             self._emb_modules.append(
-                self._create_embedding_kernel(config, device, pg, sharding_type)
+                self._create_embedding_kernel(config, device, pg, sharding_type, env)
             )
 
         self._feature_splits: List[int] = []
@@ -555,6 +588,7 @@ def _create_embedding_kernel(
         device: Optional[torch.device],
         pg: Optional[dist.ProcessGroup],
         sharding_type: Optional[ShardingType],
+        env: Optional[ShardingEnv],
     ) -> BaseEmbedding:
         if config.compute_kernel == EmbeddingComputeKernel.DENSE:
             return BatchedDenseEmbeddingBag(
@@ -564,12 +598,26 @@ def _create_embedding_kernel(
                 sharding_type=sharding_type,
             )
         elif config.compute_kernel == EmbeddingComputeKernel.FUSED:
-            return BatchedFusedEmbeddingBag(
-                config=config,
-                pg=pg,
-                device=device,
-                sharding_type=sharding_type,
-            )
+            if (
+                env
+                and isinstance(env, ShardingEnv2D)
+                and env.sharding_strategy == ShardingStrategy.FULLY_SHARDED
+            ):
+                return ShardedBatchedFusedEmbeddingBag(
+                    config=config,
+                    pg=pg,
+                    device=device,
+                    sharding_type=sharding_type,
+                    env=env,
+                )
+            else:
+                return BatchedFusedEmbeddingBag(
+                    config=config,
+                    pg=pg,
+                    device=device,
+                    sharding_type=sharding_type,
+                    env=env,
+                )
         elif config.compute_kernel in {
             EmbeddingComputeKernel.KEY_VALUE,
         }:
@@ -744,6 +792,14 @@ def forward(
             dim=1,
         )
 
+    def get_resize_awaitables(self) -> List[LazyAwaitable[torch.Tensor]]:
+        # TODO - we can probably do some smart grouping to make this more efficient
+        return [
+            emb_module.get_rs_awaitable()  # pyre-ignore[29]
+            for emb_module in self._emb_modules
+            if hasattr(emb_module, "get_rs_awaitable")
+        ]
+
     # pyre-fixme[14]: `state_dict` overrides method defined in `Module` inconsistently.
     def state_dict(
         self,

torchrec/distributed/embedding_lookup.py

@@ -407,13 +407,15 @@ def __init__(
         embedding_names: List[str],
         module_fqn: Optional[str] = None,
         sharding_types: Optional[List[str]] = None,
+        resize_awaitables: Optional[List[Awaitable[torch.Tensor]]] = None,
     ) -> None:
         super().__init__()
         self._awaitables = awaitables
         self._embedding_dims = embedding_dims
         self._embedding_names = embedding_names
         self._module_fqn = module_fqn
         self._sharding_types = sharding_types
+        self._resize_awaitables = resize_awaitables
 
     def _wait_impl(self) -> KeyedTensor:
         embeddings = []
@@ -425,6 +427,12 @@ def _wait_impl(self) -> KeyedTensor:
             ):
                 embeddings.append(w.wait())
 
+        # free memory and resize
+        if self._resize_awaitables is not None:
+            # pyre-ignore[16]
+            for awaitable in self._resize_awaitables:
+                awaitable.wait()
+
         return construct_output_kt(
             embeddings=embeddings,
             embedding_names=self._embedding_names,
@@ -1655,6 +1663,7 @@ def compute_and_output_dist(
         """
         batch_size_per_feature_pre_a2a = []
         awaitables = []
+        resize_awaitables = []
 
         # No usage of zip for dynamo
         for i in range(len(self._lookups)):
@@ -1669,7 +1678,11 @@ def compute_and_output_dist(
                 self._module_fqn,
                 sharding_type,
             ):
+                # with fully sharded 2D enabled, it returns an awaitable for the reduce scatter and resize operation
                 embs = lookup(features)
+                if hasattr(lookup, "get_resize_awaitables"):
+                    # pyre-ignore[29]
+                    resize_awaitables.extend(lookup.get_resize_awaitables())
                 if self.post_lookup_tracker_fn is not None:
                     self.post_lookup_tracker_fn(features, embs, self, None)
 
@@ -1710,6 +1723,7 @@ def compute_and_output_dist(
                 embedding_names=self._embedding_names,
                 module_fqn=self._module_fqn,
                 sharding_types=self._sharding_types,
+                resize_awaitables=resize_awaitables,
             )
 
         # register callback if there are features that need mean pooling