Cleanups for the EEG-based TBE benchmark CLI, pt 2

q10 · facebook-github-bot · commit a556812204ec · 2025-03-13T16:28:42.000-07:00
Summary: X-link: facebookresearch/FBGEMM#890 - Cleanups for the EEG-based TBE benchmark CLI, pt 2 Reviewed By: jiawenliu64 Differential Revision: D70426271
diff --git a/fbgemm_gpu/bench/tbe/tbe_training_benchmark.py b/fbgemm_gpu/bench/tbe/tbe_training_benchmark.py
@@ -32,6 +32,7 @@
 )
 from fbgemm_gpu.tbe.bench import (
     benchmark_requests,
+    EmbeddingOpsCommonConfigLoader,
     TBEBenchmarkingConfigLoader,
     TBEDataConfigLoader,
 )
@@ -50,50 +51,39 @@ def cli() -> None:
 
 
 @cli.command()
-@click.option("--weights-precision", type=SparseType, default=SparseType.FP32)
-@click.option("--cache-precision", type=SparseType, default=None)
-@click.option("--stoc", is_flag=True, default=False)
-@click.option(
-    "--managed",
-    default="device",
-    type=click.Choice(["device", "managed", "managed_caching"], case_sensitive=False),
-)
 @click.option(
     "--emb-op-type",
     default="split",
     type=click.Choice(["split", "dense", "ssd"], case_sensitive=False),
+    help="The type of the embedding op to benchmark",
+)
+@click.option(
+    "--row-wise/--no-row-wise",
+    default=True,
+    help="Whether to use row-wise adagrad optimzier or not",
 )
-@click.option("--row-wise/--no-row-wise", default=True)
-@click.option("--pooling", type=str, default="sum")
-@click.option("--weighted-num-requires-grad", type=int, default=None)
-@click.option("--bounds-check-mode", type=int, default=BoundsCheckMode.NONE.value)
-@click.option("--output-dtype", type=SparseType, default=SparseType.FP32)
 @click.option(
-    "--uvm-host-mapped",
-    is_flag=True,
-    default=False,
-    help="Use host mapped UVM buffers in SSD-TBE (malloc+cudaHostRegister)",
+    "--weighted-num-requires-grad",
+    type=int,
+    default=None,
+    help="The number of weighted tables that require gradient",
 )
 @click.option(
-    "--ssd-prefix", type=str, default="/tmp/ssd_benchmark", help="SSD directory prefix"
+    "--ssd-prefix",
+    type=str,
+    default="/tmp/ssd_benchmark",
+    help="SSD directory prefix",
 )
 @click.option("--cache-load-factor", default=0.2)
 @TBEBenchmarkingConfigLoader.options
 @TBEDataConfigLoader.options
