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

fbgemm_gpu/tbe/bench/tbe_data_config_loader.py

@@ -0,0 +1,289 @@
+#!/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 enum import Enum
+
+import click
+import torch
+import yaml
+
+from fbgemm_gpu.tbe.bench.tbe_data_config import (
+    BatchParams,
+    IndicesParams,
+    PoolingParams,
+    TBEDataConfig,
+)
+
+
+@dataclasses.dataclass(frozen=True)
+class TBEDataConfigHelperText(Enum):
+    # Config File
+    TBE_CONFIG = "TBE data configuration filepath.  If provided, all other `--tbe-*` options are ignored."
+
+    # Table Parameters
+    TBE_NUM_TABLES = "Number of tables (T)"
+    TBE_NUM_EMBEDDINGS = "Number of embeddings (E)"
+    TBE_EMBEDDING_DIM = "Embedding dimensions (D)"
+    TBE_MIXED_DIM = "Use mixed dimensions"
+    TBE_WEIGHTED = "Flag to indicate if the table is weighted"
+
+    # Batch Parameters
+    TBE_BATCH_SIZE = "Batch size (B)"
+    TBE_BATCH_VBE_SIGMA = "Standard deviation of B for VBE"
+    TBE_BATCH_VBE_DIST = "VBE distribution (choices: 'uniform', 'normal')"
+    TBE_BATCH_VBE_RANKS = "Number of ranks for VBE"
+
+    # Indices Parameters
+    TBE_INDICES_HITTERS = "Heavy hitters for indices (comma-delimited list of floats)"
+    TBE_INDICES_ZIPF = "Zipf distribution parameters for indices generation (q, s)"
+    TBE_INDICES_DTYPE = "The dtype of the table indices (choices: '32', '64')"
+    TBE_OFFSETS_DTYPE = "The dtype of the table indices (choices: '32', '64')"
+
+    # Pooling Parameters
+    TBE_POOLING_SIZE = "Bag size / pooling factor (L)"
+    TBE_POOLING_VL_SIGMA = "Standard deviation of B for VBE"
+    TBE_POOLING_VL_DIST = "VBE distribution (choices: 'uniform', 'normal')"
+
+
+class TBEDataConfigLoader:
+    @classmethod
+    # pyre-ignore [2]
+    def options(cls, func) -> click.Command:
+        options = [
+            # Config File
+            click.option(
+                "--tbe-config",
+                type=str,
+                required=False,
+                help=TBEDataConfigHelperText.TBE_CONFIG.value,
+            ),
+            # Table Parameters
+            click.option(
+                "--tbe-num-tables",
+                type=int,
+                default=32,
+                help=TBEDataConfigHelperText.TBE_NUM_TABLES.value,
+            ),
+            click.option(
+                "--tbe-num-embeddings",
+                type=int,
+                default=int(1e5),
+                help=TBEDataConfigHelperText.TBE_NUM_EMBEDDINGS.value,
+            ),
+            click.option(
+                "--tbe-num-embeddings-list",
+                type=str,
+                required=False,
+                default=None,
+                help="Comma-separated list of number of embeddings (Es)",
+            ),
+            click.option(
+                "--tbe-embedding-dim",
+                type=int,
+                default=128,
+                help=TBEDataConfigHelperText.TBE_EMBEDDING_DIM.value,
+            ),
+            click.option(
+                "--tbe-embedding-dim-list",
+                type=str,
+                required=False,
+                default=None,
+                help="Comma-separated list of number of Embedding dimensions (D)",
+            ),
+            click.option(
+                "--tbe-mixed-dim",
+                is_flag=True,
+                default=False,
+                help=TBEDataConfigHelperText.TBE_MIXED_DIM.value,
+            ),
+            click.option(
+                "--tbe-weighted",
+                is_flag=True,
+                default=False,
+                help=TBEDataConfigHelperText.TBE_WEIGHTED.value,
+            ),
+            click.option(
+                "--tbe-max-indices",
+                type=int,
+                required=False,
+                default=None,
+                help="(Optional) Maximum number of indices, will be calculated if not provided",
+            ),
+            # Batch Parameters
+            click.option(
+                "--tbe-batch-size",
+                type=int,
+                default=512,
+                help=TBEDataConfigHelperText.TBE_BATCH_SIZE.value,
+            ),
+            click.option(
+                "--tbe-batch-sizes-list",
+                type=str,
+                required=False,
+                default=None,
+                help="List Batch sizes per feature (Bs)",
+            ),
+            click.option(
+                "--tbe-batch-vbe-sigma",
+                type=int,
+                required=False,
+                help=TBEDataConfigHelperText.TBE_BATCH_VBE_SIGMA.value,
+            ),
+            click.option(
+                "--tbe-batch-vbe-dist",
+                type=click.Choice(["uniform", "normal"]),
+                required=False,
+                help=TBEDataConfigHelperText.TBE_BATCH_VBE_DIST.value,
+            ),
+            click.option(
+                "--tbe-batch-vbe-ranks",
+                type=int,
+                required=False,
+                help=TBEDataConfigHelperText.TBE_BATCH_VBE_RANKS.value,
+            ),
+            # Indices Parameters
+            click.option(
+                "--tbe-indices-hitters",
+                type=str,
+                default="",
+                help=TBEDataConfigHelperText.TBE_INDICES_HITTERS.value,
+            ),
+            click.option(
+                "--tbe-indices-zipf",
+                type=(float, float),
+                default=(0.1, 0.1),
+                help=TBEDataConfigHelperText.TBE_INDICES_ZIPF.value,
+            ),
+            click.option(
+                "--tbe-indices-dtype",
+                type=click.Choice(["32", "64"]),
+                default="64",
+                help=TBEDataConfigHelperText.TBE_INDICES_DTYPE.value,
+            ),
+            click.option(
+                "--tbe-offsets-dtype",
+                type=click.Choice(["32", "64"]),
+                default="64",
+                help=TBEDataConfigHelperText.TBE_OFFSETS_DTYPE.value,
+            ),
+            # Pooling Parameters
+            click.option(
+                "--tbe-pooling-size",
+                type=int,
+                default=20,
+                help=TBEDataConfigHelperText.TBE_POOLING_SIZE.value,
+            ),
+            click.option(
+                "--tbe-pooling-vl-sigma",
+                type=int,
+                required=False,
+                help=TBEDataConfigHelperText.TBE_POOLING_VL_SIGMA.value,
+            ),
+            click.option(
+                "--tbe-pooling-vl-dist",
+                type=click.Choice(["uniform", "normal"]),
+                required=False,
+                help=TBEDataConfigHelperText.TBE_POOLING_VL_DIST.value,
+            ),
+        ]
+
+        for option in reversed(options):
+            func = option(func)
+        return func
+
+    @classmethod
+    def load_from_file(cls, filepath: str) -> TBEDataConfig:
+        with open(filepath, "r") as f:
+            if filepath.endswith(".yaml") or filepath.endswith(".yml"):
+                data = yaml.safe_load(f)
+                return TBEDataConfig.from_dict(data).validate()
+            else:
+                return TBEDataConfig.from_json(f.read()).validate()
+
+    @classmethod
+    def load_from_context(cls, context: click.Context) -> TBEDataConfig:
+        params = context.params
+
+        # Read table parameters
+        T = params["tbe_num_tables"]
+        E = params["tbe_num_embeddings"]
+        if params["tbe_num_embeddings_list"] is not None:
+            Es = [int(x) for x in params["tbe_num_embeddings_list"].split(",")]
+        else:
+            Es = None
+        D = params["tbe_embedding_dim"]
+        if params["tbe_embedding_dim_list"] is not None:
+            Ds = [int(x) for x in params["tbe_embedding_dim_list"].split(",")]
+        else:
+            Ds = None
+
+        mixed_dim = params["tbe_mixed_dim"]
+        weighted = params["tbe_weighted"]
+        if params["tbe_max_indices"] is not None:
+            max_indices = params["tbe_max_indices"]
+        else:
+            max_indices = None
+
+        # Read batch parameters
+        B = params["tbe_batch_size"]
+        sigma_B = params["tbe_batch_vbe_sigma"]
+        vbe_distribution = params["tbe_batch_vbe_dist"]
+        vbe_num_ranks = params["tbe_batch_vbe_ranks"]
+        if params["tbe_batch_sizes_list"] is not None:
+            Bs = [int(x) for x in params["tbe_batch_sizes_list"].split(",")]
+        else:
+            Bs = None
+        batch_params = BatchParams(B, sigma_B, vbe_distribution, vbe_num_ranks, Bs)
+
+        # Read indices parameters
+        heavy_hitters = (
+            torch.tensor([float(x) for x in params["tbe_indices_hitters"].split(",")])
+            if params["tbe_indices_hitters"]
+            else torch.tensor([])
+        )
+        zipf_q, zipf_s = params["tbe_indices_zipf"]
+        index_dtype = (
+            torch.int32 if int(params["tbe_indices_dtype"]) == 32 else torch.int64
+        )
+        offset_dtype = (
+            torch.int32 if int(params["tbe_offsets_dtype"]) == 32 else torch.int64
+        )
+        indices_params = IndicesParams(
+            heavy_hitters, zipf_q, zipf_s, index_dtype, offset_dtype
+        )
+
+        # Read pooling parameters
+        L = params["tbe_pooling_size"]
+        sigma_L = params["tbe_pooling_vl_sigma"]
+        length_distribution = params["tbe_pooling_vl_dist"]
+        pooling_params = PoolingParams(L, sigma_L, length_distribution)
+
+        return TBEDataConfig(
+            T,
+            E,
+            D,
+            mixed_dim,
+            weighted,
+            batch_params,
+            indices_params,
+            pooling_params,
+            not torch.cuda.is_available(),
+            Es,
+            Ds,
+            max_indices,
+        ).validate()
+
+    @classmethod
+    def load(cls, context: click.Context) -> TBEDataConfig:
+        tbe_config_filepath = context.params["tbe_config"]
+        if tbe_config_filepath is not None:
+            return cls.load_from_file(tbe_config_filepath)
+        else:
+            return cls.load_from_context(context)

