Add input KJT validation in EBC input_dist (#3328)

Summary:
Pull Request resolved: #3328

Validates input KJT in EBC `input_dist`.

This is executed when initializing `input_dist` of `ShardedEmbeddingBagCollection`, so that input features are validated exactly **once** per EBC per rank, assuming the shape of the first batch is representative across all following batches.

Reviewed By: TroyGarden

Differential Revision: D71752961

Author: Shuangping Liu

torchrec/distributed/embeddingbag.py

@@ -8,6 +8,7 @@
 # pyre-strict
 
 import copy
+import logging
 from collections import defaultdict, OrderedDict
 from dataclasses import dataclass, field
 from functools import partial
@@ -109,6 +110,7 @@
 from torchrec.optim.fused import EmptyFusedOptimizer, FusedOptimizerModule
 from torchrec.optim.keyed import CombinedOptimizer, KeyedOptimizer
 from torchrec.sparse.jagged_tensor import _to_offsets, KeyedJaggedTensor, KeyedTensor
+from torchrec.sparse.jagged_tensor_validator import validate_keyed_jagged_tensor
 from torchrec.sparse.tensor_dict import maybe_td_to_kjt
 
 try:
@@ -119,6 +121,9 @@
     pass
 
 
+logger: logging.Logger = logging.getLogger(__name__)
+
+
 def _pin_and_move(tensor: torch.Tensor, device: torch.device) -> torch.Tensor:
     return (
         tensor
@@ -1515,13 +1520,21 @@ def input_dist(
             features = maybe_td_to_kjt(features, feature_keys)  # pyre-ignore[6]
         ctx.variable_batch_per_feature = features.variable_stride_per_key()
         ctx.inverse_indices = features.inverse_indices_or_none()
+
         if self._has_uninitialized_input_dist:
+            if torch._utils_internal.justknobs_check(
+                "pytorch/torchrec:enable_kjt_validation"
+            ):
+                logger.info("Validating input features...")
+                validate_keyed_jagged_tensor(features)
+
             self._create_input_dist(features.keys())
             self._has_uninitialized_input_dist = False
             if ctx.variable_batch_per_feature:
                 self._create_inverse_indices_permute_indices(ctx.inverse_indices)
             if self._has_mean_pooling_callback:
                 self._init_mean_pooling_callback(features.keys(), ctx.inverse_indices)
+
         with torch.no_grad():
             if self._has_features_permute:
                 features = features.permute(

torchrec/distributed/test_utils/test_model_parallel_base.py

@@ -20,6 +20,7 @@
     PartiallyMaterializedTensor,
 )
 from hypothesis import assume, given, settings, strategies as st, Verbosity
+from pyjk import PyPatchJustKnobs
 from torch import distributed as dist
 from torch.distributed._shard.sharded_tensor import ShardedTensor
 from torch.distributed._tensor import DTensor
@@ -28,7 +29,10 @@
     EmbeddingComputeKernel,
     EmbeddingTableConfig,
 )
-from torchrec.distributed.embeddingbag import ShardedEmbeddingBagCollection
+from torchrec.distributed.embeddingbag import (
+    logger as embeddingbag_logger,
+    ShardedEmbeddingBagCollection,
+)
 from torchrec.distributed.fused_embeddingbag import ShardedFusedEmbeddingBagCollection
 from torchrec.distributed.model_parallel import DistributedModelParallel
 from torchrec.distributed.planner import (
@@ -65,6 +69,7 @@
 from torchrec.modules.embedding_modules import EmbeddingBagCollection
 from torchrec.modules.fused_embedding_modules import FusedEmbeddingBagCollection
 from torchrec.optim.rowwise_adagrad import RowWiseAdagrad
+from torchrec.sparse.jagged_tensor import KeyedJaggedTensor
 from torchrec.test_utils import get_free_port, seed_and_log
 
 
@@ -205,13 +210,23 @@ def setUp(self, backend: str = "nccl") -> None:
 
         dist.init_process_group(backend=self.backend)
 
+    @classmethod
+    def setUpClass(cls) -> None:
+        super().setUpClass()
+        cls.patcher = PyPatchJustKnobs()
+
     def tearDown(self) -> None:
         dist.destroy_process_group()
 
     def test_sharding_ebc_as_top_level(self) -> None:
+        model = self._create_sharded_model()
+
+        self.assertTrue(isinstance(model.module, ShardedEmbeddingBagCollection))
+
+    def test_sharding_fused_ebc_as_top_level(self) -> None:
         embedding_dim = 128
         num_embeddings = 256
-        ebc = EmbeddingBagCollection(
+        ebc = FusedEmbeddingBagCollection(
             device=torch.device("meta"),
             tables=[
                 EmbeddingBagConfig(
@@ -222,16 +237,67 @@ def test_sharding_ebc_as_top_level(self) -> None:
                     pooling=PoolingType.SUM,
                 ),
             ],
+            optimizer_type=torch.optim.SGD,
+            optimizer_kwargs={"lr": 0.02},
         )
 
         model = DistributedModelParallel(ebc, device=self.device)
 
-        self.assertTrue(isinstance(model.module, ShardedEmbeddingBagCollection))
+        self.assertTrue(isinstance(model.module, ShardedFusedEmbeddingBagCollection))
 
-    def test_sharding_fused_ebc_as_top_level(self) -> None:
-        embedding_dim = 128
-        num_embeddings = 256
-        ebc = FusedEmbeddingBagCollection(
+    def test_sharding_ebc_input_validation_enabled(self) -> None:
+        model = self._create_sharded_model()
+        kjt = KeyedJaggedTensor(
+            keys=["my_feature", "my_feature"],
+            values=torch.tensor([1, 2, 3, 4, 5]),
+            lengths=torch.tensor([1, 2, 0, 2]),
+            offsets=torch.tensor([0, 1, 3, 3, 5]),
+        )
+
+        with self.patcher.patch("pytorch/torchrec:enable_kjt_validation", True):
+            with self.assertRaisesRegex(ValueError, "keys must be unique"):
+                model(kjt)
+
+    def test_sharding_ebc_validate_input_only_once(self) -> None:
+        model = self._create_sharded_model()
+        kjt = KeyedJaggedTensor(
+            keys=["my_feature"],
+            values=torch.tensor([1, 2, 3, 4, 5]),
+            lengths=torch.tensor([1, 2, 0, 2]),
+            offsets=torch.tensor([0, 1, 3, 3, 5]),
+        ).to(self.device)
+
+        with self.patcher.patch("pytorch/torchrec:enable_kjt_validation", True):
+            with self.assertLogs(embeddingbag_logger, level="INFO") as logs:
+                model(kjt)
+                model(kjt)
+                model(kjt)
+
+        matched_logs = list(
+            filter(lambda s: "Validating input features..." in s, logs.output)
+        )
+        self.assertEqual(1, len(matched_logs))
+
+    def test_sharding_ebc_input_validation_disabled(self) -> None:
+        model = self._create_sharded_model()
+        kjt = KeyedJaggedTensor(
+            keys=["my_feature", "my_feature"],
+            values=torch.tensor([1, 2, 3, 4, 5]),
+            lengths=torch.tensor([1, 2, 0, 2]),
+            offsets=torch.tensor([0, 1, 3, 3, 5]),
+        ).to(self.device)
+
+        # Without KJT validation, input_dist will not raise exceptions
+        with self.patcher.patch("pytorch/torchrec:enable_kjt_validation", False):
+            try:
+                model(kjt)
+            except ValueError:
+                self.fail("Input validation should not be enabled.")
+
+    def _create_sharded_model(
+        self, embedding_dim: int = 128, num_embeddings: int = 256
+    ) -> DistributedModelParallel:
+        ebc = EmbeddingBagCollection(
             device=torch.device("meta"),
             tables=[
                 EmbeddingBagConfig(
@@ -242,13 +308,8 @@ def test_sharding_fused_ebc_as_top_level(self) -> None:
                     pooling=PoolingType.SUM,
                 ),
             ],
-            optimizer_type=torch.optim.SGD,
-            optimizer_kwargs={"lr": 0.02},
         )
-
-        model = DistributedModelParallel(ebc, device=self.device)
-
-        self.assertTrue(isinstance(model.module, ShardedFusedEmbeddingBagCollection))
+        return DistributedModelParallel(ebc, device=self.device)
 
 
 class ModelParallelSingleRankBase(unittest.TestCase):