Fix return results of forward method of DecoupleEmbeddingColletion (#2861)

Summary:
Pull Request resolved: #2861

The decouple_di_pass will replace ManagedCollisionEmbeddingCollection with DecoupleEmbeddingColletion. The forward method of the former returns a tuple of the embedding results and the features inself (Tuple[Dict[str, KJT], KJT]), but the latter only returns the embedding results (Dict[str, KJT]). This caused incompatibility issues in following TGIF transform passes, namely QuantizationPass.

Reviewed By: kausv

Differential Revision: D71412636

fbshipit-source-id: e1489422ace0e3d3cb88e0aa040b5c1c09ed65ad

Author: joshuuuasu

torchrec/distributed/quant_embedding.py

@@ -1189,7 +1189,9 @@ def __init__(
     def forward(
         self,
         features: KeyedJaggedTensor,
-    ) -> torch.Tensor:
+    ) -> Tuple[
+        Union[KeyedTensor, Dict[str, JaggedTensor]], Optional[KeyedJaggedTensor]
+    ]:
         remapped_kjt = self._mcc_remapper(features)
         return self._ec_lookup(remapped_kjt)
 

torchrec/quant/embedding_modules.py

@@ -1091,15 +1091,16 @@ def to(
         )
         return self
 
+    # pyre-ignore
     def forward(
         self,
         features: KeyedJaggedTensor,
-    ) -> Dict[str, JaggedTensor]:
+    ) -> Tuple[
+        Union[KeyedTensor, Dict[str, JaggedTensor]], Optional[KeyedJaggedTensor]
+    ]:
         features = self._managed_collision_collection(features)
 
-        # mcec expects Tuple return type
-        # pyre-ignore
-        return (super().forward(features),)
+        return (super().forward(features), features)
 
     def _get_name(self) -> str:
         return "QuantManagedCollisionEmbeddingCollection"