fbgemm_gpu/tbe/bench/tbe_data_config_param_models.py

@@ -0,0 +1,170 @@
+#!/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 typing import Any, Optional
+
+import torch
+
+
+def str_to_int_dtype(dtype: str) -> torch.dtype:
+    if dtype == "torch.int32":
+        return torch.int32
+    elif dtype == "torch.int64":
+        return torch.int64
+    else:
+        raise ValueError(f"Unsupported dtype: {dtype}")
+
+
+@dataclasses.dataclass(frozen=True, eq=False)
+class IndicesParams:
+    # Heavy hitters for the Zipf distribution, i.e. a probability density map
+    # for the most hot indices.  There should not ever be more than 100
+    # elements, and currently it is limited to 20 entries (kHeavyHittersMaxSize)
+    heavy_hitters: torch.Tensor
+    # zipf*: parameters for the Zipf distribution (x+q)^{-s}
+    zipf_q: float
+    # zipf_s is synonymous with alpha in the literature
+    zipf_s: float
+    # [Optional] dtype for indices tensor
+    index_dtype: Optional[torch.dtype] = None
+    # [Optional] dtype for offsets tensor
+    offset_dtype: Optional[torch.dtype] = None
+
+    @classmethod
+    # pyre-ignore [3]
+    def from_dict(cls, data: dict[str, Any]):
+        if not isinstance(data["heavy_hitters"], torch.Tensor):
+            data["heavy_hitters"] = torch.tensor(
+                data["heavy_hitters"], dtype=torch.float32
+            )
+            data["index_dtype"] = str_to_int_dtype(data["index_dtype"])
+            data["offset_dtype"] = str_to_int_dtype(data["offset_dtype"])
+        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]:
+        # https://stackoverflow.com/questions/73735974/convert-dataclass-of-dataclass-to-json-string
+        tmp = dataclasses.asdict(self)
+        # Convert tensor to list for JSON serialization
+        tmp["heavy_hitters"] = self.heavy_hitters.tolist()
+        tmp["index_dtype"] = str(self.index_dtype)
+        tmp["offset_dtype"] = str(self.offset_dtype)
+        return tmp
+
+    def json(self, format: bool = False) -> str:
+        return json.dumps(self.dict(), indent=(2 if format else -1), sort_keys=True)
+
+    # pyre-ignore [2]
+    def __eq__(self, other) -> bool:
+        return (
+            (self.zipf_q, self.zipf_s, self.index_dtype, self.offset_dtype)
+            == (other.zipf_q, other.zipf_s, other.index_dtype, other.offset_dtype)
+        ) and bool((self.heavy_hitters - other.heavy_hitters).abs().max() < 1e-6)
+
+    # pyre-ignore [3]
+    def validate(self):
+        assert self.zipf_q > 0, "zipf_q must be positive"
+        assert self.zipf_s > 0, "zipf_s must be positive"
+        assert self.index_dtype is None or self.index_dtype in [
+            torch.int32,
+            torch.int64,
+        ], "index_dtype must be one of [torch.int32, torch.int64]"
+        assert self.offset_dtype is None or self.offset_dtype in [
+            torch.int32,
+            torch.int64,
+        ], "offset_dtype must be one of [torch.int32, torch.int64]"
+        return self
+
+
+@dataclasses.dataclass(frozen=True)
+class BatchParams:
+    # Target batch size, i.e. number of batch lookups per table
+    B: int
+    # [Optional] Standard deviation of B (for variable batch size configuration)
+    sigma_B: Optional[int] = None
+    # [Optional] Distribution of batch sizes (normal, uniform)
+    vbe_distribution: Optional[str] = "normal"
+    # Number of ranks for variable batch size generation
+    vbe_num_ranks: Optional[int] = None
+    # List of target batch sizes, i.e. number of batch lookups per table
+    Bs: Optional[list[int]] = None
+
+    @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):
+        if self.Bs is not None:
+            assert all(b > 0 for b in self.Bs), "All elements in Bs must be positive"
+        else:
+            assert self.B > 0, "B must be positive"
+        assert not self.sigma_B or self.sigma_B > 0, "sigma_B must be positive"
+        assert (
+            self.vbe_num_ranks is None or self.vbe_num_ranks > 0
+        ), "vbe_num_ranks must be positive"
+        assert self.vbe_distribution is None or self.vbe_distribution in [
+            "normal",
+            "uniform",
+        ], "vbe_distribution must be one of [normal, uniform]"
+        return self
+
+
+@dataclasses.dataclass(frozen=True)
+class PoolingParams:
+    # Target bag size, i.e. pooling factor, or number of indices per batch lookup
+    L: int
+    # [Optional] Standard deviation of L (for variable bag size configuration)
+    sigma_L: Optional[int] = None
+    # [Optional] Distribution of embedding sequence lengths (normal, uniform)
+    length_distribution: Optional[str] = "normal"
+
+    @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.L > 0, "L must be positive"
+        assert not self.sigma_L or self.sigma_L > 0, "sigma_L must be positive"
+        assert self.length_distribution is None or self.length_distribution in [
+            "normal",
+            "uniform",
+        ], "length_distribution must be one of [normal, uniform]"
+        return self

