fbgemm-gpu-genai-nightly 2025.12.19__cp310-cp310-manylinux_2_28_x86_64.whl

fbgemm_gpu/split_table_batched_embeddings_ops_training_common.py

@@ -0,0 +1,146 @@
+#!/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-unsafe
+
+from typing import Optional
+
+import torch
+from torch import Tensor
+
+try:
+    try:
+        from torch.compiler import is_compiling
+
+        def is_torchdynamo_compiling() -> bool:  # type: ignore[misc]
+            # at least one test fails if we import is_compiling as a different name
+            return is_compiling()
+
+    except Exception:
+        # torch.compiler.is_compiling is not available in torch 1.10
+        from torch._dynamo import is_compiling as is_torchdynamo_compiling
+except Exception:
+
+    def is_torchdynamo_compiling() -> bool:  # type: ignore[misc]
+        return False
+
+
+# @manual=//deeplearning/fbgemm/fbgemm_gpu/codegen:split_embedding_codegen_lookup_invokers
+import fbgemm_gpu.split_embedding_codegen_lookup_invokers as invokers
+from fbgemm_gpu.split_table_batched_embeddings_ops_common import PoolingMode
+
+
+def generate_vbe_metadata(
+    offsets: Tensor,
+    batch_size_per_feature_per_rank: Optional[list[list[int]]],
+    pooling_mode: PoolingMode,
+    feature_dims_cpu: Tensor,
+    device: torch.device,
+) -> invokers.lookup_args.VBEMetadata:
+    """
+    Generate VBE metadata based on batch_size_per_feature_per_rank.
+    Metadata includes:
+        1) B_offsets - A tensor that contains batch size offsets for each
+                        feature
+        2) output_offsets_feature_rank - A tensor that contains output
+                                            offsets for each feature
+        3) B_offsets_per_rank_per_feature - A tensor that contains batch
+                                            size offsets for each feature
+                                            and rank
+        4) max_B - The maximum batch size for all features
+        5) max_B_feature_rank - The maximum batch size for all ranks and
+                                features
+        6) output_size - The output size (number of elements)
+    """
+    if batch_size_per_feature_per_rank is not None:
+        assert (
+            pooling_mode != PoolingMode.NONE
+        ), "Variable batch size TBE support is not enabled for PoolingMode.NONE"
+        # TODO: Add input check
+        zero_tensor = torch.zeros(1, device="cpu", dtype=torch.int32)
+
+        # Create B offsets
+        total_batch_size_per_feature = torch.tensor(
+            batch_size_per_feature_per_rank, dtype=torch.int32, device="cpu"
+        ).sum(dim=1)
+
+        max_B = total_batch_size_per_feature.max().item()
+        if not torch.jit.is_scripting() and is_torchdynamo_compiling():
+            torch._check_is_size(max_B)
+            torch._check(max_B < offsets.numel())
+
+        Bs = torch.concat([zero_tensor, total_batch_size_per_feature])
+        B_offsets = Bs.cumsum(dim=0).to(torch.int)
+
+        # Create output offsets
+        B_feature_rank = torch.tensor(
+            batch_size_per_feature_per_rank,
+            device="cpu",
+            dtype=torch.int64,
+        )
+        max_B_feature_rank = B_feature_rank.max().item()
+        if not torch.jit.is_scripting() and is_torchdynamo_compiling():
+            torch._check_is_size(max_B_feature_rank)
+            torch._check(max_B_feature_rank <= offsets.size(0))
+        output_sizes_feature_rank = B_feature_rank.transpose(
+            0, 1
+        ) * feature_dims_cpu.view(1, -1)
+        output_offsets_feature_rank = torch.concat(
+            [
+                zero_tensor.to(torch.int64),
+                output_sizes_feature_rank.flatten().cumsum(dim=0),
+            ]
+        )
+        output_size = output_offsets_feature_rank[-1].item()
+        if not torch.jit.is_scripting() and is_torchdynamo_compiling():
+            torch._check_is_size(output_size)
+
+        # TODO: Support INT8 output
+        # B_offsets_rank_per_feature is for rank and (b, t) mapping
+        B_offsets_rank_per_feature = (
+            torch.tensor(
+                [
+                    [0] + batch_size_per_feature
+                    for batch_size_per_feature in batch_size_per_feature_per_rank
+                ],
+                device="cpu",
+                dtype=torch.int32,
+            )
+            .cumsum(dim=1)
+            .to(torch.int)
+        )
+
+        B_offsets = B_offsets.to(device, non_blocking=True)
+        output_offsets_feature_rank = output_offsets_feature_rank.to(
+            device, non_blocking=True
+        )
+        B_offsets_rank_per_feature = B_offsets_rank_per_feature.to(
+            device, non_blocking=True
+        )
+
+        # TODO: Use int32 for B_offsets and int64 for output_offsets_feature_rank
+        vbe_metadata = invokers.lookup_args.VBEMetadata(
+            B_offsets=B_offsets,
+            output_offsets_feature_rank=output_offsets_feature_rank,
+            B_offsets_rank_per_feature=B_offsets_rank_per_feature,
+            # pyre-ignore
+            max_B=max_B,
+            # pyre-ignore
+            max_B_feature_rank=max_B_feature_rank,
+            # pyre-ignore
+            output_size=output_size,
+        )
+    else:
+        vbe_metadata = invokers.lookup_args.VBEMetadata(
+            B_offsets=None,
+            output_offsets_feature_rank=None,
+            B_offsets_rank_per_feature=None,
+            max_B=-1,
+            max_B_feature_rank=-1,
+            output_size=-1,
+        )
+    return vbe_metadata

