Cleanup rollout code for sharding optimization (#2921)

Summary:
Pull Request resolved: #2921

As per title, the optimization has been rolled out to production jobs for several months so it it is time to delete the rollout code and old implementation.

Reviewed By: ilyas409

Differential Revision: D73693963

fbshipit-source-id: f31db8ada9eafb6c2df346ce3e60fc9798298f44

Author: Boris Sarana

torchrec/distributed/embedding_types.py

@@ -31,9 +31,7 @@
 from torch.distributed._tensor.placement_types import Placement
 from torch.nn.modules.module import _addindent
 from torch.nn.parallel import DistributedDataParallel
-from torchrec.distributed.global_settings import (
-    construct_sharded_tensor_from_metadata_enabled,
-)
+
 from torchrec.distributed.types import (
     get_tensor_size_bytes,
     ModuleSharder,
@@ -358,11 +356,6 @@ def __init__(
         self._lookups: List[nn.Module] = []
         self._output_dists: List[nn.Module] = []
 
-        # option to construct ShardedTensor from metadata avoiding expensive all-gather
-        self._construct_sharded_tensor_from_metadata: bool = (
-            construct_sharded_tensor_from_metadata_enabled()
-        )
-
     def prefetch(
         self,
         dist_input: KJTList,

torchrec/distributed/embeddingbag.py

@@ -1006,46 +1006,32 @@ def _initialize_torch_state(self, skip_registering: bool = False) -> None:  # no
                 # access is allowed on them.
 
                 # create ShardedTensor from local shards and metadata avoding all_gather collective
-                if self._construct_sharded_tensor_from_metadata:
-                    sharding_spec = none_throws(
-                        self.module_sharding_plan[table_name].sharding_spec
-                    )
-
-                    tensor_properties = TensorProperties(
-                        dtype=(
-                            data_type_to_dtype(
-                                self._table_name_to_config[table_name].data_type
-                            )
-                        ),
-                    )
+                sharding_spec = none_throws(
+                    self.module_sharding_plan[table_name].sharding_spec
+                )
 
-                    self._model_parallel_name_to_sharded_tensor[table_name] = (
-                        ShardedTensor._init_from_local_shards_and_global_metadata(
-                            local_shards=local_shards,
-                            sharded_tensor_metadata=sharding_spec.build_metadata(
-                                tensor_sizes=self._name_to_table_size[table_name],
-                                tensor_properties=tensor_properties,
-                            ),
-                            process_group=(
-                                self._env.sharding_pg
-                                if isinstance(self._env, ShardingEnv2D)
-                                else self._env.process_group
-                            ),
-                        )
-                    )
-                else:
-                    # create ShardedTensor from local shards using all_gather collective
-                    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.sharding_pg
-                                if isinstance(self._env, ShardingEnv2D)
-                                else self._env.process_group
-                            ),
+                tensor_properties = TensorProperties(
+                    dtype=(
+                        data_type_to_dtype(
+                            self._table_name_to_config[table_name].data_type
                         )
+                    ),
+                )
+
+                self._model_parallel_name_to_sharded_tensor[table_name] = (
+                    ShardedTensor._init_from_local_shards_and_global_metadata(
+                        local_shards=local_shards,
+                        sharded_tensor_metadata=sharding_spec.build_metadata(
+                            tensor_sizes=self._name_to_table_size[table_name],
+                            tensor_properties=tensor_properties,
+                        ),
+                        process_group=(
+                            self._env.sharding_pg
+                            if isinstance(self._env, ShardingEnv2D)
+                            else self._env.process_group
+                        ),
                     )
+                )
 
         def extract_sharded_kvtensors(
             module: ShardedEmbeddingBagCollection,

torchrec/distributed/global_settings.py

@@ -7,14 +7,8 @@
 
 # pyre-strict
 
-import os
-
 PROPOGATE_DEVICE: bool = False
 
-TORCHREC_CONSTRUCT_SHARDED_TENSOR_FROM_METADATA_ENV = (
-    "TORCHREC_CONSTRUCT_SHARDED_TENSOR_FROM_METADATA"
-)
-
 
 def set_propogate_device(val: bool) -> None:
     global PROPOGATE_DEVICE
@@ -24,13 +18,3 @@ def set_propogate_device(val: bool) -> None:
 def get_propogate_device() -> bool:
     global PROPOGATE_DEVICE
     return PROPOGATE_DEVICE
-
-
-def construct_sharded_tensor_from_metadata_enabled() -> bool:
-    return (
-        os.environ.get(TORCHREC_CONSTRUCT_SHARDED_TENSOR_FROM_METADATA_ENV, "0") == "1"
-    )
-
-
-def enable_construct_sharded_tensor_from_metadata() -> None:
-    os.environ[TORCHREC_CONSTRUCT_SHARDED_TENSOR_FROM_METADATA_ENV] = "1"