Revert D81366596: Add configs for write dist

Differential Revision:
D81366596

Original commit changeset: 2cc120644d31

Original Phabricator Diff: D81366596

fbshipit-source-id: 9ca96c29db0efa9da04335cb4e2c021e2c445600

Author: Pooja Agarwal

torchrec/distributed/embedding_sharding.py

@@ -565,7 +565,6 @@ def _group_tables_per_rank(
                 ),
                 _prefetch_and_cached(table),
                 table.use_virtual_table if is_inference else None,
-                table.enable_embedding_update,
             )
             # micromanage the order of we traverse the groups to ensure backwards compatibility
             if grouping_key not in groups:
@@ -582,7 +581,6 @@ def _group_tables_per_rank(
                 _,
                 _,
                 use_virtual_table,
-                enable_embedding_update,
             ) = grouping_key
             grouped_tables = groups[grouping_key]
             # remove non-native fused params
@@ -604,7 +602,6 @@ def _group_tables_per_rank(
                     compute_kernel=compute_kernel_type,
                     embedding_tables=grouped_tables,
                     fused_params=per_tbe_fused_params,
-                    enable_embedding_update=enable_embedding_update,
                 )
             )
         return grouped_embedding_configs

torchrec/distributed/embedding_types.py

@@ -251,7 +251,6 @@ class GroupedEmbeddingConfig:
     compute_kernel: EmbeddingComputeKernel
     embedding_tables: List[ShardedEmbeddingTable]
     fused_params: Optional[Dict[str, Any]] = None
-    enable_embedding_update: bool = False
 
     def feature_hash_sizes(self) -> List[int]:
         feature_hash_sizes = []

torchrec/distributed/sharding/rw_sharding.py

@@ -223,7 +223,6 @@ def _shard(
                         total_num_buckets=info.embedding_config.total_num_buckets,
                         use_virtual_table=info.embedding_config.use_virtual_table,
                         virtual_table_eviction_policy=info.embedding_config.virtual_table_eviction_policy,
-                        enable_embedding_update=info.embedding_config.enable_embedding_update,
                     )
                 )
         return tables_per_rank
@@ -279,20 +278,6 @@ def _get_feature_hash_sizes(self) -> List[int]:
             feature_hash_sizes.extend(group_config.feature_hash_sizes())
         return feature_hash_sizes
 
-    def _get_num_writable_features(self) -> int:
-        return sum(
-            group_config.num_features()
-            for group_config in self._grouped_embedding_configs
-            if group_config.enable_embedding_update
-        )
-
-    def _get_writable_feature_hash_sizes(self) -> List[int]:
-        feature_hash_sizes: List[int] = []
-        for group_config in self._grouped_embedding_configs:
-            if group_config.enable_embedding_update:
-                feature_hash_sizes.extend(group_config.feature_hash_sizes())
-        return feature_hash_sizes
-
 
 class RwSparseFeaturesDist(BaseSparseFeaturesDist[KeyedJaggedTensor]):
     """

torchrec/modules/embedding_configs.py

@@ -370,7 +370,6 @@ class BaseEmbeddingConfig:
     total_num_buckets: Optional[int] = None
     use_virtual_table: bool = False
     virtual_table_eviction_policy: Optional[VirtualTableEvictionPolicy] = None
-    enable_embedding_update: bool = False
 
     def get_weight_init_max(self) -> float:
         if self.weight_init_max is None:

torchrec/schema/api_tests/test_embedding_config_schema.py

@@ -43,7 +43,6 @@ class StableEmbeddingBagConfig:
     total_num_buckets: Optional[int] = None
     use_virtual_table: bool = False
     virtual_table_eviction_policy: Optional[VirtualTableEvictionPolicy] = None
-    enable_embedding_update: bool = False
     pooling: PoolingType = PoolingType.SUM