fbgemm_gpu/ssd_split_table_batched_embeddings_ops.py

@@ -0,0 +1,26 @@
+#!/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
+# pyre-ignore-all-errors[56]
+
+import warnings
+
+from fbgemm_gpu.tbe.ssd import (  # noqa: F401
+    ASSOC,  # noqa: F401
+    SSDIntNBitTableBatchedEmbeddingBags,  # noqa: F401
+    SSDTableBatchedEmbeddingBags,  # noqa: F401
+)
+
+
+warnings.warn(  # noqa: B028
+    f"""\033[93m
+    The Python module {__name__} is now DEPRECATED and will be removed in the
+    future.  Users should import fbgemm_gpu.tbe.ssd into their scripts instead.
+    \033[0m""",
+    DeprecationWarning,
+)

fbgemm_gpu/tbe/__init__.py

@@ -0,0 +1,6 @@
+#!/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.

fbgemm_gpu/tbe/bench/__init__.py

@@ -0,0 +1,55 @@
+#!/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-unsafe
+
+import torch
+
+from .bench_config import (  # noqa F401
+    TBEBenchmarkingConfig,
+    TBEBenchmarkingConfigLoader,
+    TBEBenchmarkingHelperText,
+)
+from .bench_runs import (  # noqa F401
+    bench_warmup,
+    benchmark_cpu_requests,
+    benchmark_cpu_requests_mp,
+    benchmark_pipelined_requests,
+    benchmark_requests,
+    benchmark_requests_refer,
+    benchmark_requests_with_spec,
+    benchmark_vbe,
+)
+from .benchmark_click_interface import TbeBenchClickInterface  # 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 (  # noqa F401
+    TBEDataConfigHelperText,
+    TBEDataConfigLoader,
+)
+from .tbe_data_config_param_models import (  # noqa F401
+    BatchParams,
+    IndicesParams,
+    PoolingParams,
+)
+from .utils import fill_random_scale_bias  # noqa F401
+
+try:
+    torch.ops.load_library(
+        "//deeplearning/fbgemm/fbgemm_gpu/src/tbe/eeg:indices_estimator"
+    )
+except Exception:
+    pass
+
+#: The max number of heavy heavy hitters, as defined in
+#: fbgemm_gpu/src/tbe/eeg/indices_estimator.h
+EEG_MAX_HEAVY_HITTERS: int = 20

fbgemm_gpu/tbe/bench/bench_config.py

