Refactoring and moving the EBC-only logic into ‎embedding_collection_wrappers.py (#3251)

Summary:
Pull Request resolved: #3251

- Moved the EBC-only logic into a separate file `embedding_collection_wrappers.py`
- Removed linter issues

Reviewed By: spmex

Differential Revision: D79512602

fbshipit-source-id: 89357ef2383afabcbfacfbb4ff30f49154e4bea1

Author: SSYernar

examples/zch/main.py

@@ -13,7 +13,6 @@
 import torch
 
 from torchrec import EmbeddingConfig, KeyedJaggedTensor
-from torchrec.distributed.benchmark.benchmark_utils import get_inputs
 from tqdm import tqdm
 
 from .sparse_arch import SparseArch

torchrec/distributed/benchmark/benchmark_inference.py

@@ -19,15 +19,17 @@
 import torch
 
 from torchrec.distributed.benchmark.benchmark_utils import (
-    benchmark_module,
     BenchmarkResult,
     CompileMode,
     DLRM_NUM_EMBEDDINGS_PER_FEATURE,
     EMBEDDING_DIM,
-    get_tables,
     init_argparse_and_args,
     write_report,
 )
+from torchrec.distributed.benchmark.embedding_collection_wrappers import (
+    benchmark_ebc_module,
+    get_tables,
+)
 from torchrec.distributed.embedding_types import EmbeddingComputeKernel, ShardingType
 from torchrec.distributed.test_utils.infer_utils import (
     TestQuantEBCSharder,
@@ -84,7 +86,7 @@ def benchmark_qec(args: argparse.Namespace, output_dir: str) -> List[BenchmarkRe
         if not argname.startswith("_") and argname not in IGNORE_ARGNAME
     }
 
-    return benchmark_module(
+    return benchmark_ebc_module(
         module=module,
         sharder=sharder,
         sharding_types=BENCH_SHARDING_TYPES,
@@ -118,7 +120,7 @@ def benchmark_qebc(args: argparse.Namespace, output_dir: str) -> List[BenchmarkR
         if not argname.startswith("_") and argname not in IGNORE_ARGNAME
     }
 
-    return benchmark_module(
+    return benchmark_ebc_module(
         module=module,
         sharder=sharder,
         sharding_types=BENCH_SHARDING_TYPES,
@@ -153,7 +155,7 @@ def benchmark_qec_unsharded(
         if not argname.startswith("_") and argname not in IGNORE_ARGNAME
     }
 
-    return benchmark_module(
+    return benchmark_ebc_module(
         module=module,
         sharder=sharder,
         sharding_types=[],
@@ -190,7 +192,7 @@ def benchmark_qebc_unsharded(
         if not argname.startswith("_") and argname not in IGNORE_ARGNAME
     }
 
-    return benchmark_module(
+    return benchmark_ebc_module(
         module=module,
         sharder=sharder,
         sharding_types=[],

torchrec/distributed/benchmark/benchmark_train.py

@@ -17,16 +17,17 @@
 from typing import List, Optional, Tuple
 
 import torch
-
 from torchrec.distributed.benchmark.benchmark_utils import (
-    benchmark_module,
     BenchmarkResult,
     CompileMode,
-    get_tables,
     init_argparse_and_args,
     set_embedding_config,
     write_report,
 )
+from torchrec.distributed.benchmark.embedding_collection_wrappers import (
+    benchmark_ebc_module,
+    get_tables,
+)
 from torchrec.distributed.embedding_types import EmbeddingComputeKernel, ShardingType
 from torchrec.distributed.test_utils.test_model import TestEBCSharder
 from torchrec.distributed.types import DataType
@@ -106,7 +107,7 @@ def benchmark_ebc(
 
     args_kwargs["variable_batch_embeddings"] = variable_batch_embeddings
 
-    return benchmark_module(
+    return benchmark_ebc_module(
         module=module,
         sharder=sharder,
         sharding_types=BENCH_SHARDING_TYPES,