enable customized emb lookup kernel for TorchRec

Summary:
# context
* NVIDIA dynamicemb package depends on an old TorchRec release (r0.7) plus a PR ([meta-pytorch#2533](meta-pytorch#2533))
* The goal is to refactor the PR ([meta-pytorch#2533](meta-pytorch#2533)) on trunk so that torchrec can accept customized kernel.

# design rationales
* Given the fact that the [`EmbeddingComputeKernel`](https://github.com/pytorch/torchrec/blob/main/torchrec/distributed/embedding_types.py#L64-L72) is a Enum class which can't be dynamically extended outside of TorchRec codebase, we are adding a placeholder type named `customized_kernel` for all customized compute kernels.
* `compute_kernel` is set in [ParameterSharding](https://github.com/pytorch/torchrec/blob/main/torchrec/distributed/types.py#L694), along with `sharding_type`, `sharding_specs`, etc. User can subclass the `ParameterSharding` dataclass to add more configs and parameters needed by the customized compute kernel, including something like `customized_compute_kernel` to specify the exact one in case there are many. 
* In order to propagate some [extra config](https://fburl.com/code/bnwp44sz) to the customized kernel, we add a `get_additional_fused_params` to propagate the params to `fused_params`. (we might consider to move the [`add_params_from_parameter_sharding`](https://github.com/pytorch/torchrec/blob/main/torchrec/distributed/utils.py#L359) function to a class function of ParameterSharding, so that the user can modify the function when necessary.

NOTE: `fused_params` is originally used for passing necessary parameters to the fbgemm lookup kernels (e.g., TBE, see below). It now seems to be just a convenient way of [propagating configs to the kernel from `ParametersSharding`](https://github.com/pytorch/torchrec/blob/main/torchrec/distributed/utils.py#L359).
```
(Pdb) group_fused_params
{'optimizer': <EmbOptimType.EXACT_ADAGRAD: 'exact_adagrad'>, 'learning_rate': 0.1}
```

* besides the lookup module, very often the customized kernel also needs a customized input_dist and/or a customized output_dist. they all come from [EmbeddingSharding](https://github.com/pytorch/torchrec/blob/main/torchrec/distributed/embedding_sharding.py#L964) and its [child classes](https://github.com/pytorch/torchrec/tree/main/torchrec/distributed/sharding) like cw_sharding, tw_sharding, etc.
* we make it public for the main API [`create_embedding_sharding`](https://github.com/pytorch/torchrec/blob/main/torchrec/distributed/embedding.py#L150) function that return a subclass of EmbeddingSharding, which further creates the user-defined input_dist, output_dist, lookup modules and so on.

WARNING: somehow the HKV-based customized compute kernel can't handle `_initialize_torch_state` likely due to the table.weight tensor is no long on the GPU, so it can't really be represented with sharded tensor or DTensor. It's the user's responsibility to correctly handle the state_dict by overriding the `_initialize_torch_state` function.

Differential Revision: D70723583

Author: TroyGarden

torchrec/distributed/embedding.py

@@ -147,46 +147,6 @@ def get_ec_index_dedup() -> bool:
     return EC_INDEX_DEDUP
 
 
-def create_embedding_sharding(
-    sharding_type: str,
-    sharding_infos: List[EmbeddingShardingInfo],
-    env: ShardingEnv,
-    device: Optional[torch.device] = None,
-    qcomm_codecs_registry: Optional[Dict[str, QuantizedCommCodecs]] = None,
-) -> EmbeddingSharding[
-    SequenceShardingContext, KeyedJaggedTensor, torch.Tensor, torch.Tensor
-]:
-    if sharding_type == ShardingType.TABLE_WISE.value:
-        return TwSequenceEmbeddingSharding(
-            sharding_infos=sharding_infos,
-            env=env,
-            device=device,
-            qcomm_codecs_registry=qcomm_codecs_registry,
-        )
-    elif sharding_type == ShardingType.ROW_WISE.value:
-        return RwSequenceEmbeddingSharding(
-            sharding_infos=sharding_infos,
-            env=env,
-            device=device,
-            qcomm_codecs_registry=qcomm_codecs_registry,
-        )
-    elif sharding_type == ShardingType.DATA_PARALLEL.value:
-        return DpSequenceEmbeddingSharding(
-            sharding_infos=sharding_infos,
-            env=env,
-            device=device,
-        )
-    elif sharding_type == ShardingType.COLUMN_WISE.value:
-        return CwSequenceEmbeddingSharding(
-            sharding_infos=sharding_infos,
-            env=env,
-            device=device,
-            qcomm_codecs_registry=qcomm_codecs_registry,
-        )
-    else:
-        raise ValueError(f"Sharding not supported {sharding_type}")
-
-
 def create_sharding_infos_by_sharding(
     module: EmbeddingCollectionInterface,
     table_name_to_parameter_sharding: Dict[str, ParameterSharding],
@@ -557,7 +517,7 @@ def __init__(
                 SequenceShardingContext, KeyedJaggedTensor, torch.Tensor, torch.Tensor
             ],
         ] = {
-            sharding_type: create_embedding_sharding(
+            sharding_type: self.create_embedding_sharding(
                 sharding_type=sharding_type,
                 sharding_infos=embedding_confings,
                 env=env,
@@ -637,6 +597,51 @@ def __init__(
         if module.device != torch.device("meta"):
             self.load_state_dict(module.state_dict())
 
+    @classmethod
+    def create_embedding_sharding(
+        cls,
+        sharding_type: str,
+        sharding_infos: List[EmbeddingShardingInfo],
+        env: ShardingEnv,
+        device: Optional[torch.device] = None,
+        qcomm_codecs_registry: Optional[Dict[str, QuantizedCommCodecs]] = None,
+    ) -> EmbeddingSharding[
+        SequenceShardingContext, KeyedJaggedTensor, torch.Tensor, torch.Tensor
+    ]:
+        """
+        This is the main function to generate `EmbeddingSharding` instances based on sharding_type
+        so that the same sharding_type in one EC would be fused.
+        """
+        if sharding_type == ShardingType.TABLE_WISE.value:
+            return TwSequenceEmbeddingSharding(
+                sharding_infos=sharding_infos,
+                env=env,
+                device=device,
+                qcomm_codecs_registry=qcomm_codecs_registry,
+            )
+        elif sharding_type == ShardingType.ROW_WISE.value:
+            return RwSequenceEmbeddingSharding(
+                sharding_infos=sharding_infos,
+                env=env,
+                device=device,
+                qcomm_codecs_registry=qcomm_codecs_registry,
+            )
+        elif sharding_type == ShardingType.DATA_PARALLEL.value:
+            return DpSequenceEmbeddingSharding(
+                sharding_infos=sharding_infos,
+                env=env,
+                device=device,
+            )
+        elif sharding_type == ShardingType.COLUMN_WISE.value:
+            return CwSequenceEmbeddingSharding(
+                sharding_infos=sharding_infos,
+                env=env,
+                device=device,
+                qcomm_codecs_registry=qcomm_codecs_registry,
+            )
+        else:
+            raise ValueError(f"Sharding not supported {sharding_type}")
+
     @staticmethod
     def _pre_state_dict_hook(
         self: "ShardedEmbeddingCollection",
@@ -757,14 +762,23 @@ def _initialize_torch_state(self) -> None:  # noqa
             parameter_sharding,
         ) in self.module_sharding_plan.items():
             if parameter_sharding.sharding_type == ShardingType.DATA_PARALLEL.value:
+                # Don't need to use sharded/distributed state tensor for DATA_PARALLEL
+                # because each rank has a full copy of the table in DATA_PARALLEL
+                continue
+            _model_parallel_name_to_compute_kernel[table_name] = (
+                parameter_sharding.compute_kernel
+            )
+            if (
+                parameter_sharding.compute_kernel
+                == EmbeddingComputeKernel.CUSTOMIZED_KERNEL.value
+            ):
+                # Skip state_dict handling for CUSTOMIZED_KERNEL, this should be implemented
+                # in child class for the CUSTOMIZED_KERNEL
                 continue
             self._model_parallel_name_to_local_shards[table_name] = []
             self._model_parallel_name_to_shards_wrapper[table_name] = OrderedDict(
                 [("local_tensors", []), ("local_offsets", [])]
             )
-            _model_parallel_name_to_compute_kernel[table_name] = (
-                parameter_sharding.compute_kernel
-            )
 
         self._name_to_table_size = {}
         for table in self._embedding_configs:
@@ -783,6 +797,11 @@ def _initialize_torch_state(self) -> None:  # noqa
                 # save local_shards for transforming MP params to shardedTensor
                 for key, v in lookup.state_dict().items():
                     table_name = key[: -len(".weight")]
+                    if (
+                        _model_parallel_name_to_compute_kernel[table_name]
+                        == EmbeddingComputeKernel.CUSTOMIZED_KERNEL.value
+                    ):
+                        continue
                     if isinstance(v, DTensor):
                         shards_wrapper = self._model_parallel_name_to_shards_wrapper[
                             table_name

torchrec/distributed/embedding_lookup.py

@@ -181,46 +181,11 @@ def __init__(
         pg: Optional[dist.ProcessGroup] = None,
         device: Optional[torch.device] = None,
     ) -> None:
-        # TODO rename to _create_embedding_kernel
-        def _create_lookup(
-            config: GroupedEmbeddingConfig,
-        ) -> BaseEmbedding:
-            for table in config.embedding_tables:
-                if (
-                    table.compute_kernel == EmbeddingComputeKernel.FUSED_UVM_CACHING
-                    or table.compute_kernel == EmbeddingComputeKernel.KEY_VALUE
-                ):
-                    self._need_prefetch = True
-            if config.compute_kernel == EmbeddingComputeKernel.DENSE:
-                return BatchedDenseEmbedding(
-                    config=config,
-                    pg=pg,
-                    device=device,
-                )
-            elif config.compute_kernel == EmbeddingComputeKernel.FUSED:
-                return BatchedFusedEmbedding(
-                    config=config,
-                    pg=pg,
-                    device=device,
-                )
-            elif config.compute_kernel in {
-                EmbeddingComputeKernel.KEY_VALUE,
-            }:
-                return KeyValueEmbedding(
-                    config=config,
-                    pg=pg,
-                    device=device,
-                )
-            else:
-                raise ValueError(
-                    f"Compute kernel not supported {config.compute_kernel}"
-                )
-
         super().__init__()
         self._emb_modules: nn.ModuleList = nn.ModuleList()
         self._need_prefetch: bool = False
         for config in grouped_configs:
-            self._emb_modules.append(_create_lookup(config))
+            self._emb_modules.append(self._create_embedding_kernel(config, pg, device))
 
         self._feature_splits: List[int] = []
         for config in grouped_configs:
@@ -239,6 +204,41 @@ def _create_lookup(
 
         self.grouped_configs = grouped_configs
 
+    def _create_embedding_kernel(
+        self,
+        config: GroupedEmbeddingConfig,
+        pg: Optional[dist.ProcessGroup],
+        device: Optional[torch.device],
+    ) -> BaseEmbedding:
+        for table in config.embedding_tables:
+            if (
+                table.compute_kernel == EmbeddingComputeKernel.FUSED_UVM_CACHING
+                or table.compute_kernel == EmbeddingComputeKernel.KEY_VALUE
+            ):
+                self._need_prefetch = True
+        if config.compute_kernel == EmbeddingComputeKernel.DENSE:
+            return BatchedDenseEmbedding(
+                config=config,
+                pg=pg,
+                device=device,
+            )
+        elif config.compute_kernel == EmbeddingComputeKernel.FUSED:
+            return BatchedFusedEmbedding(
+                config=config,
+                pg=pg,
+                device=device,
+            )
+        elif config.compute_kernel in {
+            EmbeddingComputeKernel.KEY_VALUE,
+        }:
+            return KeyValueEmbedding(
+                config=config,
+                pg=pg,
+                device=device,
+            )
+        else:
+            raise ValueError(f"Compute kernel not supported {config.compute_kernel}")
+
     def prefetch(
         self,
         sparse_features: KeyedJaggedTensor,
@@ -409,44 +409,12 @@ def __init__(
         scale_weight_gradients: bool = True,
         sharding_type: Optional[ShardingType] = None,
     ) -> None:
-        # TODO rename to _create_embedding_kernel
-        def _create_lookup(
-            config: GroupedEmbeddingConfig,
-            device: Optional[torch.device] = None,
-            sharding_type: Optional[ShardingType] = None,
-        ) -> BaseEmbedding:
-            if config.compute_kernel == EmbeddingComputeKernel.DENSE:
-                return BatchedDenseEmbeddingBag(
-                    config=config,
-                    pg=pg,
-                    device=device,
-                    sharding_type=sharding_type,
-                )
-            elif config.compute_kernel == EmbeddingComputeKernel.FUSED:
-                return BatchedFusedEmbeddingBag(
-                    config=config,
-                    pg=pg,
-                    device=device,
-                    sharding_type=sharding_type,
-                )
-            elif config.compute_kernel in {
-                EmbeddingComputeKernel.KEY_VALUE,
-            }:
-                return KeyValueEmbeddingBag(
-                    config=config,
-                    pg=pg,
-                    device=device,
-                    sharding_type=sharding_type,
-                )
-            else:
-                raise ValueError(
-                    f"Compute kernel not supported {config.compute_kernel}"
-                )
-
         super().__init__()
         self._emb_modules: nn.ModuleList = nn.ModuleList()
         for config in grouped_configs:
-            self._emb_modules.append(_create_lookup(config, device, sharding_type))
+            self._emb_modules.append(
+                self._create_embedding_kernel(config, device, pg, sharding_type)
+            )
 
         self._feature_splits: List[int] = []
         for config in grouped_configs:
@@ -473,6 +441,39 @@ def _create_lookup(
             else 1
         )
 
+    def _create_embedding_kernel(
+        self,
+        config: GroupedEmbeddingConfig,
+        device: Optional[torch.device],
+        pg: Optional[dist.ProcessGroup],
+        sharding_type: Optional[ShardingType],
+    ) -> BaseEmbedding:
+        if config.compute_kernel == EmbeddingComputeKernel.DENSE:
+            return BatchedDenseEmbeddingBag(
+                config=config,
+                pg=pg,
+                device=device,
+                sharding_type=sharding_type,
+            )
+        elif config.compute_kernel == EmbeddingComputeKernel.FUSED:
+            return BatchedFusedEmbeddingBag(
+                config=config,
+                pg=pg,
+                device=device,
+                sharding_type=sharding_type,
+            )
+        elif config.compute_kernel in {
+            EmbeddingComputeKernel.KEY_VALUE,
+        }:
+            return KeyValueEmbeddingBag(
+                config=config,
+                pg=pg,
+                device=device,
+                sharding_type=sharding_type,
+            )
+        else:
+            raise ValueError(f"Compute kernel not supported {config.compute_kernel}")
+
     def prefetch(
         self,
         sparse_features: KeyedJaggedTensor,

torchrec/distributed/embedding_types.py

@@ -70,6 +70,7 @@ class EmbeddingComputeKernel(Enum):
     QUANT_UVM = "quant_uvm"
     QUANT_UVM_CACHING = "quant_uvm_caching"
     KEY_VALUE = "key_value"
+    CUSTOMIZED_KERNEL = "customized_kernel"
 
 
 def compute_kernel_to_embedding_location(