@@ -0,0 +1,156 @@
+#!/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
+import json
+from enum import Enum
+from typing import Any, Optional
+
+import click
+
+
+@dataclasses.dataclass(frozen=True)
+class TBEBenchmarkingConfig:
+    # Number of iterations
+    iterations: int
+    # Number of input TBE batches to generate for testing
+    num_requests: int
+    # Number of warmup iterations to run before making measurements
+    warmup_iterations: int
+    # Amount of memory to use for flushing the GPU cache after each iteration
+    flush_gpu_cache_size_mb: int
+    # If set, trace will be exported to the path specified in trace_url
+    export_trace: bool
+    # The path for exporting the trace
+    trace_url: Optional[str]
+    # If set and export_trace is true, the benchmark will upload performance data from the trace to Scuba
+    upload_perf_data: bool
+
+    @classmethod
+    # pyre-ignore [3]
+    def from_dict(cls, data: dict[str, Any]):
+        return cls(**data)
+
+    @classmethod
+    # pyre-ignore [3]
+    def from_json(cls, data: str):
+        return cls.from_dict(json.loads(data))
+
+    def dict(self) -> dict[str, Any]:
+        return dataclasses.asdict(self)
+
+    def json(self, format: bool = False) -> str:
+        return json.dumps(self.dict(), indent=(2 if format else -1), sort_keys=True)
+
+    # pyre-ignore [3]
+    def validate(self):
+        assert self.iterations > 0, "iterations must be positive"
+        assert self.num_requests > 0, "num_requests must be positive"
+        assert self.warmup_iterations >= 0, "warmup_iterations must be non-negative"
+        assert (
+            self.flush_gpu_cache_size_mb >= 0
+        ), "flush_gpu_cache_size_mb must be non-negative"
+        return self
+
+
+@dataclasses.dataclass(frozen=True)
+class TBEBenchmarkingHelperText(Enum):
+    BENCH_ITERATIONS = "Number of benchmark iterations to run"
+    BENCH_NUM_REQUESTS = "Number of input batches to generate. If the value is smaller than the number of benchmark iterations, input batches will be re-used"
+    BENCH_WARMUP_ITERATIONS = (
+        "Number of warmup iterations to run before making measurements"
+    )
+    BENCH_FLUSH_GPU_CACHE_SIZE = (
+        "Amount of memory to use for flushing the GPU cache after each iteration (MB)"
+    )
+    BENCH_EXPORT_TRACE = (
+        "If set, trace will be exported to the path specified in trace url"
+    )
+    BENCH_TRACE_URL = "The path for exporting the trace"
+    BENCH_UPLOAD_PERF_DATA = "If set and export_trace is true, the benchmark will upload performance data from the trace to Scuba"
+
+
+class TBEBenchmarkingConfigLoader:
+    @classmethod
+    # pyre-ignore [2]
+    def options(cls, func) -> click.Command:
+        options = [
+            click.option(
+                "--bench-iterations",
+                type=int,
+                default=100,
+                help=TBEBenchmarkingHelperText.BENCH_ITERATIONS.value,
+            ),
+            click.option(
+                "--bench-num-requests",
+                type=int,
+                default=-1,
+                help=TBEBenchmarkingHelperText.BENCH_NUM_REQUESTS.value,
+            ),
+            click.option(
+                "--bench-warmup-iterations",
+                type=int,
+                default=0,
+                help=TBEBenchmarkingHelperText.BENCH_WARMUP_ITERATIONS.value,
+            ),
+            click.option(
+                "--bench-flush-gpu-cache-size",
+                type=int,
+                default=0,
+                help=TBEBenchmarkingHelperText.BENCH_FLUSH_GPU_CACHE_SIZE.value,
+            ),
+            click.option(
+                "--bench-export-trace",
+                is_flag=True,
+                default=False,
+                help=TBEBenchmarkingHelperText.BENCH_EXPORT_TRACE.value,
+            ),
+            click.option(
+                "--bench-trace-url",
+                type=str,
+                required=False,
+                default="{emb_op_type}_tbe_{phase}_trace_{ospid}.json",
+                help=TBEBenchmarkingHelperText.BENCH_TRACE_URL.value,
+            ),
+            click.option(
+                "--upload-perf-data",
+                is_flag=True,
+                default=False,
+                help=TBEBenchmarkingHelperText.BENCH_UPLOAD_PERF_DATA.value,
+            ),
+        ]
+
+        for option in reversed(options):
+            func = option(func)
+        return func
+
+    @classmethod
+    def load(cls, context: click.Context) -> TBEBenchmarkingConfig:
+        params = context.params
+
+        iterations = params["bench_iterations"]
+        num_requests = params["bench_num_requests"]
+        warmup_iterations = params["bench_warmup_iterations"]
+        flush_gpu_cache_size = params["bench_flush_gpu_cache_size"]
+        export_trace = params["bench_export_trace"]
+        trace_url = params["bench_trace_url"]
+        upload_perf_data = params["upload_perf_data"]
+
+        # Default the number of TBE requests to number of iterations specified
+        num_requests = iterations if num_requests == -1 else num_requests
+
+        return TBEBenchmarkingConfig(
+            iterations,
+            num_requests,
+            warmup_iterations,
+            flush_gpu_cache_size,
+            export_trace,
+            trace_url,
+            upload_perf_data,
+        ).validate()