add sharding_type argument to pipeline benchmark

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

Differential Revision: D64676132

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={},