Remove long_indices argument from model.generate function (#2874)

Summary:
Pull Request resolved: #2874

Related to stack starting with: D66728661

Removing long_indices as an input for more precise input values:

```
        use_offsets: bool = False,
        indices_dtype: torch.dtype = torch.int64,
        offsets_dtype: torch.dtype = torch.int64,
        lengths_dtype: torch.dtype = torch.int64,
```

Reviewed By: TroyGarden

Differential Revision: D69883857

fbshipit-source-id: aeb6b2590af4ddd9bef85fa629b6407bfbb6567b

Author: basilwong

torchrec/distributed/benchmark/benchmark_utils.py

@@ -372,8 +372,9 @@ def get_inputs(
                 num_float_features=0,
                 tables=tables,
                 weighted_tables=[],
-                long_indices=False,
                 tables_pooling=pooling_configs,
+                indices_dtype=torch.int32,
+                lengths_dtype=torch.int32,
             )
 
         if train:

torchrec/distributed/test_utils/infer_utils.py

@@ -187,6 +187,12 @@ def prep_inputs(
     long_indices: bool = True,
 ) -> List[ModelInput]:
     inputs = []
+    if long_indices:
+        indices_dtype = torch.int64
+        lengths_dtype = torch.int64
+    else:
+        indices_dtype = torch.int32
+        lengths_dtype = torch.int32
     for _ in range(count):
         inputs.append(
             ModelInput.generate(
@@ -195,7 +201,8 @@ def prep_inputs(
                 num_float_features=model_info.num_float_features,
                 tables=model_info.tables,
                 weighted_tables=model_info.weighted_tables,
-                long_indices=long_indices,
+                indices_dtype=indices_dtype,
+                lengths_dtype=lengths_dtype,
             )[1][0],
         )
 

torchrec/distributed/test_utils/test_model.py

@@ -98,20 +98,11 @@ def generate(
         indices_dtype: torch.dtype = torch.int64,
         offsets_dtype: torch.dtype = torch.int64,
         lengths_dtype: torch.dtype = torch.int64,
-        long_indices: bool = True,  # TODO - remove this once code base is updated to support more than long_indices spec
     ) -> Tuple["ModelInput", List["ModelInput"]]:
         """
         Returns a global (single-rank training) batch
         and a list of local (multi-rank training) batches of world_size.
         """
-        if long_indices:
-            indices_dtype = torch.int64
-            lengths_dtype = torch.int64
-            use_offsets = False
-        else:
-            indices_dtype = torch.int32
-            lengths_dtype = torch.int32
-            use_offsets = False
         batch_size_by_rank = [batch_size] * world_size
         if variable_batch_size:
             batch_size_by_rank = [

torchrec/distributed/test_utils/test_model_parallel_base.py

@@ -113,7 +113,8 @@ def _test_sharded_forward(
             dense_device=cuda_device,
             sparse_device=cuda_device,
             generate=generate,
-            long_indices=False,
+            indices_dtype=torch.int32,
+            lengths_dtype=torch.int32,
         )
         global_model = quantize_callable(global_model, **quantize_callable_kwargs)
         local_input = _inputs[0][1][default_rank].to(cuda_device)

torchrec/distributed/test_utils/test_sharding.py

@@ -120,7 +120,6 @@ def __call__(
         indices_dtype: torch.dtype = torch.int64,
         offsets_dtype: torch.dtype = torch.int64,
         lengths_dtype: torch.dtype = torch.int64,
-        long_indices: bool = True,
     ) -> Tuple["ModelInput", List["ModelInput"]]: ...
 
 
@@ -166,7 +165,6 @@ def gen_model_and_input(
     global_constant_batch: bool = False,
     num_inputs: int = 1,
     input_type: str = "kjt",  # "kjt" or "td"
-    long_indices: bool = True,
 ) -> Tuple[nn.Module, List[Tuple[ModelInput, List[ModelInput]]]]:
     torch.manual_seed(0)
     if dedup_feature_names:
@@ -229,7 +227,6 @@ def gen_model_and_input(
                     indices_dtype=indices_dtype,
                     offsets_dtype=offsets_dtype,
                     lengths_dtype=lengths_dtype,
-                    long_indices=long_indices,
                 )
             )
     else:
@@ -247,7 +244,6 @@ def gen_model_and_input(
                     indices_dtype=indices_dtype,
                     offsets_dtype=offsets_dtype,
                     lengths_dtype=lengths_dtype,
-                    long_indices=long_indices,
                 )
             )
     return (model, inputs)

torchrec/distributed/tests/test_quant_sequence_model_parallel.py

@@ -198,7 +198,8 @@ def test_quant_pred_shard(
             num_float_features=10,
             tables=self.tables,
             weighted_tables=[],
-            long_indices=False,
+            indices_dtype=torch.int32,
+            lengths_dtype=torch.int32,
         )
         local_batch = local_batch.to(device)
         sharded_quant_model(local_batch.idlist_features)