Replacing ShardedTensor with DTensor for RW sharding (#2147)

Summary:
Pull Request resolved: #2147

**This is the first part of migration TorchRec state dict checkpointing from ShardedTensor to DTensor. It sets up the necessary infra to support additional sharding schemes. The general approach is to keep ShardedTensor paths and remove them once all sharding types are supported on DTensor. This includes ShardingPlan and ShardedTensor dataclasses such as ShardedTensorMetadata. Those will be migrated in a separate diff with ParameterSharding**

NOTE: This version of LocalShardsWrapper does not support empty shards, that is added in the next diff enabling CW. D57063512

**This diff includes:**
+ LocalShardsWrapper torch.tensor subclass to be used with DTensor
+ Changes in TorchRec state_dict load and creation to use DTensor for Row Wise path in both EmbeddingCollection and EmbeddingBagCollection
+ Changes to DCP to support LocalShardsWrapper for saving and reading (WriteItems and ReadItems)
+ Added DTensor paths to callsites where ShardedTensors are expected.

**LocalShardsWrapper supports the following torch ops:**
+ torch.ops._c10d_functional.all_gather_into_tensor.default
+ aten._to_copy.default
+ aten.view.default
+ aten.equal.default
+ aten.detach.default

With extensibility to add more as required by use cases.

See https://docs.google.com/document/d/16Ptl50mGFJW2cljdF2HQ6FwsiA0scwbAbjx_4dhabJw/edit?usp=drivesdk for more info regarding design and approach.

Reviewed By: XilunWu

Differential Revision: D54375878

Author: iamzainhuda

torchrec/distributed/composable/tests/test_embedding.py

@@ -14,6 +14,7 @@
 import torch
 import torch.nn as nn
 from hypothesis import given, settings, Verbosity
+from torch.distributed._tensor.api import DTensor
 from torch.distributed.optim import (
     _apply_optimizer_in_backward as apply_optimizer_in_backward,
 )
@@ -177,6 +178,8 @@ def _test_sharding(  # noqa C901
             )
             if isinstance(sharded_state, ShardedTensor):
                 sharded_state.gather(out=sharded_param)
+            elif isinstance(sharded_state, DTensor):
+                sharded_param = sharded_state.full_tensor()
             else:
                 sharded_param = sharded_state
 

torchrec/distributed/composable/tests/test_embeddingbag.py

@@ -18,6 +18,7 @@
 import torch.nn as nn
 
 from hypothesis import assume, given, settings, Verbosity
+from torch.distributed._tensor.api import DTensor
 from torch.distributed.optim import (
     _apply_optimizer_in_backward as apply_optimizer_in_backward,
 )
@@ -238,7 +239,11 @@ def _test_sharding(  # noqa C901
                     if ctx.rank == 0
                     else None
                 )
-                sharded_state.gather(out=out)
+                if isinstance(sharded_state, DTensor):
+                    out = sharded_state.full_tensor()
+                else:
+                    sharded_state.gather(out=out)
+
                 if ctx.rank == 0:
                     torch.testing.assert_close(
                         unsharded_state,

torchrec/distributed/composable/tests/test_fsdp.py

@@ -16,6 +16,7 @@
 from torch import nn
 from torch.distributed._composable import fully_shard
 from torch.distributed._shard.sharded_tensor import ShardedTensor
+from torch.distributed._tensor import DTensor
 
 from torch.distributed.checkpoint import (
     FileSystemReader,
@@ -193,6 +194,10 @@ def _run(  # noqa
                         if not p.local_shards():
                             continue
                         p = p.local_tensor()
+                    if isinstance(p, DTensor):
+                        if not p.to_local().local_shards():
+                            continue
+                        p = p.to_local().local_shards()[0]
                     p_sum += p.sum()
                     p.zero_()
                     assert p.sum() == 0
@@ -205,6 +210,10 @@ def _run(  # noqa
                         if not t.local_shards():
                             continue
                         t = t.local_tensor()
+                    if isinstance(t, DTensor):
+                        if not t.to_local().local_shards():  # pyre-ignore[16]
+                            continue
+                        t = t.to_local().local_shards()[0]
                     o_sum += t.sum()
                     t.zero_()
                     assert t.sum() == 0
@@ -228,6 +237,10 @@ def _run(  # noqa
                             continue
                         p = p.local_tensor()
                     p_sum_loaded += p.sum()
+                    if isinstance(p, DTensor):
+                        if not p.to_local().local_shards():
+                            continue
+                        p = p.to_local().local_shards()[0]
             assert p_sum.allclose(p_sum_loaded)
 
             o_sum_loaded = torch.zeros(1, device=ctx.device)
@@ -239,6 +252,10 @@ def _run(  # noqa
                         if not t.local_shards():
                             continue
                         t = t.local_tensor()
+                    if isinstance(t, DTensor):
+                        if not t.to_local().local_shards():
+                            continue
+                        t = t.to_local().local_shards()[0]
                     o_sum_loaded += t.sum()
             assert o_sum.allclose(o_sum_loaded)
 

torchrec/distributed/embedding.py

@@ -28,6 +28,7 @@
 import torch
 from torch import distributed as dist, nn
 from torch.autograd.profiler import record_function
+from torch.distributed._tensor import DTensor
 from torch.nn.parallel import DistributedDataParallel
 from torchrec.distributed.embedding_sharding import (
     EmbeddingSharding,
@@ -55,6 +56,7 @@
 from torchrec.distributed.sharding.tw_sequence_sharding import (
     TwSequenceEmbeddingSharding,
 )
+from torchrec.distributed.shards_wrapper import LocalShardsWrapper
 from torchrec.distributed.types import (
     Awaitable,
     EmbeddingModuleShardingPlan,
@@ -601,18 +603,20 @@ def _pre_load_state_dict_hook(
     ) -> None:
         """
         Modify the destination state_dict for model parallel
-        to transform from ShardedTensors into tensors
+        to transform from ShardedTensors/DTensors into tensors
         """
-        for (
-            table_name,
-            model_shards,
-        ) in self._model_parallel_name_to_local_shards.items():
+        for table_name in self._model_parallel_name_to_local_shards.keys():
             key = f"{prefix}embeddings.{table_name}.weight"
-
+            # gather model shards from both DTensor and ShardedTensor maps
+            model_shards_sharded_tensor = self._model_parallel_name_to_local_shards[
+                table_name
+            ]
+            model_shards_dtensor = self._model_parallel_name_to_shards_wrapper[
+                table_name
+            ]
             # If state_dict[key] is already a ShardedTensor, use its local shards
             if isinstance(state_dict[key], ShardedTensor):
                 local_shards = state_dict[key].local_shards()
-                # If no local shards, create an empty tensor
                 if len(local_shards) == 0:
                     state_dict[key] = torch.empty(0)
                 else:
@@ -624,27 +628,57 @@ def _pre_load_state_dict_hook(
                         ).view(-1, dim)
                     else:
                         state_dict[key] = local_shards[0].tensor.view(-1, dim)
-            else:
+            elif isinstance(state_dict[key], DTensor):
+                shards_wrapper = state_dict[key].to_local()
+                local_shards = shards_wrapper.local_shards()
+                dim = shards_wrapper.local_sizes()[0][1]
+                if len(local_shards) == 0:
+                    state_dict[key] = torch.empty(0)
+                elif len(local_shards) > 1:
+                    # TODO - add multiple shards on rank support
+                    raise RuntimeError(
+                        f"Multiple shards on rank is not supported for DTensor yet, got {len(local_shards)}"
+                    )
+                else:
+                    state_dict[key] = local_shards[0].view(-1, dim)
+            elif isinstance(state_dict[key], torch.Tensor):
                 local_shards = []
-                for shard in model_shards:
-                    # Extract shard size and offsets for splicing
-                    shard_sizes = shard.metadata.shard_sizes
-                    shard_offsets = shard.metadata.shard_offsets
-
-                    # Prepare tensor by splicing and placing on appropriate device
-                    spliced_tensor = state_dict[key][
-                        shard_offsets[0] : shard_offsets[0] + shard_sizes[0],
-                        shard_offsets[1] : shard_offsets[1] + shard_sizes[1],
-                    ].to(shard.tensor.get_device())
-
-                    # Append spliced tensor into local shards
-                    local_shards.append(spliced_tensor)
-
+                if model_shards_sharded_tensor:
+                    # splice according to sharded tensor metadata
+                    for shard in model_shards_sharded_tensor:
+                        # Extract shard size and offsets for splicing
+                        shard_size = shard.metadata.shard_sizes
+                        shard_offset = shard.metadata.shard_offsets
+
+                        # Prepare tensor by splicing and placing on appropriate device
+                        spliced_tensor = state_dict[key][
+                            shard_offset[0] : shard_offset[0] + shard_size[0],
+                            shard_offset[1] : shard_offset[1] + shard_size[1],
+                        ]
+
+                        # Append spliced tensor into local shards
+                        local_shards.append(spliced_tensor)
+                elif model_shards_dtensor:
+                    # splice according to dtensor metadata
+                    for tensor, shard_offset in zip(
+                        model_shards_dtensor["local_tensors"],
+                        model_shards_dtensor["local_offsets"],
+                    ):
+                        shard_size = tensor.size()
+                        spliced_tensor = state_dict[key][
+                            shard_offset[0] : shard_offset[0] + shard_size[0],
+                            shard_offset[1] : shard_offset[1] + shard_size[1],
+                        ]
+                        local_shards.append(spliced_tensor)
                 state_dict[key] = (
                     torch.empty(0)
                     if not local_shards
                     else torch.cat(local_shards, dim=0)
                 )
+            else:
+                raise RuntimeError(
+                    f"Unexpected state_dict key type {type(state_dict[key])} found for {key}"
+                )
 
         for lookup in self._lookups:
             while isinstance(lookup, DistributedDataParallel):
@@ -661,7 +695,9 @@ def _initialize_torch_state(self) -> None:  # noqa
         for table_name in self._table_names:
             self.embeddings[table_name] = nn.Module()
         self._model_parallel_name_to_local_shards = OrderedDict()
+        self._model_parallel_name_to_shards_wrapper = OrderedDict()
         self._model_parallel_name_to_sharded_tensor = OrderedDict()
+        self._model_parallel_name_to_dtensor = OrderedDict()
         model_parallel_name_to_compute_kernel: Dict[str, str] = {}
         for (
             table_name,
@@ -670,6 +706,9 @@ def _initialize_torch_state(self) -> None:  # noqa
             if parameter_sharding.sharding_type == ShardingType.DATA_PARALLEL.value:
                 continue
             self._model_parallel_name_to_local_shards[table_name] = []
+            self._model_parallel_name_to_shards_wrapper[table_name] = OrderedDict(
+                [("local_tensors", []), ("local_offsets", [])]
+            )
             model_parallel_name_to_compute_kernel[table_name] = (
                 parameter_sharding.compute_kernel
             )
@@ -691,18 +730,29 @@ def _initialize_torch_state(self) -> None:  # noqa
                 # save local_shards for transforming MP params to shardedTensor
                 for key, v in lookup.state_dict().items():
                     table_name = key[: -len(".weight")]
-                    self._model_parallel_name_to_local_shards[table_name].extend(
-                        v.local_shards()
-                    )
+                    if isinstance(v, DTensor):
+                        shards_wrapper = self._model_parallel_name_to_shards_wrapper[
+                            table_name
+                        ]
+                        local_shards_wrapper = v._local_tensor
+                        shards_wrapper["local_tensors"].extend(local_shards_wrapper.local_shards())  # pyre-ignore[16]
+                        shards_wrapper["local_offsets"].extend(local_shards_wrapper.local_offsets())  # pyre-ignore[16]
+                        shards_wrapper["global_size"] = v.size()
+                        shards_wrapper["global_stride"] = v.stride()
+                        shards_wrapper["placements"] = v.placements
+                    elif isinstance(v, ShardedTensor):
+                        self._model_parallel_name_to_local_shards[table_name].extend(
+                            v.local_shards()
+                        )
             for (
                 table_name,
                 tbe_slice,
             ) in lookup.named_parameters_by_table():
                 self.embeddings[table_name].register_parameter("weight", tbe_slice)
-        for (
-            table_name,
-            local_shards,
-        ) in self._model_parallel_name_to_local_shards.items():
+        for table_name in self._model_parallel_name_to_local_shards.keys():
+            local_shards = self._model_parallel_name_to_local_shards[table_name]
+            shards_wrapper_map = self._model_parallel_name_to_shards_wrapper[table_name]
+
             # for shards that don't exist on this rank, register with empty tensor
             if not hasattr(self.embeddings[table_name], "weight"):
                 self.embeddings[table_name].register_parameter(
@@ -715,18 +765,34 @@ def _initialize_torch_state(self) -> None:  # noqa
                     self.embeddings[table_name].weight._in_backward_optimizers = [
                         EmptyFusedOptimizer()
                     ]
+
             if model_parallel_name_to_compute_kernel[table_name] in {
                 EmbeddingComputeKernel.KEY_VALUE.value
             }:
                 continue
-            # created ShardedTensors once in init, use in post_state_dict_hook
-            self._model_parallel_name_to_sharded_tensor[table_name] = (
-                ShardedTensor._init_from_local_shards(
-                    local_shards,
-                    self._name_to_table_size[table_name],
-                    process_group=self._env.process_group,
+
+            if shards_wrapper_map["local_tensors"]:
+                self._model_parallel_name_to_dtensor[table_name] = DTensor.from_local(
+                    local_tensor=LocalShardsWrapper(
+                        local_shards=shards_wrapper_map["local_tensors"],
+                        local_offsets=shards_wrapper_map["local_offsets"],
+                    ),
+                    device_mesh=self._env.device_mesh,
+                    placements=shards_wrapper_map["placements"],
+                    shape=shards_wrapper_map["global_size"],
+                    stride=shards_wrapper_map["global_stride"],
+                    run_check=False,
+                )
+            else:
+                # if DTensors for table do not exist, create ShardedTensor
+                # created ShardedTensors once in init, use in post_state_dict_hook
+                self._model_parallel_name_to_sharded_tensor[table_name] = (
+                    ShardedTensor._init_from_local_shards(
+                        local_shards,
+                        self._name_to_table_size[table_name],
+                        process_group=self._env.process_group,
+                    )
                 )
-            )
 
         def post_state_dict_hook(
             module: ShardedEmbeddingCollection,
@@ -741,6 +807,12 @@ def post_state_dict_hook(
             ) in module._model_parallel_name_to_sharded_tensor.items():
                 destination_key = f"{prefix}embeddings.{table_name}.weight"
                 destination[destination_key] = sharded_t
+            for (
+                table_name,
+                d_tensor,
+            ) in module._model_parallel_name_to_dtensor.items():
+                destination_key = f"{prefix}embeddings.{table_name}.weight"
+                destination[destination_key] = d_tensor
 
         self.register_state_dict_pre_hook(self._pre_state_dict_hook)
         self._register_state_dict_hook(post_state_dict_hook)

torchrec/distributed/embedding_kernel.py

@@ -15,11 +15,14 @@
 import torch
 import torch.distributed as dist
 from torch import nn
+from torch.distributed._tensor import DTensor
 from torchrec.distributed.embedding_types import (
+    DTensorMetadata,
     EmbeddingComputeKernel,
     GroupedEmbeddingConfig,
     ShardedEmbeddingTable,
 )
+from torchrec.distributed.shards_wrapper import LocalShardsWrapper
 from torchrec.distributed.types import Shard, ShardedTensor, ShardedTensorMetadata
 from torchrec.sparse.jagged_tensor import KeyedJaggedTensor
 
@@ -73,6 +76,8 @@ def get_state_dict(
     """
     key_to_local_shards: Dict[str, List[Shard]] = defaultdict(list)
     key_to_global_metadata: Dict[str, ShardedTensorMetadata] = {}
+    key_to_dtensor_metadata: Dict[str, DTensorMetadata] = {}
+    key_to_local_tensor_shards: Dict[str, List[Any]] = defaultdict(list)  # pyre-ignore[33]
 
     def get_key_from_embedding_table(embedding_table: ShardedEmbeddingTable) -> str:
         return prefix + f"{embedding_table.name}.weight"
@@ -98,7 +103,16 @@ def get_key_from_embedding_table(embedding_table: ShardedEmbeddingTable) -> str:
         if qscale is not None:
             assert embedding_table.local_cols == param.size(1)  # pyre-ignore[16]
 
-        if embedding_table.global_metadata is not None and pg is not None:
+        if embedding_table.dtensor_metadata is not None and pg is not None:
+            # DTensor path
+            key_to_dtensor_metadata[key] = embedding_table.dtensor_metadata
+            key_to_local_tensor_shards[key].append(
+                [
+                    param,
+                    embedding_table.local_metadata.shard_offsets,  # pyre-ignore[16]
+                ]
+            )
+        elif embedding_table.global_metadata is not None and pg is not None:
             # set additional field of sharded tensor based on local tensor properties
             embedding_table.global_metadata.tensor_properties.dtype = (
                 param.dtype  # pyre-ignore[16]
@@ -133,5 +147,24 @@ def get_key_from_embedding_table(embedding_table: ShardedEmbeddingTable) -> str:
                     process_group=pg,
                 )
             )
-
+        # DTensor path
+        for key in key_to_local_tensor_shards:
+            dtensor_metadata = key_to_dtensor_metadata[key]
+            destination[key] = DTensor.from_local(
+                local_tensor=LocalShardsWrapper(
+                    local_shards=[
+                        tensor_shards[0]
+                        for tensor_shards in key_to_local_tensor_shards[key]
+                    ],
+                    local_offsets=[
+                        tensor_shards[1]
+                        for tensor_shards in key_to_local_tensor_shards[key]
+                    ],
+                ),
+                device_mesh=dtensor_metadata.mesh,
+                placements=dtensor_metadata.placements,
+                shape=torch.Size(dtensor_metadata.size),  # pyre-ignore[6]
+                stride=dtensor_metadata.stride,
+                run_check=False,
+            )
     return destination