Back out "make EBC group on sharding type and compute kernel" (#2113)

Summary:
Pull Request resolved: #2113

Original commit changeset: 920673d619e9

Original Phabricator Diff: D58160613

Jobs on MVAI failing: https://www.internalfb.com/mlhub/pipelines/runs/mast/fire-dstaay-cfr_roo_testing_v3_pp_dedic_strm?job_attempt=0&version=0&env=PRODUCTION&referrer=MAST_JOB_NOTIFICATION_BOT

backout job is fine:
fire-dstaay-cfr_roo_testing_v3_pp

Reviewed By: sarckk

Differential Revision: D58571298

fbshipit-source-id: 77078a9401803a26258a506d25bb9d4d7334dae4

Author: dstaay-fb

torchrec/distributed/embeddingbag.py

@@ -13,7 +13,6 @@
 from functools import partial
 from typing import (
     Any,
-    Callable,
     cast,
     Dict,
     Iterator,
@@ -38,7 +37,6 @@
     EmbeddingShardingInfo,
     KJTListSplitsAwaitable,
     Multistreamable,
-    USE_ONE_TBE_PER_TABLE,
 )
 from torchrec.distributed.embedding_types import (
     BaseEmbeddingSharder,
@@ -75,7 +73,6 @@
     optimizer_type_to_emb_opt_type,
 )
 from torchrec.modules.embedding_configs import (
-    BaseEmbeddingConfig,
     EmbeddingBagConfig,
     EmbeddingTableConfig,
     PoolingType,
@@ -144,6 +141,7 @@ def replace_placement_with_meta_device(
 
 
 def create_embedding_bag_sharding(
+    sharding_type: str,
     sharding_infos: List[EmbeddingShardingInfo],
     env: ShardingEnv,
     device: Optional[torch.device] = None,
@@ -152,7 +150,6 @@ def create_embedding_bag_sharding(
 ) -> EmbeddingSharding[
     EmbeddingShardingContext, KeyedJaggedTensor, torch.Tensor, torch.Tensor
 ]:
-    sharding_type = sharding_infos[0].param_sharding.sharding_type
     if device is not None and device.type == "meta":
         replace_placement_with_meta_device(sharding_infos)
     if sharding_type == ShardingType.TABLE_WISE.value:
@@ -198,48 +195,12 @@ def create_embedding_bag_sharding(
         raise ValueError(f"Sharding type not supported {sharding_type}")
 
 
-def get_sharding_group(
-    config: BaseEmbeddingConfig,
-    param_sharding: ParameterSharding,
-    fused_params: Optional[Dict[str, Any]] = None,
-) -> str:
-    if fused_params and fused_params.get(USE_ONE_TBE_PER_TABLE, False):
-        return config.name
-    if param_sharding.sharding_type in {
-        ShardingType.COLUMN_WISE.value,
-        ShardingType.TABLE_COLUMN_WISE.value,
-    }:
-        assert param_sharding.ranks
-        num_ranks = len(param_sharding.ranks)
-        assert config.embedding_dim % num_ranks == 0
-        dim = config.embedding_dim // num_ranks
-    else:
-        dim = config.embedding_dim
-
-    group = f"{param_sharding.sharding_type}@{param_sharding.compute_kernel}"
-    if (
-        param_sharding.compute_kernel == EmbeddingComputeKernel.FUSED_UVM_CACHING.value
-        and (
-            (fused_params and fused_params.get("prefetch_pipeline", False))
-            or (
-                param_sharding.cache_params
-                and param_sharding.cache_params.prefetch_pipeline
-            )
-        )
-    ):
-        group += f"@{dim}"
-    return group
-
-
-def create_sharding_infos_by_group(
+def create_sharding_infos_by_sharding(
     module: EmbeddingBagCollectionInterface,
     table_name_to_parameter_sharding: Dict[str, ParameterSharding],
     prefix: str,
     fused_params: Optional[Dict[str, Any]],
     suffix: Optional[str] = "weight",
-    group_fn: Optional[
-        Callable[[EmbeddingBagConfig, ParameterSharding, Optional[Dict[str, Any]]], str]
-    ] = None,
 ) -> Dict[str, List[EmbeddingShardingInfo]]:
 
     if fused_params is None:
@@ -255,7 +216,7 @@ def create_sharding_infos_by_group(
             else:
                 shared_feature[feature_name] = True
 
-    group_to_sharding_infos: Dict[str, List[EmbeddingShardingInfo]] = defaultdict(list)
+    sharding_type_to_sharding_infos: Dict[str, List[EmbeddingShardingInfo]] = {}
 
     # state_dict returns parameter.Tensor, which loses parameter level attributes
     parameter_by_name = dict(module.named_parameters())
@@ -288,6 +249,9 @@ def create_sharding_infos_by_group(
         assert param_name in parameter_by_name or param_name in state_dict
         param = parameter_by_name.get(param_name, state_dict[param_name])
 
+        if parameter_sharding.sharding_type not in sharding_type_to_sharding_infos:
+            sharding_type_to_sharding_infos[parameter_sharding.sharding_type] = []
+
         optimizer_params = getattr(param, "_optimizer_kwargs", [{}])
         optimizer_classes = getattr(param, "_optimizer_classes", [None])
 
@@ -309,32 +273,28 @@ def create_sharding_infos_by_group(
         )
         per_table_fused_params = convert_to_fbgemm_types(per_table_fused_params)
 
-        group = (
-            group_fn(config, parameter_sharding, fused_params)
-            if group_fn is not None
-            else parameter_sharding.sharding_type
-        )
-        sharding_info = EmbeddingShardingInfo(
-            embedding_config=EmbeddingTableConfig(
-                num_embeddings=config.num_embeddings,
-                embedding_dim=config.embedding_dim,
-                name=config.name,
-                data_type=config.data_type,
-                feature_names=copy.deepcopy(config.feature_names),
-                pooling=config.pooling,
-                is_weighted=module.is_weighted(),
-                has_feature_processor=False,
-                embedding_names=embedding_names,
-                weight_init_max=config.weight_init_max,
-                weight_init_min=config.weight_init_min,
-                pruning_indices_remapping=config.pruning_indices_remapping,
-            ),
-            param_sharding=parameter_sharding,
-            param=param,
-            fused_params=per_table_fused_params,
+        sharding_type_to_sharding_infos[parameter_sharding.sharding_type].append(
+            EmbeddingShardingInfo(
+                embedding_config=EmbeddingTableConfig(
+                    num_embeddings=config.num_embeddings,
+                    embedding_dim=config.embedding_dim,
+                    name=config.name,
+                    data_type=config.data_type,
+                    feature_names=copy.deepcopy(config.feature_names),
+                    pooling=config.pooling,
+                    is_weighted=module.is_weighted(),
+                    has_feature_processor=False,
+                    embedding_names=embedding_names,
+                    weight_init_max=config.weight_init_max,
+                    weight_init_min=config.weight_init_min,
+                    pruning_indices_remapping=config.pruning_indices_remapping,
+                ),
+                param_sharding=parameter_sharding,
+                param=param,
+                fused_params=per_table_fused_params,
+            )
         )
-        group_to_sharding_infos[group].append(sharding_info)
-    return group_to_sharding_infos
+    return sharding_type_to_sharding_infos
 
 
 def create_sharding_infos_by_sharding_device_group(
@@ -611,30 +571,31 @@ def __init__(
         )
         self._env = env
 
-        group_to_sharding_infos = create_sharding_infos_by_group(
+        sharding_type_to_sharding_infos = create_sharding_infos_by_sharding(
             module,
             table_name_to_parameter_sharding,
             "embedding_bags.",
             fused_params,
-            group_fn=get_sharding_group,
         )
-        self._embedding_shardings: List[
+        self._sharding_type_to_sharding: Dict[
+            str,
             EmbeddingSharding[
                 EmbeddingShardingContext,
                 KeyedJaggedTensor,
                 torch.Tensor,
                 torch.Tensor,
-            ]
-        ] = [
-            create_embedding_bag_sharding(
+            ],
+        ] = {
+            sharding_type: create_embedding_bag_sharding(
+                sharding_type,
                 embedding_configs,
                 env,
                 device,
                 permute_embeddings=True,
                 qcomm_codecs_registry=self.qcomm_codecs_registry,
             )
-            for embedding_configs in group_to_sharding_infos.values()
-        ]
+            for sharding_type, embedding_configs in sharding_type_to_sharding_infos.items()
+        }
 
         self._is_weighted: bool = module.is_weighted()
         self._device = device
@@ -679,12 +640,15 @@ def __init__(
                     optims.append(("", tbe_module.fused_optimizer))
         self._optim: CombinedOptimizer = CombinedOptimizer(optims)
 
-        for i, (sharding, lookup) in enumerate(
-            zip(self._embedding_shardings, self._lookups)
+        for index, (sharding, lookup) in enumerate(
+            zip(
+                self._sharding_type_to_sharding.values(),
+                self._lookups,
+            )
         ):
             # TODO: can move this into DpPooledEmbeddingSharding once all modules are composable
             if isinstance(sharding, DpPooledEmbeddingSharding):
-                self._lookups[i] = DistributedDataParallel(
+                self._lookups[index] = DistributedDataParallel(
                     module=lookup,
                     device_ids=(
                         [device]
@@ -806,8 +770,10 @@ def _initialize_torch_state(self) -> None:  # noqa
                 table.embedding_dim,
             )
 
-        for lookup, sharding in zip(self._lookups, self._embedding_shardings):
-            if isinstance(sharding, DpPooledEmbeddingSharding):
+        for sharding_type, lookup in zip(
+            self._sharding_type_to_sharding.keys(), self._lookups
+        ):
+            if sharding_type == ShardingType.DATA_PARALLEL.value:
                 # unwrap DDP
                 lookup = lookup.module
             else:
@@ -898,7 +864,7 @@ def _create_input_dist(
         input_feature_names: List[str],
     ) -> None:
         feature_names: List[str] = []
-        for sharding in self._embedding_shardings:
+        for sharding in self._sharding_type_to_sharding.values():
             self._input_dists.append(sharding.create_input_dist())
             feature_names.extend(sharding.feature_names())
             self._feature_splits.append(len(sharding.feature_names()))
@@ -924,7 +890,7 @@ def _init_mean_pooling_callback(
         # account for shared features
         feature_names: List[str] = [
             feature_name
-            for sharding in self._embedding_shardings
+            for sharding in self._sharding_type_to_sharding.values()
             for feature_name in sharding.feature_names()
         ]
 
@@ -951,12 +917,12 @@ def _init_mean_pooling_callback(
     def _create_lookups(
         self,
     ) -> None:
-        for sharding in self._embedding_shardings:
+        for sharding in self._sharding_type_to_sharding.values():
             self._lookups.append(sharding.create_lookup())
 
     def _create_output_dist(self) -> None:
         embedding_shard_metadata: List[Optional[ShardMetadata]] = []
-        for sharding in self._embedding_shardings:
+        for sharding in self._sharding_type_to_sharding.values():
             self._output_dists.append(sharding.create_output_dist(device=self._device))
             self._embedding_names.extend(sharding.embedding_names())
             self._embedding_dims.extend(sharding.embedding_dims())
@@ -1270,6 +1236,7 @@ def __init__(
         self._embedding_sharding: EmbeddingSharding[
             EmbeddingShardingContext, KeyedJaggedTensor, torch.Tensor, torch.Tensor
         ] = create_embedding_bag_sharding(
+            sharding_type=self.parameter_sharding.sharding_type,
             sharding_infos=[
                 EmbeddingShardingInfo(
                     embedding_config=embedding_table_config,

torchrec/distributed/fused_embeddingbag.py

@@ -65,7 +65,7 @@ def __init__(
         )
 
         for index, (sharding, lookup) in enumerate(
-            zip(self._embedding_shardings, self._lookups)
+            zip(self._sharding_type_to_sharding.values(), self._lookups)
         ):
             if isinstance(sharding, DpPooledEmbeddingSharding):
                 self._lookups[index] = DistributedDataParallel(

torchrec/distributed/mc_embedding_modules.py

@@ -109,18 +109,14 @@ def __init__(
         # TODO: This is a hack since _embedding_module doesn't need input
         # dist, so eliminating it so all fused a2a will ignore it.
         self._embedding_module._has_uninitialized_input_dist = False
-        embedding_shardings = (
-            self._embedding_module._embedding_shardings
-            if isinstance(self._embedding_module, ShardedEmbeddingBagCollection)
-            else list(self._embedding_module._sharding_type_to_sharding.values())
-        )
         self._managed_collision_collection: ShardedManagedCollisionCollection = (
             mc_sharder.shard(
                 module._managed_collision_collection,
                 table_name_to_parameter_sharding,
                 env=env,
                 device=device,
-                embedding_shardings=embedding_shardings,
+                # pyre-ignore
+                sharding_type_to_sharding=self._embedding_module._sharding_type_to_sharding,
             )
         )
         self._return_remapped_features: bool = module._return_remapped_features