Clean up KJT validator killswitch (#3615)

Summary:
Pull Request resolved: #3615

The killswitch `pytorch/torchrec:enable_kjt_validation` has been switched ON for a couple of months and it is working normally, so it should be safe to clean it up.

Reviewed By: TroyGarden

Differential Revision: D89088119

fbshipit-source-id: 9f339867f192d76224155bceb14619e98f35ff0e

Author: Shuangping Liu

torchrec/distributed/embeddingbag.py

@@ -1539,11 +1539,8 @@ def input_dist(
         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._embedding_bag_configs)
+            logger.info("Validating input features...")
+            validate_keyed_jagged_tensor(features, self._embedding_bag_configs)
 
             self._create_input_dist(features.keys())
             self._has_uninitialized_input_dist = False

torchrec/distributed/test_utils/test_model_parallel_base.py

@@ -240,102 +240,37 @@ def test_sharding_fused_ebc_as_top_level(self) -> None:
 
         self.assertTrue(isinstance(model.module, ShardedFusedEmbeddingBagCollection))
 
-    @patch("torch._utils_internal.justknobs_check")
-    def test_sharding_ebc_input_validation_enabled(self, mock_jk: Mock) -> None:
+    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]),
         )
-        mock_jk.return_value = True
 
         with self.assertRaisesRegex(ValueError, "keys must be unique"):
             model(kjt)
 
-        # Count only calls with the input "pytorch/torchrec:enable_kjt_validation"
-        # This ignores any other calls to justknobs_check() with other inputs
-        # and protects the test from breaking when new JK checks are added.
-        validation_calls = [
-            call
-            for call in mock_jk.call_args_list
-            if len(call[0]) > 0
-            and call[0][0] == "pytorch/torchrec:enable_kjt_validation"
-        ]
-        self.assertEqual(
-            1,
-            len(validation_calls),
-            "There should be exactly one call to JK with pytorch/torchrec:enable_kjt_validation",
-        )
-
-    @patch("torch._utils_internal.justknobs_check")
-    def test_sharding_ebc_validate_input_only_once(self, mock_jk: Mock) -> None:
+    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)
-        mock_jk.return_value = True
 
         with self.assertLogs(embeddingbag_logger, level="INFO") as logs:
             model(kjt)
             model(kjt)
             model(kjt)
 
-        # Count only calls with the input "pytorch/torchrec:enable_kjt_validation"
-        # This ignores any other calls to justknobs_check() with other inputs
-        # and protects the test from breaking when new JK checks are added.
-        validation_calls = [
-            call
-            for call in mock_jk.call_args_list
-            if len(call[0]) > 0
-            and call[0][0] == "pytorch/torchrec:enable_kjt_validation"
-        ]
-        self.assertEqual(
-            1,
-            len(validation_calls),
-            "There should be exactly one call to JK with pytorch/torchrec:enable_kjt_validation",
-        )
         matched_logs = list(
             filter(lambda s: "Validating input features..." in s, logs.output)
         )
         self.assertEqual(1, len(matched_logs))
 
-    @patch("torch._utils_internal.justknobs_check")
-    def test_sharding_ebc_input_validation_disabled(self, mock_jk: Mock) -> 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)
-        mock_jk.return_value = False
-
-        # Without KJT validation, input_dist will not raise exceptions
-        try:
-            model(kjt)
-        except ValueError:
-            self.fail("Input validation should not be enabled.")
-
-        # Count only calls with the input "pytorch/torchrec:enable_kjt_validation"
-        # This ignores any other calls to justknobs_check() with other inputs
-        # and protects the test from breaking when new JK checks are added.
-        validation_calls = [
-            call
-            for call in mock_jk.call_args_list
-            if len(call[0]) > 0
-            and call[0][0] == "pytorch/torchrec:enable_kjt_validation"
-        ]
-        self.assertEqual(
-            1,
-            len(validation_calls),
-            "There should be exactly one call to JK with pytorch/torchrec:enable_kjt_validation",
-        )
-
     def _create_sharded_model(
         self, embedding_dim: int = 128, num_embeddings: int = 256
     ) -> DistributedModelParallel: