refactor pipeline benchmark (#3432)

Summary:
Pull Request resolved: #3432

# context
* refactor and clean up the benchmark_uitls and benchmark_pipeline_utils files
* run-related functions and wrappers are now in the benchmark_base.py, such as `_run_benchmark_core`, `multi_process_benchmark`, and `BenchmarkResult`.
* components generation related functions are now in the benchmark_uilts.py, such as `generate_pipeline`, `generate_planner`, and `generate_sharded_model_and_optimizer`.
* also remove some redundent functions such as `_init_module_and_run_benchmark` and `benchmark_module`, which are available in embedding_collection_wrappers.py (D79512602)

Reviewed By: spmex

Differential Revision: D83883554

fbshipit-source-id: b98759582041477d36200a8a540740bdb6015bc7

Author: TroyGarden

torchrec/distributed/benchmark/benchmark_base.py

@@ -18,7 +18,7 @@
 
 import torch
 
-from torchrec.distributed.benchmark.benchmark_utils import (
+from torchrec.distributed.benchmark.benchmark_base import (
     BenchmarkResult,
     CompileMode,
     DLRM_NUM_EMBEDDINGS_PER_FEATURE,

torchrec/distributed/benchmark/benchmark_inference.py

@@ -13,7 +13,7 @@
 
 import click
 import torch
-from torchrec.distributed.benchmark.benchmark_utils import benchmark_func
+from torchrec.distributed.benchmark.benchmark_base import benchmark_func
 from torchrec.distributed.embedding import EmbeddingCollectionContext
 from torchrec.distributed.embedding_sharding import _set_sharding_context_post_a2a
 from torchrec.distributed.sharding.sequence_sharding import SequenceShardingContext

torchrec/distributed/benchmark/benchmark_pipeline_utils.py

@@ -26,7 +26,7 @@
     ComputeDevice,
     SplitTableBatchedEmbeddingBagsCodegen,
 )
-from torchrec.distributed.benchmark.benchmark_utils import benchmark_func
+from torchrec.distributed.benchmark.benchmark_base import benchmark_func
 from torchrec.distributed.test_utils.test_model import ModelInput
 from torchrec.modules.embedding_configs import EmbeddingBagConfig
 from torchrec.sparse.jagged_tensor import KeyedJaggedTensor

torchrec/distributed/benchmark/benchmark_set_sharding_context_post_a2a.py

@@ -17,7 +17,7 @@
 from typing import List, Optional, Tuple
 
 import torch
-from torchrec.distributed.benchmark.benchmark_utils import (
+from torchrec.distributed.benchmark.benchmark_base import (
     BenchmarkResult,
     CompileMode,
     init_argparse_and_args,