fx friendly jagged tensor construction (#1757)

Summary:
Pull Request resolved: #1757

* introduce a new fx friendly jagged_tensor construction. so that the graph does not get broken by kjt nodes, this partially facilitate a tensor only fx spliited graph.

Reviewed By: 842974287

Differential Revision: D54473827

fbshipit-source-id: 42c941ce3f54bb5bc9335f1c71302263822eb139

torchrec/modules/utils.py

@@ -14,6 +14,11 @@
 from torchrec.sparse.jagged_tensor import JaggedTensor, KeyedJaggedTensor
 
 
+@torch.fx.wrap
+def _fx_to_list(tensor: torch.Tensor) -> List[int]:
+    return tensor.long().tolist()
+
+
 def extract_module_or_tensor_callable(
     module_or_callable: Union[
         Callable[[], torch.nn.Module],
@@ -167,6 +172,54 @@ def construct_jagged_tensors(
         return ret
 
 
+def construct_jagged_tensors_inference(
+    embeddings: torch.Tensor,
+    lengths: torch.Tensor,
+    values: torch.Tensor,
+    embedding_names: List[str],
+    need_indices: bool = False,
+    features_to_permute_indices: Optional[Dict[str, List[int]]] = None,
+    reverse_indices: Optional[torch.Tensor] = None,
+) -> Dict[str, JaggedTensor]:
+    with record_function("## construct_jagged_tensors_inference ##"):
+        if reverse_indices is not None:
+            embeddings = torch.index_select(
+                embeddings, 0, reverse_indices.to(torch.int32)
+            )
+
+        ret: Dict[str, JaggedTensor] = {}
+        length_per_key: List[int] = _fx_to_list(
+            torch.sum(lengths.view(len(embedding_names), -1), dim=1)
+        )
+
+        lengths = lengths.view(len(embedding_names), -1)
+        lengths_tuple = torch.unbind(lengths, dim=0)
+        embeddings_list = torch.split(embeddings, length_per_key, dim=0)
+        values_list = torch.split(values, length_per_key) if need_indices else None
+
+        key_indices = defaultdict(list)
+        for i, key in enumerate(embedding_names):
+            key_indices[key].append(i)
+        for key, indices in key_indices.items():
+            # combines outputs in correct order for CW sharding
+            indices = (
+                _permute_indices(indices, features_to_permute_indices[key])
+                if features_to_permute_indices and key in features_to_permute_indices
+                else indices
+            )
+            ret[key] = JaggedTensor(
+                lengths=lengths_tuple[indices[0]],
+                values=(
+                    embeddings_list[indices[0]]
+                    if len(indices) == 1
+                    else torch.cat([embeddings_list[i] for i in indices], dim=1)
+                ),
+                # pyre-ignore
+                weights=values_list[indices[0]] if need_indices else None,
+            )
+        return ret
+
+
 def _permute_indices(indices: List[int], permute: List[int]) -> List[int]:
     permuted_indices = [0] * len(indices)
     for i, permuted_index in enumerate(permute):

torchrec/quant/embedding_modules.py

@@ -43,7 +43,7 @@
     FeatureProcessedEmbeddingBagCollection as OriginalFeatureProcessedEmbeddingBagCollection,
 )
 
-from torchrec.modules.utils import construct_jagged_tensors
+from torchrec.modules.utils import construct_jagged_tensors_inference
 from torchrec.sparse.jagged_tensor import JaggedTensor, KeyedJaggedTensor, KeyedTensor
 from torchrec.tensor_types import UInt2Tensor, UInt4Tensor
 from torchrec.types import ModuleNoCopyMixin
@@ -865,16 +865,18 @@ def forward(
             indices = f.values()
             lengths = f.lengths()
             offsets = f.offsets()
+            lengths = f.lengths()
             lookup = (
                 emb_module(indices=indices, offsets=offsets)
                 if self.register_tbes
                 else emb_module.forward(indices=indices, offsets=offsets)
             )
             lookup = _get_batching_hinted_output(lengths=lengths, output=lookup)
             embedding_names = self._embedding_names_by_batched_tables[key]
-            jt = construct_jagged_tensors(
+            jt = construct_jagged_tensors_inference(
                 embeddings=lookup,
-                features=f,
+                lengths=lengths,
+                values=indices,
                 embedding_names=embedding_names,
                 need_indices=self.need_indices(),
             )