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

fbgemm_gpu/__init__.py

@@ -0,0 +1,186 @@
+#!/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.
+
+import json
+import logging
+import os
+import re
+
+import torch
+
+# Based on the FBGEMM-PyTorch compatibility table at
+# https://docs.pytorch.org/FBGEMM/general/Releases.html#fbgemm-releases-compatibility
+_fbgemm_torch_compat_table = {
+    "1.5": "2.10",
+    "1.4": "2.9",
+    "1.3": "2.8",
+    "1.2": "2.7",
+    "1.1": "2.6",
+    "1.0": "2.5",
+    "0.8": "2.4",
+    "0.7": "2.3",
+    "0.6": "2.2",
+    "0.5": "2.1",
+    "0.4": "2.0",
+}
+
+
+def _load_target_info(target: str) -> dict[str, str]:
+    try:
+        filepath = os.path.join(
+            os.path.dirname(__file__), "docs", f"target.{target}.json.py"
+        )
+        with open(filepath, "r") as file:
+            data = json.load(file)
+    except Exception:
+        data = {}
+
+    return data
+
+
+def _load_library(filename: str, version: str, no_throw: bool = False) -> None:
+    """Load a shared library from the given filename."""
+
+    # Check if the version of PyTorch is compatible with the version of FBGEMM
+    # that we are trying to load, and print a loud warning if not.  This is
+    # useful for the OSS build, where we have a single FBGEMM library that is
+    # compatible with multiple versions of PyTorch.
+    #
+    # Based on: https://github.com/pytorch/ao/blob/main/torchao/__init__.py#L30
+
+    keys = [
+        key
+        for key in _fbgemm_torch_compat_table.keys()
+        if version.startswith(f"{key}.")
+    ]
+
+    if version == "INTERNAL" or "+git" in version:
+        # if FBGEMM version has "+git", assume it's locally built and we don't know
+        #   anything about the PyTorch version used to build it
+        logging.info(
+            "FBGEMM version is INTERNAL or local, ignoring version compatibility check with PyTorch"
+        )
+
+    elif re.match(r"^\d{4}\.\d{1,2}\.\d{1,2}.*$", version):
+        # if FBGEMM version is a date, assume it's a nightly build and that we
+        # know what we're doing
+        logging.info(
+            "FBGEMM version is a nightly version, ignoring version compatibility check with PyTorch"
+        )
+
+    elif not keys:
+        logging.warning(
+            f"""
+            \033[33m
+            _fbgemm_torch_compat_table has no entry for {version} of FBGEMM;
+            cannot determine compatibility with PyTorch {torch.__version__}
+            \033[0m
+            """
+        )
+
+    elif not str(torch.__version__).startswith(_fbgemm_torch_compat_table[keys[0]]):
+        logging.warning(
+            f"""
+            \033[31m
+            FBGEMM_GPU version is {version}, which is not guaranteed to be
+            compatible with PyTorch {torch.__version__}; library loading might
+            crash!
+
+            Please refer to
+            https://docs.pytorch.org/FBGEMM/general/Releases.html#fbgemm-releases-compatibility
+            for the FBGEMM-PyTorch compatibility table.
+            \033[0m
+            """
+        )
+
+    try:
+        torch.ops.load_library(os.path.join(os.path.dirname(__file__), filename))
+        logging.info(f"Successfully loaded: '{filename}'")
+
+    except Exception as error:
+        logging.error(f"Could not load the library '{filename}'!\n\n\n{error}\n\n\n")
+        if not no_throw:
+            raise error
+
+
+# Since __init__.py is only used in OSS context, we define `open_source` here
+# and use its existence to determine whether or not we are in OSS context
+open_source: bool = True
+
+# Trigger the manual addition of docstrings to pybind11-generated operators
+import fbgemm_gpu.docs  # noqa: F401, E402
+
+
+__targets_infos__ = {
+    target: _load_target_info(target) for target in ["default", "genai", "hstu"]
+}
+__targets_infos__ = {k: v for (k, v) in __targets_infos__.items() if v}
+
+try:
+    __target__, __info__ = next(iter(__targets_infos__.items()))
+    __variant__ = __info__["variant"]
+    __version__ = __info__["version"]
+except Exception:
+    __variant__: str = "INTERNAL"
+    __version__: str = "INTERNAL"
+    __target__: str = "INTERNAL"
+
+fbgemm_gpu_libraries = [
+    "fbgemm_gpu_config",
+    "fbgemm_gpu_tbe_utils",
+    "fbgemm_gpu_tbe_index_select",
+    "fbgemm_gpu_tbe_cache",
+    "fbgemm_gpu_tbe_optimizers",
+    "fbgemm_gpu_tbe_inference",
+    "fbgemm_gpu_tbe_training_forward",
+    "fbgemm_gpu_tbe_training_backward",
+    "fbgemm_gpu_tbe_training_backward_pt2",
+    "fbgemm_gpu_tbe_training_backward_dense",
+    "fbgemm_gpu_tbe_training_backward_split_host",
+    "fbgemm_gpu_tbe_training_backward_gwd",
+    "fbgemm_gpu_tbe_training_backward_vbe",
+    "fbgemm_gpu_py",
+]
+
+fbgemm_genai_libraries = [
+    "experimental/gen_ai/fbgemm_gpu_experimental_gen_ai",
+]
+
+# NOTE: While FBGEMM_GPU GenAI is not available for ROCm yet, we would like to
+# be able to install the existing CUDA variant of the package onto ROCm systems,
+# so that we can at least use the Triton GEMM libraries from experimental/gemm.
+# But loading fbgemm_gpu package will trigger load-checking the .SO file for the
+# GenAI libraries, which will fail.  This workaround ignores check-loading the
+# .SO file for the ROCm case, so that clients can import
+# fbgemm_gpu.experimental.gemm without triggering an error.
+if torch.cuda.is_available() and torch.version.hip:
+    fbgemm_genai_libraries = []
+
+libraries_to_load = {
+    "default": fbgemm_gpu_libraries,
+    "genai": fbgemm_genai_libraries,
+}
+
+for target, info in __targets_infos__.items():
+    for library in libraries_to_load.get(target, []):
+        # NOTE: In all cases, we want to throw an error if we cannot load the
+        # library.  However, this appears to break the OSS documentation build,
+        # where the Python documentation doesn't show up in the generated docs.
+        #
+        # To work around this problem, we introduce a fake build variant called
+        # `docs` and we only throw a library load error when the variant is not
+        # `docs`.  For more information, see:
+        #
+        #   https://github.com/pytorch/FBGEMM/pull/3477
+        #   https://github.com/pytorch/FBGEMM/pull/3717
+        _load_library(f"{library}.so", info["version"], info["variant"] == "docs")
+
+try:
+    # Trigger meta operator registrations
+    from . import sparse_ops  # noqa: F401, E402
+except Exception:
+    pass

