add sharding_type argument to pipeline benchmark (#2495)

Summary:
Pull Request resolved: #2495

# context
* add sharding_type argument to the pipeline benchmark
* better control of different sharding types

Reviewed By: iamzainhuda

Differential Revision: D64676132

fbshipit-source-id: 8ceeba667b6b2f8aded3dcf3d894cf8cbca31d3e

Author: TroyGarden

torchrec/distributed/batched_embedding_kernel.py

@@ -476,7 +476,8 @@ def get_sharded_optim_state(
                     momentum_local_shards: List[Shard] = []
                     optimizer_sharded_tensor_metadata: ShardedTensorMetadata
 
-                    optim_state = shard_params.optimizer_states[0][momentum_idx - 1]  # pyre-ignore[16]
+                    # pyre-ignore [16]
+                    optim_state = shard_params.optimizer_states[0][momentum_idx - 1]
                     if (
                         optim_state.nelement() == 1 and state_key != "momentum1"
                     ):  # special handling for backward compatibility, momentum1 is rowwise state for rowwise_adagrad

torchrec/distributed/train_pipeline/tests/pipeline_benchmarks.py

@@ -93,6 +93,12 @@ def _gen_pipelines(
     default=8192,
     help="Batch size.",
 )
+@click.option(
+    "--sharding_type",
+    type=ShardingType,
+    default=ShardingType.TABLE_WISE,
+    help="ShardingType.",
+)
 @click.option(
     "--pooling_factor",
     type=int,
@@ -129,6 +135,7 @@ def main(
     dim_emb: int,
     n_batches: int,
     batch_size: int,
+    sharding_type: ShardingType,
     pooling_factor: int,
     input_type: str,
     pipeline: str,
@@ -178,7 +185,7 @@ def main(
             callable=runner,
             tables=tables,
             weighted_tables=weighted_tables,
-            sharding_type=ShardingType.TABLE_WISE.value,
+            sharding_type=sharding_type.value,
             kernel_type=EmbeddingComputeKernel.FUSED.value,
             batches=batches,
             fused_params={},
@@ -190,7 +197,7 @@ def main(
         single_runner(
             tables=tables,
             weighted_tables=weighted_tables,
-            sharding_type=ShardingType.TABLE_WISE.value,
+            sharding_type=sharding_type.value,
             kernel_type=EmbeddingComputeKernel.FUSED.value,
             batches=batches,
             fused_params={},