Back out "fix flaky test due to input_jkt.weight dtype" (#2784)

Summary:

Original commit changeset: 52fc46ced5a3

Original Phabricator Diff: D70126859

To fix failures like this one - f703883945

Differential Revision: D70706946

Author: pushpakrajgautam

torchrec/distributed/test_utils/test_model.py

@@ -243,7 +243,8 @@ def _validate_pooling_factor(
             global_idlist_indices.append(indices)
             global_idlist_offsets.append(offsets)
 
-        for idx, ind_range in enumerate(idscore_ind_ranges):
+        for idx in range(len(idscore_ind_ranges)):
+            ind_range = idscore_ind_ranges[idx]
             lengths_ = torch.abs(
                 torch.randn(batch_size * world_size, device=device)
                 + (

torchrec/distributed/test_utils/test_sharding.py

@@ -59,11 +59,7 @@
     ShardingPlan,
     ShardingType,
 )
-from torchrec.modules.embedding_configs import (
-    BaseEmbeddingConfig,
-    DataType,
-    EmbeddingBagConfig,
-)
+from torchrec.modules.embedding_configs import BaseEmbeddingConfig, EmbeddingBagConfig
 from torchrec.optim.keyed import CombinedOptimizer, KeyedOptimizerWrapper
 from torchrec.optim.optimizers import in_backward_optimizer_filter
 
@@ -558,7 +554,9 @@ def _custom_hook(input: List[torch.Tensor]) -> None:
             )
 
             # Compare predictions of sharded vs unsharded models.
-            if qcomms_config is not None:
+            if qcomms_config is None:
+                torch.testing.assert_close(global_pred, torch.cat(all_local_pred))
+            else:
                 # With quantized comms, we can relax constraints a bit
                 rtol = 0.003
                 if CommType.FP8 in [
@@ -570,18 +568,6 @@ def _custom_hook(input: List[torch.Tensor]) -> None:
                 torch.testing.assert_close(
                     global_pred, torch.cat(all_local_pred), rtol=rtol, atol=atol
                 )
-            elif (
-                weighted_tables is not None
-                and weighted_tables[0].data_type == DataType.FP16
-            ):  # https://www.internalfb.com/intern/diffing/?paste_number=1740410921
-                torch.testing.assert_close(
-                    global_pred,
-                    torch.cat(all_local_pred),
-                    atol=1e-4,  # relaxed atol due to FP16 in weights
-                    rtol=1e-4,  # relaxed rtol due to FP16 in weights
-                )
-            else:
-                torch.testing.assert_close(global_pred, torch.cat(all_local_pred))
 
 
 def create_device_mesh_for_2D(

torchrec/modules/embedding_modules.py

@@ -248,11 +248,7 @@ def forward(
                 res = embedding_bag(
                     input=f.values(),
                     offsets=f.offsets(),
-                    per_sample_weights=(
-                        f.weights().to(embedding_bag.weight.dtype)
-                        if self._is_weighted
-                        else None
-                    ),
+                    per_sample_weights=f.weights() if self._is_weighted else None,
                 ).float()
                 pooled_embeddings.append(res)
         return KeyedTensor(