fbgemm_gpu/batched_unary_embeddings_ops.py

@@ -0,0 +1,87 @@
+#!/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
+
+
+from math import sqrt
+
+import torch
+
+from fbgemm_gpu.utils.loader import load_torch_module
+
+try:
+    # pyre-ignore[21]
+    from fbgemm_gpu import open_source  # noqa: F401
+except Exception:
+    load_torch_module("//deeplearning/fbgemm/fbgemm_gpu:sparse_ops")
+
+
+def wrap_weight_to_parameter(weights: list[torch.Tensor]) -> list[torch.Tensor]:
+    for i, v in enumerate(weights):
+        if not isinstance(v, torch.nn.Parameter):
+            weights[i] = torch.nn.Parameter(v)
+    return weights
+
+
+class BatchedUnaryEmbeddingBag(torch.nn.Module):
+    # pyre-fixme[3]: Return type must be annotated.
+    def __init__(self, num_tasks: int, hash_sizes: list[int], long_index: bool = False):
+        super().__init__()
+        self.num_tasks = num_tasks
+        self.hash_sizes = hash_sizes
+        # [N][sum(E)][1]
+        embedding_data = torch.randn(size=(num_tasks, sum(self.hash_sizes), 1))
+        self.weight = torch.nn.Parameter(embedding_data)
+        index_dtype = torch.int64 if long_index else torch.int32
+        table_offsets_tensor = torch.cat(
+            [
+                torch.tensor([0], dtype=index_dtype),
+                torch.cumsum(
+                    torch.tensor(hash_sizes),
+                    dim=0,
+                    dtype=index_dtype,
+                ),
+            ]
+        )
+        self.register_buffer("table_offsets_tensor", table_offsets_tensor)
+        self.init_parameters()
+
+    # pyre-fixme[3]: Return type must be annotated.
+    def forward(self, offsets: torch.Tensor, input: torch.Tensor):
+        # output is [N][B][T]
+        return torch.ops.fbgemm.batched_unary_embeddings(
+            self.weight,
+            self.table_offsets_tensor,
+            offsets,
+            input,
+        )
+
+    @torch.jit.export
+    # pyre-fixme[3]: Return type must be annotated.
+    def split_embedding_weights(self):
+        embedding_weights = []
+        for n in range(self.num_tasks):
+            for t in range(len(self.hash_sizes)):
+                embedding_weights.append(
+                    self.weight.detach()[
+                        n,
+                        self.table_offsets_tensor[t] : self.table_offsets_tensor[t + 1],
+                        :,
+                    ]
+                )
+        return embedding_weights
+
+    @torch.jit.export
+    # pyre-fixme[3]: Return type must be annotated.
+    def init_parameters(self):
+        for num_emb, param in zip(
+            self.hash_sizes * self.num_tasks,
+            wrap_weight_to_parameter(self.split_embedding_weights()),
+        ):
+            assert param.shape == (num_emb, 1)
+            param.data.uniform_(-sqrt(1 / num_emb), sqrt(1 / num_emb))

