D65489998 Optimize performance of embeddings sharding (pytorch#2615)

Summary: Backing out "D65489998 Optimize performance of embeddings sharding" due to dtype mismatch during Transfer Learning TransferAll Reviewed By: gag1jain, ge0405 Differential Revision: D66800554
gag1jain · Dec 5, 2024 · 8f6dcd9 · 8f6dcd9
1 parent 5b4a5e1
commit 8f6dcd9
Showing 1 changed file with 4 additions and 26 deletions.
diff --git a/torchrec/distributed/embeddingbag.py b/torchrec/distributed/embeddingbag.py
@@ -65,10 +65,8 @@
     QuantizedCommCodecs,
     ShardedTensor,
     ShardingEnv,
-    ShardingEnv2D,
     ShardingType,
     ShardMetadata,
-    TensorProperties,
 )
 from torchrec.distributed.utils import (
     add_params_from_parameter_sharding,
@@ -936,31 +934,11 @@ def _initialize_torch_state(self) -> None:  # noqa
                 # created ShardedTensors once in init, use in post_state_dict_hook
                 # note: at this point kvstore backed tensors don't own valid snapshots, so no read
                 # access is allowed on them.
-                sharding_spec = none_throws(
-                    self.module_sharding_plan[table_name].sharding_spec
-                )
-                metadata = sharding_spec.build_metadata(
-                    tensor_sizes=self._name_to_table_size[table_name],
-                    tensor_properties=(
-                        TensorProperties(
-                            dtype=local_shards[0].tensor.dtype,
-                            layout=local_shards[0].tensor.layout,
-                            requires_grad=local_shards[0].tensor.requires_grad,
-                        )
-                        if local_shards
-                        else TensorProperties()
-                    ),
-                )
-
                 self._model_parallel_name_to_sharded_tensor[table_name] = (
-                    ShardedTensor._init_from_local_shards_and_global_metadata(
-                        local_shards=local_shards,
-                        sharded_tensor_metadata=metadata,
-                        process_group=(
-                            self._env.sharding_pg
-                            if isinstance(self._env, ShardingEnv2D)
-                            else self._env.process_group
-                        ),
+                    ShardedTensor._init_from_local_shards(
+                        local_shards,
+                        self._name_to_table_size[table_name],
+                        process_group=self._env.process_group,
                     )
                 )