fbgemm_gpu/tbe/bench/utils.py

@@ -0,0 +1,48 @@
+# 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 logging
+
+import numpy as np
+import torch
+
+from fbgemm_gpu.split_embedding_configs import SparseType
+
+logging.basicConfig(level=logging.DEBUG)
+
+
+def fill_random_scale_bias(
+    emb: torch.nn.Module,
+    T: int,
+    weights_precision: SparseType,
+) -> None:
+    for t in range(T):
+        # pyre-fixme[29]: `Union[Module, Tensor]` is not a function.
+        (weights, scale_shift) = emb.split_embedding_weights()[t]
+        if scale_shift is not None:
+            (E, R) = scale_shift.shape
+            assert R == 4
+            scales = None
+            shifts = None
+            if weights_precision == SparseType.INT8:
+                scales = np.random.uniform(0.001, 0.01, size=(E,)).astype(np.float16)
+                shifts = np.random.normal(-2, 2, size=(E,)).astype(np.float16)
+            elif weights_precision == SparseType.INT4:
+                scales = np.random.uniform(0.01, 0.1, size=(E,)).astype(np.float16)
+                shifts = np.random.normal(-2, 2, size=(E,)).astype(np.float16)
+            elif weights_precision == SparseType.INT2:
+                scales = np.random.uniform(0.1, 1, size=(E,)).astype(np.float16)
+                shifts = np.random.normal(-2, 2, size=(E,)).astype(np.float16)
+            scale_shift.copy_(
+                torch.tensor(
+                    np.stack([scales, shifts], axis=1)
+                    .astype(np.float16)
+                    .view(np.uint8),
+                    device=scale_shift.device,
+                )
+            )

fbgemm_gpu/tbe/cache/__init__.py

@@ -0,0 +1,11 @@
+#!/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 .kv_embedding_ops_inference import KVEmbeddingInference  # noqa: F401
+from .split_embeddings_cache_ops import get_unique_indices_v2  # noqa: F401