fbgemm_gpu/config/__init__.py

@@ -0,0 +1,9 @@
+# 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
+
+from .feature_list import FeatureGate, FeatureGateName  # noqa F401

fbgemm_gpu/config/feature_list.py

@@ -0,0 +1,88 @@
+# 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
+
+from enum import auto, Enum
+
+import torch
+
+try:
+    torch.ops.load_library("//deeplearning/fbgemm/fbgemm_gpu:config_cpp_torch_op")
+except Exception:
+    import fbgemm_gpu  # noqa F401
+
+
+# Note: ENUM name must match EXACTLY with the JK knob name in the UI
+class FeatureGateName(Enum):
+    """
+    FBGEMM_GPU feature gates enum (Python).
+
+    **Code Example:**
+
+    .. code-block:: python
+
+        from fbgemm_gpu.config import FeatureGateName
+
+        def foo():
+            if FeatureGateName.TBE_V2.is_enabled():
+
+                # Do something if feature is enabled
+                ...
+            else:
+                # Do something different if feature is disabled
+                ...
+
+    Note:
+        While not required, it is best to mirror the enum values in C++,
+        in `fbgemm_gpu/config/feature_gates.h`.
+
+        For fbcode: The ENUM name must match EXACTLY with the JK knob name in the UI
+        For OSS: The environment variable will be evaluated as f"FBGEMM_{ENUM}"
+
+    """
+
+    # Enable TBE V2 APIs
+    TBE_V2 = auto()
+
+    # Enable Ensemble Rowwise Adagrad (D60189486 stack)
+    TBE_ENSEMBLE_ROWWISE_ADAGRAD = auto()
+
+    # Enable ROCm packed bags optimization in TBE inference
+    TBE_ROCM_INFERENCE_PACKED_BAGS = auto()
+
+    # Enable HIP-based backward kernels in TBE training
+    TBE_ROCM_HIP_BACKWARD_KERNEL = auto()
+
+    # Enable bounds_check_indices_v2
+    BOUNDS_CHECK_INDICES_V2 = auto()
+
+    # Enable TBE input parameters extraction
+    TBE_REPORT_INPUT_PARAMS = auto()
+
+    def is_enabled(self) -> bool:
+        return FeatureGate.is_enabled(self)
+
+
+class FeatureGate:
+    """
+    FBGEMM_GPU feature gate.
+
+    This class exists because methods defined on enums cannot be invoked when
+    the enum is packaged into a model (the mechanism is unclear).
+
+    **Code Example:**
+
+    .. code-block:: python
+
+        from deeplearning.fbgemm.fbgemm_gpu.config import FeatureGate, FeatureGateName
+
+        FeatureGate.is_enabled(FeatureGateName.TBE_V2)
+    """
+
+    @classmethod
+    def is_enabled(cls, feature: FeatureGateName) -> bool:
+        return torch.ops.fbgemm.check_feature_gate_key(feature.name)

