Enable Output Dist module when DI + Lowering is required (#2846)

Summary:
Pull Request resolved: #2846

Enable Output Dist module when DI + Lowering is required. This is mainly to handle empty or zero tensors being passed across the boundary of Intermodules and merge as DI look up results during publishing may lead to this kind of set up.

Reviewed By: jiayisuse

Differential Revision: D71671275

fbshipit-source-id: 94936476d361b299ee7e5ffb9884f5265e364765

Author: faran928

torchrec/distributed/quant_embedding.py

@@ -300,6 +300,15 @@ def _construct_jagged_tensors_tw(
     return ret
 
 
+@torch.fx.wrap
+def _fx_marker_construct_jagged_tensor(
+    values: torch.Tensor,
+    lengths: torch.Tensor,
+    weights: Optional[torch.Tensor],
+) -> JaggedTensor:
+    return JaggedTensor(values=values, lengths=lengths, weights=weights)
+
+
 def _construct_jagged_tensors_rw(
     embeddings: List[torch.Tensor],
     feature_keys: List[str],
@@ -326,7 +335,7 @@ def _construct_jagged_tensors_rw(
             length_per_key,
         )
     for i, key in enumerate(feature_keys):
-        ret[key] = JaggedTensor(
+        ret[key] = _fx_marker_construct_jagged_tensor(
             values=embs_split_per_key[i],
             lengths=lengths_list[i],
             weights=values_list[i] if need_indices else None,
@@ -455,6 +464,7 @@ def _construct_jagged_tensors(
                 rw_bucket_mapping_tensor,
             ),
         )
+
     elif sharding_type == ShardingType.COLUMN_WISE.value:
         return _construct_jagged_tensors_cw(
             embeddings,
@@ -1112,6 +1122,7 @@ def forward(self, features: KeyedJaggedTensor) -> Tuple[
         List[Optional[torch.Tensor]],
         List[Optional[torch.Tensor]],
     ]:
+
         with torch.no_grad():
             ret: List[KJTList] = []
             unbucketize_permute_tensor = []