[bugfix] fix correctness of kjt.lengths when ShardedEmbeddingBag’s pooling_type is mean and shard_type is row_wise (#106)

Author: tiankongdeguiji

tzrec/tests/configs/multi_tower_din_fg_mock.config

@@ -59,6 +59,8 @@ feature_configs {
         hash_bucket_size: 100
         embedding_dim: 16
         embedding_name: "id_4_emb"
+        value_dim: 0
+        pooling: "mean"
     }
 }
 feature_configs {
@@ -68,6 +70,8 @@ feature_configs {
         hash_bucket_size: 100
         embedding_dim: 16
         embedding_name: "id_4_emb"
+        value_dim: 0
+        pooling: "mean"
     }
 }
 feature_configs {

tzrec/utils/dist_util.py

@@ -9,40 +9,15 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from collections import OrderedDict
-from typing import List
+from typing import Dict, List, Optional, Tuple
 
 import torch
 from torch import distributed as dist
-from torch import nn
-from torchrec.distributed.types import ShardingPlan, ShardingType
-
-
-def sync_dp_emb_table(model: nn.Module, plan: ShardingPlan) -> None:
-    """Sync data parallel embedding table params."""
-    dp_param_names = []
-    for _, module_plan in plan.plan.items():
-        # pyre-ignore [16]
-        for param_name, param_sharding in module_plan.items():
-            if param_sharding.sharding_type == ShardingType.DATA_PARALLEL.value:
-                dp_param_names.append(param_name)
-    dp_params = OrderedDict()
-    for name, param in model.named_parameters():
-        name_parts = name.split(".")
-        if (
-            len(name_parts) > 2
-            and name_parts[-1] == "weight"
-            and name_parts[-2] in dp_param_names
-        ):
-            # pyre-ignore [16]
-            ori_t = param._original_tensor
-            if ori_t not in dp_params:
-                dp_params[ori_t] = 1
-    broadcast_works = []
-    for t in dp_params:
-        broadcast_works.append(dist.broadcast(t.detach(), src=0, async_op=True))
-    for w in broadcast_works:
-        w.wait()
+from torch.autograd.profiler import record_function
+from torchrec.distributed import embeddingbag
+from torchrec.distributed.utils import none_throws
+from torchrec.modules.embedding_configs import PoolingType
+from torchrec.sparse.jagged_tensor import _to_offsets
 
 
 def broadcast_string(s: str, src: int = 0) -> str:
@@ -106,3 +81,80 @@ def gather_strings(s: str, dst: int = 0) -> List[str]:
             gathered_strings.append(string)
 
     return gathered_strings
+
+
+# lengths of kjt will be modified by create_mean_pooling_divisor, we fix it
+# with lengths = lengths.clone() temporarily.
+def _create_mean_pooling_divisor(
+    lengths: torch.Tensor,
+    keys: List[str],
+    offsets: torch.Tensor,
+    stride: int,
+    stride_per_key: List[int],
+    dim_per_key: torch.Tensor,
+    pooling_type_to_rs_features: Dict[str, List[str]],
+    embedding_names: List[str],
+    embedding_dims: List[int],
+    variable_batch_per_feature: bool,
+    kjt_inverse_order: torch.Tensor,
+    kjt_key_indices: Dict[str, int],
+    kt_key_ordering: torch.Tensor,
+    inverse_indices: Optional[Tuple[List[str], torch.Tensor]] = None,
+    weights: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    with record_function("## ebc create mean pooling callback ##"):
+        batch_size = (
+            none_throws(inverse_indices)[1].size(dim=1)
+            if variable_batch_per_feature
+            else stride
+        )
+
+        if weights is not None:
+            # if we have weights, lengths is the sum of weights by offsets for feature
+            lengths = torch.ops.fbgemm.segment_sum_csr(1, offsets.int(), weights)
+
+        if variable_batch_per_feature:
+            inverse_indices = none_throws(inverse_indices)
+            device = inverse_indices[1].device
+            inverse_indices_t = inverse_indices[1]
+            if len(keys) != len(inverse_indices[0]):
+                inverse_indices_t = torch.index_select(
+                    inverse_indices[1], 0, kjt_inverse_order
+                )
+            offsets = _to_offsets(torch.tensor(stride_per_key, device=device))[
+                :-1
+            ].unsqueeze(-1)
+            indices = (inverse_indices_t + offsets).flatten()
+            lengths = torch.index_select(input=lengths, dim=0, index=indices)
+
+        # only convert the sum pooling features to be 1 lengths
+        lengths = lengths.clone()
+        for feature in pooling_type_to_rs_features[PoolingType.SUM.value]:
+            feature_index = kjt_key_indices[feature]
+            feature_index = feature_index * batch_size
+            lengths[feature_index : feature_index + batch_size] = 1
+
+        if len(embedding_names) != len(keys):
+            lengths = torch.index_select(
+                lengths.reshape(-1, batch_size),
+                0,
+                kt_key_ordering,
+            ).reshape(-1)
+
+        # transpose to align features with keyed tensor dim_per_key
+        lengths = lengths.reshape(-1, batch_size).T  # [batch_size, num_features]
+        output_size = sum(embedding_dims)
+
+        divisor = torch.repeat_interleave(
+            input=lengths,
+            repeats=dim_per_key,
+            dim=1,
+            output_size=output_size,
+        )
+        eps = 1e-6  # used to safe guard against 0 division
+        divisor = divisor + eps
+        return divisor.detach()
+
+
+# pyre-ignore [9]
+embeddingbag._create_mean_pooling_divisor = _create_mean_pooling_divisor