fbgemm_gpu/docs/__init__.py

@@ -0,0 +1,18 @@
+#!/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.
+
+# Trigger the manual addition of docstrings to pybind11-generated operators
+try:
+    from . import (  # noqa: F401
+        jagged_tensor_ops,
+        merge_pooled_embedding_ops,
+        permute_pooled_embedding_ops,
+        quantize_ops,
+        sparse_ops,
+    )
+except Exception:
+    pass

fbgemm_gpu/docs/common.py

@@ -0,0 +1,9 @@
+# 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.
+
+
+def add_docs(method, docstr: str):
+    method.__doc__ = docstr

fbgemm_gpu/docs/examples.py

@@ -0,0 +1,73 @@
+# 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.
+
+from ctypes import c_size_t
+
+
+# [fbgemm-gpu.python.docs.examples.docstring.start]
+def example_method(alignment: c_size_t, param: float) -> int:
+    """
+    This class is an example of how you can write docstrings.
+    You can add multiple lines of those descriptions. Make sure to include
+    useful information about your method.
+
+    **Code Example:**
+
+    .. code-block:: cpp
+
+        // Here is a C++ code block
+        std::vector<int32_t> foo(const std::vector<int32_t> &lst) {
+            std::vector<int32_t> ret;
+            for (const auto x : lst) {
+                ret.emplace_back(x * 2);
+            }
+            return ret;
+        }
+
+    And here is a verbatim-text diagram example:
+
+    .. code-block:: text
+
+        .------+---------------------------------.-----------------------------
+        |            Block A (first)             |       Block B (second)
+        +------+------+--------------------------+------+------+---------------
+        | Next | Prev |   usable space           | Next | Prev | usable space..
+        +------+------+--------------------------+------+--+---+---------------
+        ^  |                                     ^         |
+        |  '-------------------------------------'         |
+        |                                                  |
+        '----------- Block B's prev points to Block A -----'
+
+    Todo:
+        * This is a TODO item.
+        * And a second TODO item.
+
+    Args:
+        alignment (c_size_t): Description of the `alignment` value.
+        param (float): Description of `param1`.
+
+    Returns:
+        Description of the method's return value.
+
+    Raises:
+        AttributeError: If there is an error with the attributes.
+        ValueError: If `param` is equal to 3.14.
+
+    Example:
+        This is how you can use this function
+
+        >>> print("Code blocks are supported")
+
+    Note:
+        For more info on reStructuredText docstrings, see
+        `here <https://www.sphinx-doc.org/en/master/usage/restructuredtext/basics.html>`__
+        and
+        `here <https://peps.python.org/pep-0287/>`__.
+    """
+    return 42
+
+
+# [fbgemm-gpu.python.docs.examples.docstring.end]