Fix per-feature boundary lost by ZCH index dedup (#4525)

Summary:

`fbgemm.jagged_unique_indices` dedups a whole hash range at once, and
`_create_dedup_indices` puts every feature of a table into one range, so the op
returns the group's total unique count spread evenly over the group's
(feature, batch) slots rather than the real per-feature counts. A table binding a
single feature is unaffected. A table binding two features of very unequal length
gets ids relabeled across the boundary.

Managed collision modules read the feature an id arrives under as a control
signal, so this is a correctness bug rather than just wrong bookkeeping. On IG
Reels ESR the `media_embbedding_cache` HASH_ZCH table binds a ~33 long write
feature and a ~1280 long `_readonly` history feature; both come back as ~656, so
roughly 620 history ids per step are relabeled as candidates and inserted past
the read-only gate. The table reaches 100% occupancy within ~100 steps and, with
eviction disabled, only collides from then on.

Re-attribute every surviving row to the first feature it appeared in and regroup
the values so each feature is contiguous again. Rows shared by two features still
collapse into one, so the dedup saving is kept. Gated on a table actually binding
more than one feature, so the common path is untouched. Only per-feature totals
are made exact: the layout within a feature stays approximate, as it already was,
because reverse_indices undoes it downstream.

Also declares the `reverse_indices` field that `_dedup_indices` already appends
to. `mc_modules` shadows `EmbeddingCollectionContext` with a local dataclass that
lacks it, so a standalone sharded `ManagedCollisionCollection` with dedup on
raised `AttributeError`; only the `ManagedCollisionEmbeddingCollection` path
worked, since that context derives from the `embedding.py` class instead.

Differential Revision: D115457342

Author: yingufan

torchrec/distributed/mc_modules.py

@@ -95,6 +95,7 @@ class EmbeddingCollectionContext(Multistreamable):
         Union[InferSequenceShardingContext, SequenceShardingContext]
     ]
     input_features: List[KeyedJaggedTensor] = field(default_factory=list)
+    reverse_indices: List[torch.Tensor] = field(default_factory=list)
     # VBE-Attributes for EBC
     inverse_indices: Optional[Tuple[List[str], torch.Tensor]] = None
     variable_batch_per_feature: bool = False
@@ -106,6 +107,8 @@ def record_stream(self, stream: torch.Stream) -> None:
         for f in self.input_features:
             # pyrefly: ignore[bad-argument-type]
             f.record_stream(stream)
+        for r in self.reverse_indices:
+            r.record_stream(stream)
         if self.inverse_indices is not None:
             self.inverse_indices[1].record_stream(stream)
 
