revert sharding grouping logic for vbe

Summary:
reverting sharding grouping logic in EBC/VLE modules that supported specific UVM caching + prefetch pipeline uses cases to circumvent VBE TBE output concatenation.

As concatenation is implemented in the preceding diff, this diff cleans up the logic left behind from grouping sharding by UVM caching kernel conditions to avoid VBE TBE output concatenation.

Differential Revision: D58989195

Author: joshuadeng

torchrec/distributed/embeddingbag.py

@@ -13,7 +13,6 @@
 from functools import partial
 from typing import (
     Any,
-    Callable,
     cast,
     Dict,
     Iterator,
@@ -39,7 +38,6 @@
     EmbeddingShardingInfo,
     KJTListSplitsAwaitable,
     Multistreamable,
-    USE_ONE_TBE_PER_TABLE,
 )
 from torchrec.distributed.embedding_types import (
     BaseEmbeddingSharder,
@@ -77,7 +75,6 @@
     optimizer_type_to_emb_opt_type,
 )
 from torchrec.modules.embedding_configs import (
-    BaseEmbeddingConfig,
     EmbeddingBagConfig,
     EmbeddingTableConfig,
     PoolingType,
@@ -200,15 +197,7 @@ 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:
-    return param_sharding.sharding_type
-
-
-def create_sharding_infos_by_group(
+def create_sharding_infos_by_sharding(
     module: EmbeddingBagCollectionInterface,
     table_name_to_parameter_sharding: Dict[str, ParameterSharding],
     prefix: str,
@@ -229,7 +218,9 @@ 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]] = (
+        defaultdict(list)
+    )
 
     # state_dict returns parameter.Tensor, which loses parameter level attributes
     parameter_by_name = dict(module.named_parameters())
@@ -283,7 +274,6 @@ def create_sharding_infos_by_group(
         )
         per_table_fused_params = convert_to_fbgemm_types(per_table_fused_params)
 
-        group = get_sharding_group(config, parameter_sharding, fused_params)
         sharding_info = EmbeddingShardingInfo(
             embedding_config=EmbeddingTableConfig(
                 num_embeddings=config.num_embeddings,
@@ -303,8 +293,10 @@ def create_sharding_infos_by_group(
             param=param,
             fused_params=per_table_fused_params,
         )
-        group_to_sharding_infos[group].append(sharding_info)
-    return group_to_sharding_infos
+        sharding_type_to_sharding_infos[parameter_sharding.sharding_type].append(
+            sharding_info
+        )
+    return sharding_type_to_sharding_infos
 
 
 def create_sharding_infos_by_sharding_device_group(
@@ -581,7 +573,7 @@ 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.",
@@ -602,7 +594,7 @@ def __init__(
                 permute_embeddings=True,
                 qcomm_codecs_registry=self.qcomm_codecs_registry,
             )
-            for embedding_configs in group_to_sharding_infos.values()
+            for embedding_configs in sharding_type_to_sharding_infos.values()
         ]
 
         self._is_weighted: bool = module.is_weighted()

torchrec/distributed/planner/tests/test_embeddingbag_utils.py

@@ -11,7 +11,7 @@
 import unittest
 
 from torchrec.distributed.embeddingbag import (
-    create_sharding_infos_by_group,
+    create_sharding_infos_by_sharding,
     EmbeddingBagCollectionSharder,
 )
 from torchrec.distributed.planner import (
@@ -79,7 +79,7 @@ def setUp(self) -> None:
         )
         self.expected_plan = planner.plan(self.model, [self.sharder])  # pyre-ignore[6]
 
-        self.expected_sharding_infos = create_sharding_infos_by_group(
+        self.expected_sharding_infos = create_sharding_infos_by_sharding(
             self.model,
             self.expected_plan.get_plan_for_module(""),  # pyre-ignore[6]
             prefix="embedding_bags.",
@@ -93,7 +93,7 @@ def test_create_sharding_infos_by_group_override(self) -> None:
 
         # with sharder fused params that will get overridden
         sharder_fused_params = {"enforce_hbm": False}
-        overriden_sharding_infos = create_sharding_infos_by_group(
+        overriden_sharding_infos = create_sharding_infos_by_sharding(
             self.model,
             self.expected_plan.get_plan_for_module(""),
             prefix="embedding_bags.",
@@ -106,7 +106,7 @@ def test_create_sharding_infos_by_group_override(self) -> None:
 
         # with sharder fused params that won't get overridden
         sharder_fused_params = {"ABC": True}
-        not_overriden_sharding_infos = create_sharding_infos_by_group(
+        not_overriden_sharding_infos = create_sharding_infos_by_sharding(
             self.model,
             self.expected_plan.get_plan_for_module(""),
             prefix="embedding_bags.",
@@ -141,7 +141,7 @@ def test_create_sharding_infos_by_group_combine(self) -> None:
         # provide that two fused params from sharder
         sharder_fused_params = {"enforce_hbm": True, "stochastic_rounding": False}
 
-        combined_sharding_infos = create_sharding_infos_by_group(
+        combined_sharding_infos = create_sharding_infos_by_sharding(
             self.model,
             new_plan.get_plan_for_module(""),  # pyre-ignore[6]
             prefix="embedding_bags.",
@@ -156,7 +156,7 @@ def test_create_sharding_infos_by_group_combine(self) -> None:
 
         # provide that two fused params from sharder wrongly
         sharder_fused_params = {"enforce_hbm": True, "stochastic_rounding": True}
-        wrong_combined_sharding_infos = create_sharding_infos_by_group(
+        wrong_combined_sharding_infos = create_sharding_infos_by_sharding(
             self.model,
             new_plan.get_plan_for_module(""),  # pyre-ignore[6]
             prefix="embedding_bags.",