Update the generate_variable_batch_input function to enable configuring index/offset/length type (#2765)

Summary:
Pull Request resolved: #2765

# Diff Specific Changes

Adding new parameters for the index/offset/length type, as well as modifying the [generate_variable_batch_input](https://www.internalfb.com/code/fbsource/[4217c068fa966d569d2042a7263cefe1a06dc87a]/fbcode/torchrec/distributed/test_utils/test_model.py?lines=589) function + helper functions to use these parameters.

# Context Doc:
 https://docs.google.com/document/d/1YVfxsafqXkxAAdRyXbjmSH4AEz3-6DBiTGjs1rT8ZHQ/edit?usp=sharing

Updating the TorchRec unit test suite to cover int32 and int64 indices/offets support.

# Summary

Specifically for the [test_model_parallel](https://www.internalfb.com/code/fbsource/[3505ccb75a649a7d21218bcda126d1e8392afc5a]/fbcode/torchrec/distributed/test_utils/test_model_parallel.py?lines=34) suite that I am looking at the change appears to be fairly straightforward.

1.The [ModelParallelTestShared](https://www.internalfb.com/code/fbsource/fbcode/torchrec/distributed/test_utils/test_model_parallel.py?lines=34) class defines a [test suite python library](https://www.internalfb.com/code/fbsource/[cbd0bd0020a7afbec4922d8abc0d88b7d45cba56]/fbcode/torchrec/distributed/test_utils/TARGETS?lines=65-69) referenced by multiple unit tests in the TorchRec codebase including [test_model_parallel_nccl](https://www.internalfb.com/code/fbsource/[cbd0bd0020a7afbec4922d8abc0d88b7d45cba56]/fbcode/torchrec/distributed/tests/TARGETS?lines=85-100) in which we are particularly interested in for this particular case. The method all of the unit tests in this class use is [`_test_sharding`](https://www.internalfb.com/code/fbsource/[fa9508a29b62ce57681ee73cd6d4cac56f153a58]/fbcode/torchrec/distributed/test_utils/test_model_parallel.py?lines=132). Within the `_test_sharding` function, the "callable" argument input to the [`_run_multi_process_test`](https://www.internalfb.com/code/symbol/fbsource/py/fbcode/caffe2.torch.fb.hpc.tests.sparse_data_dist_test.SparseDataDistTest._run_multi_process_test) function is [`sharding_single_rank_test`](https://www.internalfb.com/code/fbsource/[fa9508a29b62ce57681ee73cd6d4cac56f153a58]/fbcode/torchrec/distributed/test_utils/test_sharding.py?lines=296) which shows us how the input data/model is generated. Additional arguments will need to be added to both the [`_test_sharding`](https://www.internalfb.com/code/fbsource/[fa9508a29b62ce57681ee73cd6d4cac56f153a58]/fbcode/torchrec/distributed/test_utils/test_model_parallel.py?lines=132) and [`_run_multi_process_test`](https://www.internalfb.com/code/symbol/fbsource/py/fbcode/caffe2.torch.fb.hpc.tests.sparse_data_dist_test.SparseDataDistTest._run_multi_process_test) functions.

2.The [`sharding_single_rank_test`](https://www.internalfb.com/code/fbsource/[fa9508a29b62ce57681ee73cd6d4cac56f153a58]/fbcode/torchrec/distributed/test_utils/test_sharding.py?lines=296) function is where we define additional kwargs. This function leverages the [`gen_model_and_input`](https://www.internalfb.com/code/fbsource/[f7e6a3281d924b465e0e90ff079aa9df83ae9530]/fbcode/torchrec/distributed/test_utils/test_sharding.py?lines=131) to define the test model and more importantly for our purposes the input tables.

```
    generate=(cast(VariableBatchModelInputCallable, ModelInput.generate_variable_batch_input) if variable_batch_per_feature else ModelInput.generate),
```

3.The [ModelInput](https://www.internalfb.com/code/fbsource/[4217c068fa966d569d2042a7263cefe1a06dc87a]/fbcode/torchrec/distributed/test_utils/test_model.py?lines=48) class' [`generate`](https://www.internalfb.com/code/fbsource/[4217c068fa966d569d2042a7263cefe1a06dc87a]/fbcode/torchrec/distributed/test_utils/test_model.py?lines=55) and [`generate_variable_batch_input`](https://www.internalfb.com/code/fbsource/[4217c068fa966d569d2042a7263cefe1a06dc87a]/fbcode/torchrec/distributed/test_utils/test_model.py?lines=589) methods are used to generate the input tensors used in the unit tests. All we need to do is add new arguments that enable configuring the index/offset type of the tables.

# Diff stack change summary:

a. Update the generate_variable_batch_input to enable configuring index/offset/length type
b. Update the generate to enable configuring index/offset/length type
c. Update Model Input Callable Protocol to Enable Configuring index/offset/length type
d. test_model_parallel: new test for different table index types
e. Deprecate long_indices argument for torch.dtype arguments

Reviewed By: TroyGarden

Differential Revision: D70054939

fbshipit-source-id: 9bee9478d21cfcba694ae866437e9d1ee4910c75

Author: basilwong

torchrec/distributed/test_utils/test_model.py

@@ -396,29 +396,42 @@ def _generate_variable_batch_local_features(
         strides_per_rank_per_feature: Dict[int, Dict[str, int]],
         inverse_indices_per_rank_per_feature: Dict[int, Dict[str, torch.Tensor]],
         weights_per_rank_per_feature: Optional[Dict[int, Dict[str, torch.Tensor]]],
+        use_offsets: bool,
+        indices_dtype: torch.dtype,
+        offsets_dtype: torch.dtype,
+        lengths_dtype: torch.dtype,
     ) -> List[KeyedJaggedTensor]:
         local_kjts = []
         keys = list(feature_num_embeddings.keys())
+
         for rank in range(world_size):
             lengths_per_rank_per_feature[rank] = {}
             values_per_rank_per_feature[rank] = {}
             strides_per_rank_per_feature[rank] = {}
             inverse_indices_per_rank_per_feature[rank] = {}
+
             if weights_per_rank_per_feature is not None:
                 weights_per_rank_per_feature[rank] = {}
 
             for key, num_embeddings in feature_num_embeddings.items():
                 batch_size = random.randint(1, average_batch_size * dedup_factor - 1)
-                lengths = torch.randint(low=0, high=5, size=(batch_size,))
+                lengths = torch.randint(
+                    low=0, high=5, size=(batch_size,), dtype=lengths_dtype
+                )
                 lengths_per_rank_per_feature[rank][key] = lengths
                 lengths_sum = sum(lengths.tolist())
-                values = torch.randint(0, num_embeddings, (lengths_sum,))
+                values = torch.randint(
+                    0, num_embeddings, (lengths_sum,), dtype=indices_dtype
+                )
                 values_per_rank_per_feature[rank][key] = values
                 if weights_per_rank_per_feature is not None:
                     weights_per_rank_per_feature[rank][key] = torch.rand(lengths_sum)
                 strides_per_rank_per_feature[rank][key] = batch_size
                 inverse_indices_per_rank_per_feature[rank][key] = torch.randint(
-                    0, batch_size, (dedup_factor * average_batch_size,)
+                    0,
+                    batch_size,
+                    (dedup_factor * average_batch_size,),
+                    dtype=indices_dtype,
                 )
 
             values = torch.cat(list(values_per_rank_per_feature[rank].values()))
@@ -428,23 +441,40 @@ def _generate_variable_batch_local_features(
                 if weights_per_rank_per_feature is not None
                 else None
             )
-            stride_per_key_per_rank = [
-                [stride] for stride in strides_per_rank_per_feature[rank].values()
-            ]
-            inverse_indices = (
-                keys,
-                torch.stack(list(inverse_indices_per_rank_per_feature[rank].values())),
-            )
-            local_kjts.append(
-                KeyedJaggedTensor(
-                    keys=keys,
-                    values=values,
-                    lengths=lengths,
-                    weights=weights,
-                    stride_per_key_per_rank=stride_per_key_per_rank,
-                    inverse_indices=inverse_indices,
+
+            if use_offsets:
+                offsets = torch.cat(
+                    [torch.tensor([0], dtype=offsets_dtype), lengths.cumsum(0)]
                 )
-            )
+                local_kjts.append(
+                    KeyedJaggedTensor(
+                        keys=keys,
+                        values=values,
+                        offsets=offsets,
+                        weights=weights,
+                    )
+                )
+            else:
+                stride_per_key_per_rank = [
+                    [stride] for stride in strides_per_rank_per_feature[rank].values()
+                ]
+                inverse_indices = (
+                    keys,
+                    torch.stack(
+                        list(inverse_indices_per_rank_per_feature[rank].values())
+                    ),
+                )
+                local_kjts.append(
+                    KeyedJaggedTensor(
+                        keys=keys,
+                        values=values,
+                        lengths=lengths,
+                        weights=weights,
+                        stride_per_key_per_rank=stride_per_key_per_rank,
+                        inverse_indices=inverse_indices,
+                    )
+                )
+
         return local_kjts
 
     @staticmethod
@@ -457,6 +487,10 @@ def _generate_variable_batch_global_features(
         strides_per_rank_per_feature: Dict[int, Dict[str, int]],
         inverse_indices_per_rank_per_feature: Dict[int, Dict[str, torch.Tensor]],
         weights_per_rank_per_feature: Optional[Dict[int, Dict[str, torch.Tensor]]],
+        use_offsets: bool,
+        indices_dtype: torch.dtype,
+        offsets_dtype: torch.dtype,
+        lengths_dtype: torch.dtype,
     ) -> KeyedJaggedTensor:
         global_values = []
         global_lengths = []
@@ -476,31 +510,41 @@ def _generate_variable_batch_global_features(
                 inverse_indices_per_feature_per_rank.append(
                     inverse_indices_per_rank_per_feature[rank][key]
                 )
+
             global_stride_per_key_per_rank.append([sum_stride])
 
         inverse_indices_list: List[torch.Tensor] = []
+
         for key in keys:
             accum_batch_size = 0
             inverse_indices = []
+
             for rank in range(world_size):
                 inverse_indices.append(
                     inverse_indices_per_rank_per_feature[rank][key] + accum_batch_size
                 )
                 accum_batch_size += strides_per_rank_per_feature[rank][key]
+
             inverse_indices_list.append(torch.cat(inverse_indices))
+
         global_inverse_indices = (keys, torch.stack(inverse_indices_list))
 
         if global_constant_batch:
             global_offsets = []
+
             for length in global_lengths:
                 global_offsets.append(_to_offsets(length))
+
             reindexed_lengths = []
+
             for length, indices in zip(
                 global_lengths, inverse_indices_per_feature_per_rank
             ):
                 reindexed_lengths.append(torch.index_select(length, 0, indices))
+
             lengths = torch.cat(reindexed_lengths)
             reindexed_values, reindexed_weights = [], []
+
             for i, (values, offsets, indices) in enumerate(
                 zip(global_values, global_offsets, inverse_indices_per_feature_per_rank)
             ):
@@ -510,25 +554,40 @@ def _generate_variable_batch_global_features(
                         reindexed_weights.append(
                             global_weights[i][offsets[idx] : offsets[idx + 1]]
                         )
+
             values = torch.cat(reindexed_values)
             weights = (
                 torch.cat(reindexed_weights) if global_weights is not None else None
             )
             global_stride_per_key_per_rank = None
             global_inverse_indices = None
+
         else:
             values = torch.cat(global_values)
             lengths = torch.cat(global_lengths)
             weights = torch.cat(global_weights) if global_weights is not None else None
 
-        return KeyedJaggedTensor(
-            keys=keys,
-            values=values,
-            lengths=lengths,
-            weights=weights,
-            stride_per_key_per_rank=global_stride_per_key_per_rank,
-            inverse_indices=global_inverse_indices,
-        )
+        if use_offsets:
+            offsets = torch.cat(
+                [torch.tensor([0], dtype=offsets_dtype), lengths.cumsum(0)]
+            )
+            return KeyedJaggedTensor(
+                keys=keys,
+                values=values,
+                offsets=offsets,
+                weights=weights,
+                stride_per_key_per_rank=global_stride_per_key_per_rank,
+                inverse_indices=global_inverse_indices,
+            )
+        else:
+            return KeyedJaggedTensor(
+                keys=keys,
+                values=values,
+                lengths=lengths,
+                weights=weights,
+                stride_per_key_per_rank=global_stride_per_key_per_rank,
+                inverse_indices=global_inverse_indices,
+            )
 
     @staticmethod
     def _generate_variable_batch_features(
@@ -539,11 +598,17 @@ def _generate_variable_batch_features(
         world_size: int,
         dedup_factor: int,
         global_constant_batch: bool,
+        use_offsets: bool,
+        indices_dtype: torch.dtype,
+        offsets_dtype: torch.dtype,
+        lengths_dtype: torch.dtype,
     ) -> Tuple[KeyedJaggedTensor, List[KeyedJaggedTensor]]:
         is_weighted = (
             True if tables and getattr(tables[0], "is_weighted", False) else False
         )
+
         feature_num_embeddings = {}
+
         for table in tables:
             for feature_name in table.feature_names:
                 feature_num_embeddings[feature_name] = (
@@ -553,33 +618,42 @@ def _generate_variable_batch_features(
                 )
 
         local_kjts = []
+
         values_per_rank_per_feature = {}
         lengths_per_rank_per_feature = {}
         strides_per_rank_per_feature = {}
         inverse_indices_per_rank_per_feature = {}
         weights_per_rank_per_feature = {} if is_weighted else None
 
         local_kjts = ModelInput._generate_variable_batch_local_features(
-            feature_num_embeddings,
-            average_batch_size,
-            world_size,
-            dedup_factor,
-            values_per_rank_per_feature,
-            lengths_per_rank_per_feature,
-            strides_per_rank_per_feature,
-            inverse_indices_per_rank_per_feature,
-            weights_per_rank_per_feature,
+            feature_num_embeddings=feature_num_embeddings,
+            average_batch_size=average_batch_size,
+            world_size=world_size,
+            dedup_factor=dedup_factor,
+            values_per_rank_per_feature=values_per_rank_per_feature,
+            lengths_per_rank_per_feature=lengths_per_rank_per_feature,
+            strides_per_rank_per_feature=strides_per_rank_per_feature,
+            inverse_indices_per_rank_per_feature=inverse_indices_per_rank_per_feature,
+            weights_per_rank_per_feature=weights_per_rank_per_feature,
+            use_offsets=use_offsets,
+            indices_dtype=indices_dtype,
+            offsets_dtype=offsets_dtype,
+            lengths_dtype=lengths_dtype,
         )
 
         global_kjt = ModelInput._generate_variable_batch_global_features(
-            list(feature_num_embeddings.keys()),
-            world_size,
-            global_constant_batch,
-            values_per_rank_per_feature,
-            lengths_per_rank_per_feature,
-            strides_per_rank_per_feature,
-            inverse_indices_per_rank_per_feature,
-            weights_per_rank_per_feature,
+            keys=list(feature_num_embeddings.keys()),
+            world_size=world_size,
+            global_constant_batch=global_constant_batch,
+            values_per_rank_per_feature=values_per_rank_per_feature,
+            lengths_per_rank_per_feature=lengths_per_rank_per_feature,
+            strides_per_rank_per_feature=strides_per_rank_per_feature,
+            inverse_indices_per_rank_per_feature=inverse_indices_per_rank_per_feature,
+            weights_per_rank_per_feature=weights_per_rank_per_feature,
+            use_offsets=use_offsets,
+            indices_dtype=indices_dtype,
+            offsets_dtype=offsets_dtype,
+            lengths_dtype=lengths_dtype,
         )
 
         return (global_kjt, local_kjts)
@@ -601,30 +675,51 @@ def generate_variable_batch_input(
         ] = None,
         pooling_avg: int = 10,
         global_constant_batch: bool = False,
+        use_offsets: bool = False,
+        indices_dtype: torch.dtype = torch.int64,
+        offsets_dtype: torch.dtype = torch.int64,
+        lengths_dtype: torch.dtype = torch.int64,
     ) -> Tuple["ModelInput", List["ModelInput"]]:
         torch.manual_seed(100)
         random.seed(100)
         dedup_factor = 2
+
         global_kjt, local_kjts = ModelInput._generate_variable_batch_features(
-            tables, average_batch_size, world_size, dedup_factor, global_constant_batch
+            tables=tables,
+            average_batch_size=average_batch_size,
+            world_size=world_size,
+            dedup_factor=dedup_factor,
+            global_constant_batch=global_constant_batch,
+            use_offsets=use_offsets,
+            indices_dtype=indices_dtype,
+            offsets_dtype=offsets_dtype,
+            lengths_dtype=lengths_dtype,
         )
+
         if weighted_tables:
             global_score_kjt, local_score_kjts = (
                 ModelInput._generate_variable_batch_features(
-                    weighted_tables,
-                    average_batch_size,
-                    world_size,
-                    dedup_factor,
-                    global_constant_batch,
+                    tables=weighted_tables,
+                    average_batch_size=average_batch_size,
+                    world_size=world_size,
+                    dedup_factor=dedup_factor,
+                    global_constant_batch=global_constant_batch,
+                    use_offsets=use_offsets,
+                    indices_dtype=indices_dtype,
+                    offsets_dtype=offsets_dtype,
+                    lengths_dtype=lengths_dtype,
                 )
             )
         else:
             global_score_kjt, local_score_kjts = None, []
+
         global_float = torch.rand(
             (dedup_factor * average_batch_size * world_size, num_float_features)
         )
+
         local_model_input = []
         label_per_rank = []
+
         for rank in range(world_size):
             label_per_rank.append(torch.rand(dedup_factor * average_batch_size))
             local_float = global_float[
@@ -644,12 +739,14 @@ def generate_variable_batch_input(
                     float_features=local_float,
                 ),
             )
+
         global_model_input = ModelInput(
             idlist_features=global_kjt,
             idscore_features=global_score_kjt,
             label=torch.cat(label_per_rank),
             float_features=global_float,
         )
+
         return (global_model_input, local_model_input)
 
     def to(self, device: torch.device, non_blocking: bool = False) -> "ModelInput":