add emb cast in KTRegroupAsDict module (#3008)

Summary:
Pull Request resolved: #3008

# context
* add `emb_dtype` to the `KTRegroupAsDict` module because in APS some customized regroup module there is a casting operation
* to make the model ir-compatible with the "short-circuit" solution, we'll need to absorb this casting function inside the KTRegroupAsDict module

Reviewed By: malaybag

Differential Revision: D75326034

fbshipit-source-id: 1a32e7c1195b062d19f2b9107ac3af190ebdeb89

Author: TroyGarden

torchrec/ir/schema.py

@@ -54,3 +54,4 @@ class PositionWeightedModuleCollectionMetadata:
 class KTRegroupAsDictMetadata:
     groups: List[List[str]]
     keys: List[str]
+    emb_dtype: Optional[str]

torchrec/ir/serializer.py

@@ -411,6 +411,11 @@ def serialize_to_dict(
             # pyre-fixme[6]: For 2nd argument expected `List[List[str]]` but got
             #  `Union[Module, Tensor]`.
             groups=module._groups,
+            emb_dtype=(
+                module._emb_dtype.value  # pyre-ignore[16]
+                if module._emb_dtype is not None
+                else None
+            ),
         )
         return metadata.__dict__
 
@@ -425,6 +430,9 @@ def deserialize_from_dict(
         return KTRegroupAsDict(
             keys=metadata.keys,
             groups=metadata.groups,
+            emb_dtype=(
+                DataType(metadata.emb_dtype) if metadata.emb_dtype is not None else None
+            ),
         )
 
 

torchrec/ir/tests/test_serializer.py

@@ -22,7 +22,6 @@
     encapsulate_ir_modules,
     mark_dynamic_kjt,
 )
-
 from torchrec.modules.embedding_configs import EmbeddingBagConfig
 from torchrec.modules.embedding_modules import EmbeddingBagCollection
 from torchrec.modules.feature_processor_ import (
@@ -32,6 +31,7 @@
 from torchrec.modules.fp_embedding_modules import FeatureProcessedEmbeddingBagCollection
 from torchrec.modules.regroup import KTRegroupAsDict
 from torchrec.sparse.jagged_tensor import KeyedJaggedTensor, KeyedTensor
+from torchrec.types import DataType
 
 
 class CompoundModule(nn.Module):
@@ -747,3 +747,80 @@ def forward(
         deserialized_out = deserialized_model(id_list_features)
         for key in eager_out.keys():
             torch.testing.assert_close(deserialized_out[key], eager_out[key])
+
+    def test_cast_in_regroup(self) -> None:
+        class Model(nn.Module):
+            def __init__(self, ebc, fpebc, regroup):
+                super().__init__()
+                self.ebc = ebc
+                self.fpebc = fpebc
+                self.regroup = regroup
+
+            def forward(
+                self,
+                features: KeyedJaggedTensor,
+            ) -> Dict[str, torch.Tensor]:
+                kt1 = self.ebc(features)
+                kt2 = self.fpebc(features)
+                return self.regroup([kt1, kt2])
+
+        tb1_config = EmbeddingBagConfig(
+            name="t1",
+            embedding_dim=3,
+            num_embeddings=10,
+            feature_names=["f1", "f2"],
+        )
+        tb2_config = EmbeddingBagConfig(
+            name="t2",
+            embedding_dim=4,
+            num_embeddings=10,
+            feature_names=["f3", "f4"],
+        )
+        tb3_config = EmbeddingBagConfig(
+            name="t3",
+            embedding_dim=5,
+            num_embeddings=10,
+            feature_names=["f5"],
+        )
+
+        ebc = EmbeddingBagCollection(
+            tables=[tb1_config, tb3_config],
+            is_weighted=False,
+        )
+        max_feature_lengths = {"f3": 100, "f4": 100}
+        fpebc = FeatureProcessedEmbeddingBagCollection(
+            EmbeddingBagCollection(
+                tables=[tb2_config],
+                is_weighted=True,
+            ),
+            PositionWeightedModuleCollection(
+                max_feature_lengths=max_feature_lengths,
+            ),
+        )
+        data_type = DataType.BF16
+
+        regroup = KTRegroupAsDict(
+            [["f1", "f3", "f5"], ["f2", "f4"]], ["odd", "even"], emb_dtype=data_type
+        )
+        model = Model(ebc, fpebc, regroup)
+        self.assertEqual(model.regroup._emb_dtype, data_type)
+
+        id_list_features = KeyedJaggedTensor.from_offsets_sync(
+            keys=["f1", "f2", "f3", "f4", "f5"],
+            values=torch.tensor([0, 1, 2, 3, 2, 3, 4, 5, 6, 7, 8, 9, 1, 1, 2]),
+            offsets=torch.tensor([0, 2, 2, 3, 4, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 15]),
+        )
+        # Serialize EBC
+        model, sparse_fqns = encapsulate_ir_modules(model, JsonSerializer)
+        ep = torch.export.export(
+            model,
+            (id_list_features,),
+            {},
+            strict=False,
+            # Allows KJT to not be unflattened and run a forward on unflattened EP
+            preserve_module_call_signature=(tuple(sparse_fqns)),
+        )
+
+        unflatten_ep = torch.export.unflatten(ep)
+        deserialized = decapsulate_ir_modules(unflatten_ep, JsonSerializer)
+        self.assertEqual(deserialized.regroup._emb_dtype, data_type)  # pyre-ignore[16]

torchrec/modules/regroup.py

@@ -12,12 +12,13 @@
 from typing import Dict, List, Optional, Tuple, Union
 
 import torch
+from torchrec.modules.embedding_configs import data_type_to_dtype
 from torchrec.sparse.jagged_tensor import (
     _desugar_keyed_tensors,
     _kt_regroup_arguments,
     KeyedTensor,
 )
-from torchrec.types import CacheMixin
+from torchrec.types import CacheMixin, DataType
 
 
 @torch.fx.wrap
@@ -131,7 +132,12 @@ class KTRegroupAsDict(torch.nn.Module, CacheMixin):
 
     """
 
-    def __init__(self, groups: List[List[str]], keys: List[str]) -> None:
+    def __init__(
+        self,
+        groups: List[List[str]],
+        keys: List[str],
+        emb_dtype: Optional[DataType] = None,
+    ) -> None:
         super().__init__()
         torch._C._log_api_usage_once(f"torchrec.modules.{self.__class__.__name__}")
         assert len(groups) == len(keys), "Groups and keys should have same length"
@@ -145,6 +151,7 @@ def __init__(self, groups: List[List[str]], keys: List[str]) -> None:
         self._splits: List[int] = []
         self._idx_key_pairs: List[Tuple[int, str]] = []
         self._permute_pooled_embs_impl = PermuteMultiEmbedding(groups)
+        self._emb_dtype = emb_dtype
 
     def _init_fbgemm_regroup(self, kts: List[KeyedTensor]) -> None:
         self._use_fbgemm_regroup = True
@@ -190,18 +197,26 @@ def _init_regroup(self, kts: List[KeyedTensor]) -> None:
         self._splits = splits
         self._idx_key_pairs = idx_key_pairs
 
+    def embedding_cast(self, embs: List[torch.Tensor]) -> List[torch.Tensor]:
+        if self._emb_dtype is None:
+            return embs
+        dtype = data_type_to_dtype(self._emb_dtype)
+        return [emb.to(dtype=dtype) for emb in embs]
+
     def forward(self, keyed_tensors: List[KeyedTensor]) -> Dict[str, torch.Tensor]:
         if not self._is_inited:
             module_init(self, keyed_tensors)
 
         if self._use_fbgemm_regroup:
             values = _get_kts_values(keyed_tensors)
+            values = self.embedding_cast(values)
             permuted_values = self._permute_pooled_embs_impl(values)
             return _to_tensor_dict(self._keys, permuted_values)
         else:
             permuted_values = _permuted_values(
                 keyed_tensors, self._idx_key_pairs, self._dim
             )
+            permuted_values = self.embedding_cast([permuted_values])[0]
             splitted_values = torch.split(permuted_values, self._splits, dim=self._dim)
             return _to_tensor_dict(self._keys, splitted_values)
 

torchrec/modules/tests/test_regroup.py

@@ -11,10 +11,12 @@
 
 import torch
 import torch.fx
-
+from hypothesis import given, settings, strategies as st, Verbosity
+from torchrec.modules.embedding_configs import data_type_to_dtype
 from torchrec.modules.regroup import KTRegroupAsDict
 from torchrec.sparse.jagged_tensor import _all_keys_used_once, KeyedTensor
 from torchrec.sparse.tests.utils import build_groups, build_kts
+from torchrec.types import DataType
 
 
 class KTRegroupAsDictTest(unittest.TestCase):
@@ -171,3 +173,25 @@ def test_fx_and_jit_regroup_skips_and_duplicates(self) -> None:
         eager_out = regroup_module(self.kts)
         for key in out.keys():
             torch.allclose(out[key], eager_out[key])
+
+    # pyre-ignore[56]
+    @given(data_type=st.sampled_from([DataType.BF16, DataType.FP16]))
+    @settings(verbosity=Verbosity.verbose, max_examples=20)
+    def test_regroup_cast(self, data_type: DataType) -> None:
+        dtype = data_type_to_dtype(data_type)
+        groups = build_groups(
+            kts=self.kts, num_groups=self.num_groups, skips=True, duplicates=True
+        )
+        assert _all_keys_used_once(self.kts, groups) is False
+
+        regroup_module = KTRegroupAsDict(groups=groups, keys=self.keys)
+        cast_regroup = KTRegroupAsDict(
+            groups=groups, keys=self.keys, emb_dtype=data_type
+        )
+
+        eager_out = regroup_module(self.kts)
+        cast_out = cast_regroup(self.kts)
+
+        for key in eager_out.keys():
+            self.assertEqual(cast_out[key].dtype, dtype)
+            torch.allclose(cast_out[key], eager_out[key].to(dtype))