fx wrap entrypoint of EmbeddingBagCollection (#1801)

Summary:
Pull Request resolved: #1801

We want to tag the part before tbe in ebc to INPUT_DIST net, in order to make the graph structure after splitting similar to sharded ebc. Using fx wrap to tag the entrypoint of ebc and we will use this in the splitting tag rules.

Reviewed By: jiayisuse, gnahzg, 842974287

Differential Revision: D55023624

fbshipit-source-id: 6d132f0b313bb8477e053f91c6be0b9e82e7b322

Author: seanx92

torchrec/quant/embedding_modules.py

@@ -230,6 +230,13 @@ def _update_embedding_configs(
         )
 
 
+@torch.fx.wrap
+def features_to_dict(
+    features: KeyedJaggedTensor,
+) -> Dict[str, JaggedTensor]:
+    return features.to_dict()
+
+
 class EmbeddingBagCollection(EmbeddingBagCollectionInterface, ModuleNoCopyMixin):
     """
     EmbeddingBagCollection represents a collection of pooled embeddings (EmbeddingBags).
@@ -452,7 +459,7 @@ def forward(
             KeyedTensor
         """
 
-        feature_dict = features.to_dict()
+        feature_dict = features_to_dict(features)
         embeddings = []
 
         # TODO ideally we can accept KJTs with any feature order. However, this will require an order check + permute, which will break torch.script.

torchrec/quant/tests/test_embedding_modules.py

@@ -29,6 +29,7 @@
 from torchrec.quant.embedding_modules import (
     EmbeddingBagCollection as QuantEmbeddingBagCollection,
     EmbeddingCollection as QuantEmbeddingCollection,
+    features_to_dict,
     quant_prep_enable_quant_state_dict_split_scale_bias,
 )
 from torchrec.sparse.jagged_tensor import JaggedTensor, KeyedJaggedTensor, KeyedTensor
@@ -446,6 +447,23 @@ def test_trace_and_script(self) -> None:
 
         gm = symbolic_trace(qebc)
 
+        non_placeholder_nodes = [
+            node for node in gm.graph.nodes if node.op != "placeholder"
+        ]
+        self.assertTrue(
+            len(non_placeholder_nodes) > 0, "Graph must have non-placeholder nodes"
+        )
+        self.assertEqual(
+            non_placeholder_nodes[0].op,
+            "call_function",
+            f"First non-placeholder node must be call_function, got {non_placeholder_nodes[0].op} instead",
+        )
+        self.assertEqual(
+            non_placeholder_nodes[0].name,
+            features_to_dict.__name__,
+            f"First non-placeholder node must be features_to_dict, got {non_placeholder_nodes[0].name} instead",
+        )
+
         features = KeyedJaggedTensor(
             keys=["f1", "f2"],
             values=torch.as_tensor([0, 1]),