+@EmbeddingOpsCommonConfigLoader.options
 @click.pass_context
 def device(  # noqa C901
     context: click.Context,
     emb_op_type: click.Choice,
-    weights_precision: SparseType,
-    cache_precision: Optional[SparseType],
-    stoc: bool,
-    managed: click.Choice,
     row_wise: bool,
-    pooling: str,
     weighted_num_requires_grad: Optional[int],
-    bounds_check_mode: int,
-    output_dtype: SparseType,
-    uvm_host_mapped: bool,
     cache_load_factor: float,
     # SSD params
     ssd_prefix: str,
@@ -110,6 +100,9 @@ def device(  # noqa C901
     # Load TBE data configuration from cli arguments
     tbeconfig = TBEDataConfigLoader.load(context)
 
+    # Load common embedding op configuration from cli arguments
+    embconfig = EmbeddingOpsCommonConfigLoader.load(context)
+
     # Generate feature_requires_grad
     feature_requires_grad = (
         tbeconfig.generate_feature_requires_grad(weighted_num_requires_grad)
@@ -123,22 +116,8 @@ def device(  # noqa C901
     # Determine the optimizer
     optimizer = OptimType.EXACT_ROWWISE_ADAGRAD if row_wise else OptimType.EXACT_ADAGRAD
 
-    # Determine the embedding location
-    embedding_location = str_to_embedding_location(str(managed))
-    if embedding_location is EmbeddingLocation.DEVICE and not torch.cuda.is_available():
-        embedding_location = EmbeddingLocation.HOST
-
-    # Determine the pooling mode
-    pooling_mode = str_to_pooling_mode(pooling)
-
     # Construct the common split arguments for the embedding op
-    common_split_args: Dict[str, Any] = {
-        "weights_precision": weights_precision,
-        "stochastic_rounding": stoc,
-        "output_dtype": output_dtype,
-        "pooling_mode": pooling_mode,
-        "bounds_check_mode": BoundsCheckMode(bounds_check_mode),
-        "uvm_host_mapped": uvm_host_mapped,
+    common_split_args: Dict[str, Any] = embconfig.split_args() | {
         "optimizer": optimizer,
         "learning_rate": 0.1,
         "eps": 0.1,
@@ -154,7 +133,7 @@ def device(  # noqa C901
                 )
                 for d in Ds
             ],
-            pooling_mode=pooling_mode,
+            pooling_mode=embconfig.pooling_mode,
             use_cpu=not torch.cuda.is_available(),
         )
     elif emb_op_type == "ssd":
@@ -177,7 +156,7 @@ def device(  # noqa C901
                 (
                     tbeconfig.E,
                     d,
-                    embedding_location,
+                    embconfig.embedding_location,
                     (
                         ComputeDevice.CUDA
                         if torch.cuda.is_available()
@@ -187,25 +166,27 @@ def device(  # noqa C901
                 for d in Ds
             ],
             cache_precision=(
-                weights_precision if cache_precision is None else cache_precision
+                embconfig.weights_dtype
+                if embconfig.cache_dtype is None
+                else embconfig.cache_dtype
             ),
             cache_algorithm=CacheAlgorithm.LRU,
             cache_load_factor=cache_load_factor,
             **common_split_args,
         )
     embedding_op = embedding_op.to(get_device())
 
-    if weights_precision == SparseType.INT8:
+    if embconfig.weights_dtype == SparseType.INT8:
         # pyre-fixme[29]: `Union[(self: DenseTableBatchedEmbeddingBagsCodegen,
         #  min_val: float, max_val: float) -> None, (self:
         #  SplitTableBatchedEmbeddingBagsCodegen, min_val: float, max_val: float) ->
         #  None, Tensor, Module]` is not a function.
         embedding_op.init_embedding_weights_uniform(-0.0003, 0.0003)
 
     nparams = sum(d * tbeconfig.E for d in Ds)
-    param_size_multiplier = weights_precision.bit_rate() / 8.0
-    output_size_multiplier = output_dtype.bit_rate() / 8.0
-    if pooling_mode.do_pooling():
+    param_size_multiplier = embconfig.weights_dtype.bit_rate() / 8.0
+    output_size_multiplier = embconfig.output_dtype.bit_rate() / 8.0
+    if embconfig.pooling_mode.do_pooling():
         read_write_bytes = (
             output_size_multiplier * tbeconfig.batch_params.B * sum(Ds)
             + param_size_multiplier
@@ -225,7 +206,7 @@ def device(  # noqa C901
             * tbeconfig.pooling_params.L
         )
 
-    logging.info(f"Managed option: {managed}")
+    logging.info(f"Managed option: {embconfig.embedding_location}")
     logging.info(
         f"Embedding parameters: {nparams / 1.0e9: .2f} GParam, "
         f"{nparams * param_size_multiplier / 1.0e9: .2f} GB"
@@ -274,11 +255,11 @@ def _context_factory(on_trace_ready: Callable[[profile], None]):
         f"T: {time_per_iter * 1.0e6:.0f}us"
     )
 
-    if output_dtype == SparseType.INT8:
+    if embconfig.output_dtype == SparseType.INT8:
         # backward bench not representative
         return
 
-    if pooling_mode.do_pooling():
+    if embconfig.pooling_mode.do_pooling():
         grad_output = torch.randn(tbeconfig.batch_params.B, sum(Ds)).to(get_device())
     else:
         grad_output = torch.randn(
diff --git a/fbgemm_gpu/fbgemm_gpu/split_table_batched_embeddings_ops_common.py b/fbgemm_gpu/fbgemm_gpu/split_table_batched_embeddings_ops_common.py
@@ -33,6 +33,21 @@ class EmbeddingLocation(enum.IntEnum):
     HOST = 3
     MTIA = 4
 
+    @classmethod
+    # pyre-ignore[3]
+    def from_str(cls, key: str):
+        lookup = {
+            "device": EmbeddingLocation.DEVICE,
+            "managed": EmbeddingLocation.MANAGED,
+            "managed_caching": EmbeddingLocation.MANAGED_CACHING,
+            "host": EmbeddingLocation.HOST,
+            "mtia": EmbeddingLocation.MTIA,
+        }
+        if key in lookup:
+            return lookup[key]
+        else:
+            raise ValueError(f"Cannot parse value into EmbeddingLocation: {key}")
+
 
 def str_to_embedding_location(key: str) -> EmbeddingLocation:
     lookup = {
diff --git a/fbgemm_gpu/fbgemm_gpu/tbe/bench/__init__.py b/fbgemm_gpu/fbgemm_gpu/tbe/bench/__init__.py
@@ -19,12 +19,17 @@
     benchmark_requests_refer,
     benchmark_vbe,
 )
-from .config import TBEDataConfig  # noqa F401
-from .config_loader import TBEDataConfigLoader  # noqa F401
-from .config_param_models import BatchParams, IndicesParams, PoolingParams  # noqa F401
+from .embedding_ops_common_config import EmbeddingOpsCommonConfigLoader  # noqa F401
 from .eval_compression import (  # noqa F401
     benchmark_eval_compression,
     EvalCompressionBenchmarkOutput,
 )
 from .reporter import BenchmarkReporter  # noqa F401
+from .tbe_data_config import TBEDataConfig  # noqa F401
+from .tbe_data_config_loader import TBEDataConfigLoader  # noqa F401
+from .tbe_data_config_param_models import (  # noqa F401
+    BatchParams,
+    IndicesParams,
+    PoolingParams,
+)
 from .utils import fill_random_scale_bias  # noqa F401
diff --git a/fbgemm_gpu/fbgemm_gpu/tbe/bench/embedding_ops_common_config.py b/fbgemm_gpu/fbgemm_gpu/tbe/bench/embedding_ops_common_config.py
@@ -0,0 +1,144 @@
+#!/usr/bin/env python3
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+# pyre-strict
+
+import dataclasses
+from typing import Any, Dict, Optional
+
+import click
+import torch
+from fbgemm_gpu.split_embedding_configs import SparseType
+from fbgemm_gpu.split_table_batched_embeddings_ops_common import (
+    BoundsCheckMode,
+    EmbeddingLocation,
+    PoolingMode,
+    str_to_embedding_location,
+    str_to_pooling_mode,
+)
+
+
+@dataclasses.dataclass(frozen=True)
+class EmbeddingOpsCommonConfig:
+    weights_dtype: SparseType
+    cache_dtype: Optional[SparseType]
+    output_dtype: SparseType
+    stochastic_rounding: bool
+    pooling_mode: PoolingMode
+    uvm_host_mapped: bool
+    embedding_location: EmbeddingLocation
+    bounds_check_mode: BoundsCheckMode
+
+    # pyre-ignore [3]
+    def validate(self):
+        return self
+
+    def split_args(self) -> Dict[str, Any]:
+        return {
+            "weights_precision": self.weights_dtype,
+            "stochastic_rounding": self.stochastic_rounding,
+            "output_dtype": self.output_dtype,
+            "pooling_mode": self.pooling_mode,
+            "bounds_check_mode": self.bounds_check_mode,
+            "uvm_host_mapped": self.uvm_host_mapped,
+        }
+
+
+class EmbeddingOpsCommonConfigLoader:
+    @classmethod
+    # pyre-ignore [2]
+    def options(cls, func) -> click.Command:
+        options = [
+            click.option(
+                "--emb-weights-dtype",
+                type=SparseType,
+                default=SparseType.FP32,
+                help="Precision of the embedding weights",
+            ),
+            click.option(
+                "--emb-cache-dtype",
+                type=SparseType,
+                default=None,
+                help="Precision of the embedding cache",
+            ),
+            click.option(
+                "--emb-output-dtype",
+                type=SparseType,
+                default=SparseType.FP32,
+                help="Precision of the embedding output",
+            ),
+            click.option(
+                "--emb-stochastic-rounding",
+                is_flag=True,
+                default=False,
+                help="Enable stochastic rounding when performing quantization",
+            ),
+            click.option(
+                "--emb-pooling-mode",
+                type=click.Choice(["sum", "mean", "none"], case_sensitive=False),
+                default="sum",
+                help="Pooling operation to perform",
+            ),
+            click.option(
+                "--emb-uvm-host-mapped",
+                is_flag=True,
+                default=False,
+                help="Use host-mapped UVM buffers",
+            ),
+            click.option(
+                "--emb-location",
+                default="device",
+                type=click.Choice(
+                    ["device", "managed", "managed_caching"], case_sensitive=False
+                ),
+                help="Memory location of the embeddings",
+            ),
+            click.option(
+                "--emb-bounds-check",
+                type=int,
+                default=BoundsCheckMode.WARNING.value,
+                help="Bounds check mode."
+                f"Available modes: FATAL={BoundsCheckMode.FATAL.value}, "
+                f"WARNING={BoundsCheckMode.WARNING.value}, "
+                f"IGNORE={BoundsCheckMode.IGNORE.value}, "
+                f"NONE={BoundsCheckMode.NONE.value}",
+            ),
+        ]
+
+        for option in reversed(options):
+            func = option(func)
+        return func
+
+    @classmethod
+    def load(cls, context: click.Context) -> EmbeddingOpsCommonConfig:
+        params = context.params
+
+        weights_dtype = params["emb_weights_dtype"]
+        cache_dtype = params["emb_cache_dtype"]
+        output_dtype = params["emb_output_dtype"]
+        stochastic_rounding = params["emb_stochastic_rounding"]
+        pooling_mode = str_to_pooling_mode(str(params["emb_pooling_mode"]))
+        uvm_host_mapped = params["emb_uvm_host_mapped"]
+        bounds_check_mode = BoundsCheckMode(params["emb_bounds_check"])
+
+        embedding_location = str_to_embedding_location(str(params["emb_location"]))
+        if (
+            embedding_location is EmbeddingLocation.DEVICE
+            and not torch.cuda.is_available()
+        ):
+            embedding_location = EmbeddingLocation.HOST
+
+        return EmbeddingOpsCommonConfig(
+            weights_dtype,
+            cache_dtype,
+            output_dtype,
+            stochastic_rounding,
+            pooling_mode,
+            uvm_host_mapped,
+            embedding_location,
+            bounds_check_mode,
+        ).validate()
diff --git a/fbgemm_gpu/fbgemm_gpu/tbe/bench/tbe_data_config.py b/fbgemm_gpu/fbgemm_gpu/tbe/bench/tbe_data_config.py
@@ -23,7 +23,7 @@
     TBERequest,
 )
 
-from .config_param_models import BatchParams, IndicesParams, PoolingParams
+from .tbe_data_config_param_models import BatchParams, IndicesParams, PoolingParams
 
 try:
     torch.ops.load_library(
diff --git a/fbgemm_gpu/fbgemm_gpu/tbe/bench/tbe_data_config_loader.py b/fbgemm_gpu/fbgemm_gpu/tbe/bench/tbe_data_config_loader.py
diff --git a/fbgemm_gpu/fbgemm_gpu/tbe/bench/tbe_data_config_param_models.py b/fbgemm_gpu/fbgemm_gpu/tbe/bench/tbe_data_config_param_models.py

Original file line number	Diff line number	Diff line change
`@@ -23,7 +23,7 @@`
`23`	`23`	`TBERequest,`
`24`	`24`	`)`
`25`	`25`
`26`		`-from .config_param_models import BatchParams, IndicesParams, PoolingParams`
	`26`	`+from .tbe_data_config_param_models import BatchParams, IndicesParams, PoolingParams`
`27`	`27`
`28`	`28`	`try:`
`29`	`29`	`torch.ops.load_library(`