@@ -205,6 +208,80 @@ def create_mc_sharding(
         raise ValueError(f"Sharding not supported {sharding_type}")
 
 
+def _restore_dedup_feature_boundary(
+    input_lengths: torch.Tensor,
+    num_features: int,
+    unique_indices: torch.Tensor,
+    reverse_indices: torch.Tensor,
+) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    """
+    Recovers the per-feature boundary that ``fbgemm.jagged_unique_indices`` drops.
+
+    Every feature of a table is given the same hash range by
+    ``_create_dedup_indices``, so the op dedups all of a table's features as one
+    group. It has no way to tell which feature a surviving row came from, and
+    returns the group's total unique count spread evenly over the group's
+    (feature, batch) slots. A table binding a single feature is unaffected, but for
+    a table binding two features of very unequal length the even split silently
+    relabels ids as belonging to the other feature: a 33-long and a 1280-long
+    feature both come back as ~656.
+
+    Managed collision modules read the feature an id arrives under as a control
+    signal - read-only gating and per-feature eviction both key off it - so the
+    relabeled ids get written to tables they should only have been read from.
+
+    Every surviving row is re-attributed to the first feature it appeared in and the
+    values regrouped so each feature is contiguous again. Rows shared by two
+    features still collapse into one, so the dedup saving is kept.
+
+    Returns the corrected ``(lengths, unique_indices, reverse_indices)``.
+    """
+    device = unique_indices.device
+    num_unique = unique_indices.numel()
+    # lengths are int64 to match what the op returns
+    if num_unique == 0:
+        return (
+            torch.zeros(input_lengths.numel(), dtype=torch.int64, device=device),
+            unique_indices,
+            reverse_indices,
+        )
+
+    stride = input_lengths.numel() // num_features
+    # KJT lengths are feature major, so segment -> feature is a plain repeat.
+    # output_size is known, and passing it keeps repeat_interleave from syncing.
+    feature_per_value = torch.repeat_interleave(
+        torch.arange(num_features, device=device).repeat_interleave(stride),
+        input_lengths.long(),
+        output_size=reverse_indices.numel(),
+    )
+    # include_self=False so the initial value is ignored wherever a row was actually
+    # referenced. Every unique row has at least one source position, but falling back
+    # to feature 0 keeps an unreferenced row in range for the scatter_add_ below
+    # rather than letting it index out of bounds.
+    feature_per_unique = torch.zeros(
+        num_unique, dtype=torch.int64, device=device
+    ).scatter_reduce_(
+        0, reverse_indices, feature_per_value, reduce="amin", include_self=False
+    )
+
+    permute = torch.argsort(feature_per_unique, stable=True)
+    inverse_permute = torch.empty_like(permute)
+    inverse_permute[permute] = torch.arange(num_unique, device=device)
+
+    # scatter_add rather than bincount, which would force a device to host sync
+    counts = torch.zeros(num_features, dtype=torch.int64, device=device).scatter_add_(
+        0, feature_per_unique, torch.ones_like(feature_per_unique)
+    )
+    # spread each feature's count over its own batch slots, the same convention the
+    # op already applies within a single feature. Only the per-feature totals carry
+    # meaning: the layout inside a feature is undone by reverse_indices downstream.
+    lengths = counts.div(stride, rounding_mode="floor").repeat_interleave(stride) + (
+        torch.arange(stride, device=device).repeat(num_features)
+        < (counts % stride).repeat_interleave(stride)
+    )
+    return lengths, unique_indices[permute], inverse_permute[reverse_indices]
+
+
 class ShardedManagedCollisionCollection(
     ShardedModule[
         KJTList,
@@ -648,6 +725,12 @@ def _create_dedup_indices(self) -> None:
             )[-1]
             <= torch.iinfo(torch.int64).max
         ), "EC Dedup requires the mc collection to have a cumuluative 'hash_input_size' kwarg to be less than max int64.  Please reduce values of individual tables to meet this constraint (ie. 2**54 is typically a good value)."
+        # per sharding group, true when any of its tables binds more than one feature
+        # and so loses its feature boundary, see _restore_dedup_feature_boundary
+        self._dedup_restore_feature_boundary: List[bool] = [
+            any(feature_count > 1 for feature_count in feature_splits)
+            for feature_splits in self._sharding_per_table_feature_splits
+        ]
         for i, (feature_splits, input_splits) in enumerate(
             zip(
                 self._sharding_per_table_feature_splits,
@@ -688,6 +771,11 @@ def _dedup_indices(
         features_by_sharding = []
 
         for i, kjt in enumerate(features):
+            # jagged_unique_indices derives the batch size as total_B / num_features,
+            # and the per-feature regrouping below relies on the same uniform stride
+            assert (
+                not kjt.variable_stride_per_key()
+            ), "EC Dedup does not support variable stride per key"
             hash_offsets = self.get_buffer(f"_dedup_hash_offsets_{i}")
             feature_offsets = self.get_buffer(f"_dedup_feature_offsets_{i}")
             (
@@ -701,6 +789,13 @@ def _dedup_indices(
                 kjt.offsets().to(torch.int64),
                 kjt.values().to(torch.int64),
             )
+            if self._dedup_restore_feature_boundary[i]:
+                lengths, unique_indices, reverse_indices = (
+                    _restore_dedup_feature_boundary(
+                        kjt.lengths(), len(kjt.keys()), unique_indices, reverse_indices
+                    )
+                )
+                offsets = torch.ops.fbgemm.asynchronous_complete_cumsum(lengths)
             # Gather weights for unique indices if present.
             # Since unique_indices[reverse_indices[i]] == indices[i],
             # we can directly assign: dedup_weights[reverse_indices[i]] = weights[i]
@@ -723,7 +818,6 @@ def _dedup_indices(
             )
 
             ctx.input_features.append(kjt)
-            # pyrefly: ignore[missing-attribute]
             ctx.reverse_indices.append(reverse_indices)
             features_by_sharding.append(dedup_features)
         return features_by_sharding

torchrec/distributed/tests/test_mc_dedup_feature_boundary.py

@@ -0,0 +1,213 @@
+#!/usr/bin/env python3
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+# pyre-strict
+
+import unittest
+from typing import cast, Dict, List, Tuple
+
+import torch
+from torchrec.distributed.mc_modules import (
+    _restore_dedup_feature_boundary,
+    ManagedCollisionCollectionContext,
+    ShardedManagedCollisionCollection,
+)
+from torchrec.sparse.jagged_tensor import KeyedJaggedTensor
+
+_HASH_SIZE = 100_000
+_STRIDE = 2  # batch size
+
+
+def _dedup(
+    values: torch.Tensor, lengths: torch.Tensor, num_features: int
+) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    """
+    Runs the dedup op with the buffers ``_create_dedup_indices`` builds for a single
+    table binding ``num_features`` features: one shared hash range, one feature group.
+    """
+    device = values.device
+    hash_offsets = torch.tensor(
+        [0] * num_features + [_HASH_SIZE], dtype=torch.int64, device=device
+    )
+    feature_offsets = torch.tensor(
+        [0] * num_features + [num_features], dtype=torch.int64, device=device
+    )
+    lengths_, _, unique_indices, reverse_indices = (
+        torch.ops.fbgemm.jagged_unique_indices(
+            hash_offsets,
+            feature_offsets,
+            torch.ops.fbgemm.asynchronous_complete_cumsum(lengths),
+            values,
+        )
+    )
+    return lengths_, unique_indices, reverse_indices
+
+
+def _per_feature(lengths: torch.Tensor, num_features: int) -> List[int]:
+    return lengths.view(num_features, -1).sum(dim=1).tolist()
+
+
+def _lengths(counts: List[int], device: torch.device) -> torch.Tensor:
+    """Per-feature totals -> a feature major KJT lengths tensor of stride _STRIDE."""
+    return torch.tensor(
+        [
+            count // _STRIDE + (1 if b < count % _STRIDE else 0)
+            for count in counts
+            for b in range(_STRIDE)
+        ],
+        dtype=torch.int64,
+        device=device,
+    )
+
+
+class _DedupOnlyCollection:
+    """
+    Carries just the state ``ShardedManagedCollisionCollection._dedup_indices`` reads
+    off ``self``, so the method can be exercised without standing up a sharded module.
+    """
+
+    def __init__(self, num_features: int, device: torch.device) -> None:
+        self._buffers: Dict[str, torch.Tensor] = {
+            "_dedup_hash_offsets_0": torch.tensor(
+                [0] * num_features + [_HASH_SIZE], dtype=torch.int64, device=device
+            ),
+            "_dedup_feature_offsets_0": torch.tensor(
+                [0] * num_features + [num_features], dtype=torch.int64, device=device
+            ),
+        }
+        self._dedup_restore_feature_boundary: List[bool] = [num_features > 1]
+
+    def get_buffer(self, name: str) -> torch.Tensor:
+        return self._buffers[name]
+
+
+@unittest.skipIf(
+    not torch.cuda.is_available(), "fbgemm.jagged_unique_indices is CUDA only"
+)
+class McDedupFeatureBoundaryTest(unittest.TestCase):
+    """
+    A managed collision table that binds more than one feature gives all of its
+    features one shared hash range, so ``fbgemm.jagged_unique_indices`` dedups them as
+    a single group and cannot say which feature a surviving row came from. These tests
+    cover the boundary that ``_restore_dedup_feature_boundary`` puts back.
+
+    Shapes mirror the IG Reels ESR memory layer, where a HASH_ZCH table binds a short
+    write feature of candidate ids and a long read-only feature of user history.
+    """
+
+    def setUp(self) -> None:
+        self.device = torch.device("cuda")
+        torch.manual_seed(0)
+        self.ids: torch.Tensor = torch.randperm(_HASH_SIZE, device=self.device)
+
+    def test_dedup_op_loses_feature_boundary(self) -> None:
+        """
+        Pins the upstream behavior being corrected: the op returns the group's unique
+        count spread evenly over the group's slots rather than the real per-feature
+        counts. If this fails, the op has learned per-feature attribution and
+        ``_restore_dedup_feature_boundary`` can be dropped.
+        """
+        write_len, read_len = 33, 1280
+        values = self.ids[: write_len + read_len]
+        lengths, unique_indices, _ = _dedup(
+            values, _lengths([write_len, read_len], self.device), 2
+        )
+
+        self.assertEqual(unique_indices.numel(), write_len + read_len)
+        self.assertEqual(_per_feature(lengths, 2), [657, 656])
+
+    def test_restore_feature_boundary(self) -> None:
+        write_len, read_len = 33, 1280
+        values = self.ids[: write_len + read_len]
+        input_lengths = _lengths([write_len, read_len], self.device)
+        _, unique_indices, reverse_indices = _dedup(values, input_lengths, 2)
+
+        lengths, unique_indices, reverse_indices = _restore_dedup_feature_boundary(
+            input_lengths, 2, unique_indices, reverse_indices
+        )
+
+        self.assertEqual(_per_feature(lengths, 2), [write_len, read_len])
+        # the write feature's segment holds exactly the ids it sent, nothing borrowed
+        # from the read-only feature
+        torch.testing.assert_close(
+            unique_indices[:write_len].sort().values, values[:write_len].sort().values
+        )
+        torch.testing.assert_close(unique_indices[reverse_indices], values)
+
+        # the 2D weights of the write path are gathered through reverse_indices, so
+        # they have to survive the regrouping too
+        weights = torch.rand(values.numel(), 4, device=self.device)
+        dedup_weights = torch.empty(
+            unique_indices.numel(), 4, dtype=weights.dtype, device=self.device
+        )
+        dedup_weights[reverse_indices] = weights
+        torch.testing.assert_close(dedup_weights[reverse_indices], weights)
+
+    def test_restore_feature_boundary_with_duplicate_and_shared_ids(self) -> None:
+        """
+        Dedup still collapses repeats, including ids sent by both features. A shared id
+        is attributed to the first feature that carried it, which is the write feature,
+        so a candidate that is also in history is still inserted rather than gated.
+        """
+        write_len, hist_len, shared_len = 33, 600, 5
+        write_ids = self.ids[:write_len]
+        hist_ids = self.ids[write_len : write_len + hist_len]
+        read_values = torch.cat([hist_ids, hist_ids, write_ids[:shared_len]])
+        values = torch.cat([write_ids, read_values])
+        input_lengths = _lengths([write_len, read_values.numel()], self.device)
+        _, unique_indices, reverse_indices = _dedup(values, input_lengths, 2)
+
+        lengths, unique_indices, reverse_indices = _restore_dedup_feature_boundary(
+            input_lengths, 2, unique_indices, reverse_indices
+        )
+
+        self.assertEqual(unique_indices.numel(), write_len + hist_len)
+        self.assertEqual(_per_feature(lengths, 2), [write_len, hist_len])
+        torch.testing.assert_close(
+            unique_indices[:write_len].sort().values, write_ids.sort().values
+        )
+        torch.testing.assert_close(unique_indices[reverse_indices], values)
+
+    def test_dedup_indices_keeps_feature_boundary(self) -> None:
+        """End to end through ``_dedup_indices``, covering the wiring and the gate."""
+        write_len, read_len = 33, 1280
+        kjt = KeyedJaggedTensor(
+            keys=["media_embbedding_cache", "merged_events_item_id_cache_readonly"],
+            values=self.ids[: write_len + read_len],
+            lengths=_lengths([write_len, read_len], self.device),
+        )
+        ctx = ManagedCollisionCollectionContext(sharding_contexts=[])
+
+        dedup_kjt = ShardedManagedCollisionCollection._dedup_indices(
+            cast(
+                ShardedManagedCollisionCollection, _DedupOnlyCollection(2, self.device)
+            ),
+            ctx,
+            [kjt],
+        )[0]
+
+        self.assertEqual(dedup_kjt.length_per_key(), [write_len, read_len])
+        torch.testing.assert_close(
+            dedup_kjt["media_embbedding_cache"].values().sort().values,
+            kjt["media_embbedding_cache"].values().sort().values,
+        )
+        torch.testing.assert_close(
+            dedup_kjt.values()[ctx.reverse_indices[0]], kjt.values()
+        )
+
+    def test_restore_feature_boundary_with_three_features(self) -> None:
+        counts = [11, 517, 42]
+        values = self.ids[: sum(counts)]
+        input_lengths = _lengths(counts, self.device)
+        _, unique_indices, reverse_indices = _dedup(values, input_lengths, 3)
+
+        lengths, unique_indices, reverse_indices = _restore_dedup_feature_boundary(
+            input_lengths, 3, unique_indices, reverse_indices
+        )
+
+        self.assertEqual(_per_feature(lengths, 3), counts)
+        torch.testing.assert_close(unique_indices[reverse_indices], values)