fbgemm-gpu-nightly-cpu 2025.3.27__cp311-cp311-manylinux_2_28_aarch64.whl → 2026.1.29__cp311-cp311-manylinux_2_28_aarch64.whl

fbgemm_gpu/split_embedding_configs.py

@@ -8,10 +8,51 @@
 # pyre-strict
 
 import enum
+import itertools
 from typing import Any, Dict  # noqa: F401
 
 import torch
 
+# fmt:skip
+from fbgemm_gpu.split_table_batched_embeddings_ops_common import (
+    EmbeddingLocation,
+    SplitState,
+)
+
+
+def pad4(value: int) -> int:
+    """
+    Compute the smallest multiple of 4 that is greater than or equal to the given value.
+
+    Parameters:
+        value (int): The integer to align (must be non-negative).
+
+    Returns:
+        int: The aligned value.
+
+    Raises:
+        ValueError: If the input is negative.
+        TypeError: If the input is not an integer.
+    """
+    return (int(value) + 3) & ~3
+
+
+def pad16(value: int) -> int:
+    """
+    Compute the smallest multiple of 16 that is greater than or equal to the given value.
+
+    Parameters:
+        value (int): The integer to align (must be non-negative).
+
+    Returns:
+        int: The aligned value.
+
+    Raises:
+        ValueError: If the input is negative.
+        TypeError: If the input is not an integer.
+    """
+    return (int(value) + 15) & ~15
+
 
 @enum.unique
 class EmbOptimType(enum.Enum):
@@ -40,6 +81,196 @@ class EmbOptimType(enum.Enum):
     def __str__(self) -> str:
         return self.value
 
+    def _extract_dtype(
+        self, optimizer_state_dtypes: dict[str, "SparseType"], name: str
+    ) -> torch.dtype:
+        if optimizer_state_dtypes is None or name not in optimizer_state_dtypes:
+            return torch.float32
+        return optimizer_state_dtypes[name].as_dtype()
+
+    def state_names(self) -> list[str]:
+        """
+        Returns the names of the optimizer states.  The order of the states will
+        be the order in which they are processed and returned in
+        SSDTableBatchedEmbeddingBags.split_optimizer_states(), but this is not
+        necessarily the same as the order they are stored in the memory layout.
+        """
+        if self == EmbOptimType.EXACT_ROWWISE_ADAGRAD:
+            return ["momentum1"]
+        elif self in [EmbOptimType.PARTIAL_ROWWISE_ADAM, EmbOptimType.ADAM]:
+            return ["momentum1", "momentum2"]
+        else:
+            return []
+
+    def state_size_table(self, D: int) -> dict[str, int]:
+        """
+        Returns the table of state names to state sizes in terms of number of
+        elements (per table row)
+        """
+        if self == EmbOptimType.EXACT_ROWWISE_ADAGRAD:
+            return {"momentum1": 1}
+        elif self == EmbOptimType.PARTIAL_ROWWISE_ADAM:
+            return {"momentum1": D, "momentum2": 1}
+        elif self == EmbOptimType.ADAM:
+            return {"momentum1": D, "momentum2": D}
+        else:
+            return {}
+
+    def state_size_nbytes(
+        self,
+        D: int,
+        optimizer_state_dtypes: dict[str, "SparseType"] = {},  # noqa: B006
+    ) -> int:
+        """
+        Returns the size of the data (in bytes) required to hold the optimizer
+        state (per table row).  This size includes byte-padding.
+        """
+        momentum1_dtype = self._extract_dtype(optimizer_state_dtypes, "momentum1")
+        momentum2_dtype = self._extract_dtype(optimizer_state_dtypes, "momentum2")
+
+        if self == EmbOptimType.EXACT_ROWWISE_ADAGRAD:
+            return momentum1_dtype.itemsize
+
+        elif self == EmbOptimType.PARTIAL_ROWWISE_ADAM:
+            return pad4(1 * momentum2_dtype.itemsize) + D * momentum1_dtype.itemsize
+
+        elif self == EmbOptimType.ADAM:
+            return (D * momentum1_dtype.itemsize) + (D * momentum2_dtype.itemsize)
+
+        else:
+            return 0
+
+    def byte_offsets_along_row(
+        self,
+        D: int,
+        weights_precision: "SparseType",
+        optimizer_state_dtypes: dict[str, "SparseType"] = {},  # noqa: B006
+    ) -> dict[str, tuple[int, int]]:
+        """
+        Returns the start and end byte offsets of each optimizer state along a
+        cache row with optimizer state offloading enabled.
+        """
+        # Extract the optimizer state dtypes
+        momentum1_dtype = self._extract_dtype(optimizer_state_dtypes, "momentum1")
+        momentum2_dtype = self._extract_dtype(optimizer_state_dtypes, "momentum2")
+
+        # This is the pointer to where the optimizer state begins in the memory
+        p0 = pad4(D) * weights_precision.as_dtype().itemsize
+
+        if self == EmbOptimType.EXACT_ROWWISE_ADAGRAD:
+            return {"momentum1": (p0, p0 + momentum1_dtype.itemsize)}
+
+        elif self == EmbOptimType.PARTIAL_ROWWISE_ADAM:
+            # momentum1 lies after momentum2
+            p1 = p0 + pad4(1 * momentum2_dtype.itemsize)
+            return {
+                "momentum2": (p0, p0 + momentum2_dtype.itemsize),
+                "momentum1": (
+                    p1,
+                    p1 + D * momentum1_dtype.itemsize,
+                ),
+            }
+
+        elif self == EmbOptimType.ADAM:
+            # momentum2 lies after momentum1
+            p1 = p0 + (D * momentum1_dtype.itemsize)
+
+            return {
+                "momentum1": (p0, p1),
+                "momentum2": (p1, p1 + D * momentum2_dtype.itemsize),
+            }
+
+        else:
+            return {}
+
+    def empty_states(
+        self,
+        rows: list[int],
+        dims: list[int],
+        optimizer_state_dtypes: dict[str, "SparseType"] = {},  # noqa: B006
+    ) -> list[list[torch.Tensor]]:
+        """
+        Creates sets of empty tensors per table to hold optimizer states based
+        on the specified optimizer type, state dtypes, embedding specs, and
+        (optionally) local row counts.
+        """
+        # Else, check that the local row count for each table is set
+        assert len(rows) == len(dims)
+
+        opt_states_set: list[list[torch.Tensor]] = []
+
+        for r, D in zip(rows, dims):
+            # Set up the table of state names to state sizes, ordered by their
+            # memory layout
+            state_size_table = self.state_size_table(D)
+            ordered_state_sizes = [(k, state_size_table[k]) for k in self.state_names()]
+
+            # Create the optimizer states for this table
+            opt_states_set.append(
+                [
+                    torch.empty(
+                        # If the state size is 1, then fix tensor to 1D to be
+                        # consistent with training.py code
+                        # pyre-ignore [6]
+                        (r, d) if d > 1 else r,
+                        dtype=self._extract_dtype(optimizer_state_dtypes, state_name),
+                        device="cpu",
+                    )
+                    for state_name, d in ordered_state_sizes
+                ]
+            )
+
+        return opt_states_set
+
+    def ssd_state_splits(
+        self,
+        embedding_specs: list[tuple[int, int]],  # Tuple of (rows, dims)
+        optimizer_state_dtypes: dict[str, "SparseType"] = {},  # noqa: B006
+        enable_optimizer_offloading: bool = False,
+    ) -> list[tuple[SplitState, str, torch.dtype]]:
+        """
+        Returns the split planning for the optimizer states
+        """
+        rows, _ = zip(*embedding_specs)
+        T_ = len(embedding_specs)
+
+        # This is the cumulative row counts for rowwise states
+        row_count_cumsum: list[int] = [0] + list(itertools.accumulate(rows))
+        # This is the cumulative element counts for elementwise states
+        table_size_cumsum: list[int] = [0] + list(
+            itertools.accumulate([r * d for r, d in embedding_specs])
+        )
+
+        if self == EmbOptimType.EXACT_ROWWISE_ADAGRAD:
+            params = {"momentum1": row_count_cumsum}
+        elif self == EmbOptimType.PARTIAL_ROWWISE_ADAM:
+            params = {"momentum1": table_size_cumsum, "momentum2": row_count_cumsum}
+        elif self == EmbOptimType.ADAM:
+            params = {
+                "momentum1": table_size_cumsum,
+                "momentum2": table_size_cumsum,
+                "row_counter": row_count_cumsum,
+            }
+        else:
+            params = {}
+
+        return [
+            (
+                SplitState(
+                    dev_size=(
+                        cumsum_table[-1] if not enable_optimizer_offloading else 0
+                    ),
+                    host_size=0,
+                    uvm_size=0,
+                    placements=[EmbeddingLocation.DEVICE for _ in range(T_)],
+                    offsets=cumsum_table[:-1],
+                ),
+                name,
+                self._extract_dtype(optimizer_state_dtypes, name),
+            )
+            for (name, cumsum_table) in params.items()
+        ]
+
 
 # Base class for quantization configuration (in case other numeric types have
 # configs)
@@ -79,14 +310,54 @@ def sparse_type_to_int(sparse_type: "SparseType") -> int:
         SparseType.BF16.value: 5,
         SparseType.FP8.value: 6,
         SparseType.MX4.value: 7,
+        SparseType.NFP8.value: 8,
     }[sparse_type.value]
 
 
+def sparse_type_int_to_dtype(ty: int) -> torch.dtype:
+    """
+    TorchScript-compatible function to convert an SparseType enum as integer) to torch.dtype.
+
+    This is a standalone function equivalent to SparseType.from_int(dtype_int).as_dtype() that works
+    with TorchScript. TorchScript does not support @staticmethod on Enum classes,
+    so this function provides a workaround.
+    """
+    if ty == 0:  # fp32
+        return torch.float32
+    elif ty == 1:  # fp16
+        return torch.float16
+    elif ty == 2:  # int8
+        return torch.uint8
+    elif ty == 3:  # int4
+        return torch.quint4x2
+    elif ty == 4:  # int2
+        return torch.quint2x4
+    elif ty == 5:  # bf16
+        return torch.bfloat16
+    elif ty == 6:  # fp8
+        return torch.uint8
+    elif ty == 7:  # mx4
+        return torch.uint8
+    elif ty == 9:
+        return (
+            torch.float8_e4m3fnuz
+            if torch.version.hip is not None
+            else torch.float8_e4m3fn
+        )
+    else:  # Invalid is 7 or non enumerated.
+        raise ValueError(f"Unsupported sparse type: {ty}")
+
+
 @enum.unique
 class SparseType(enum.Enum):
     FP32 = "fp32"
     FP16 = "fp16"
     FP8 = "fp8"
+    # NFP8 refers to "native" FP8 in that it uses the GPU implementations
+    # of E4M3 whereas the other FP8 sparsetype uses a custom format. Use of
+    # NFP8 allows us to use hardware casting intrinsics which can be much faster.
+    # Eventually, we should merge these two types.
+    NFP8 = "nfp8"
     INT8 = "int8"
     INT4 = "int4"
     INT2 = "int2"
@@ -112,9 +383,11 @@ class SparseType(enum.Enum):
             return SparseType("bf16")
         elif ty == 6:
             return SparseType("fp8")
-        elif ty == 7:
+        elif ty == 8:
             return SparseType("mx4")
-        else:
+        elif ty == 9:
+            return SparseType("nfp8")
+        else:  # Invalid is 7 or non enumerated.
             raise ValueError(f"Unsupported sparse type: {ty}")
 
     def as_int(self) -> int:
@@ -136,6 +409,8 @@ class SparseType(enum.Enum):
             return SparseType("bf16")
         elif dtype == torch.uint8:
             return SparseType("mx4")
+        elif dtype == torch.float8_e4m3fnuz or dtype == torch.float8_e4m3fn:
+            return SparseType("nfp8")
         else:
             raise ValueError(f"Unsupported sparse dtype: {dtype}")
 
@@ -149,6 +424,11 @@ class SparseType(enum.Enum):
             SparseType.INT2.value: torch.quint2x4,
             SparseType.BF16.value: torch.bfloat16,
             SparseType.MX4.value: torch.uint8,
+            SparseType.NFP8.value: (
+                torch.float8_e4m3fnuz
+                if torch.version.hip is not None
+                else torch.float8_e4m3fn
+            ),
         }[self.value]
 
     def bit_rate(self) -> int:
@@ -161,6 +441,7 @@ class SparseType(enum.Enum):
             SparseType.INT2.value: 2,
             SparseType.BF16.value: 16,
             SparseType.MX4.value: 4,
+            SparseType.NFP8.value: 8,
         }[self.value]
 
     def align_size(self) -> int:
@@ -173,6 +454,7 @@ class SparseType(enum.Enum):
             SparseType.INT2.value: 16,
             SparseType.BF16.value: 2,
             SparseType.MX4.value: 8,
+            SparseType.NFP8.value: 4,
         }[self.value]
 
     def is_float(self) -> bool:
@@ -181,6 +463,7 @@ class SparseType(enum.Enum):
             or self.value == SparseType.FP16.value
             or self.value == SparseType.FP8.value
             or self.value == SparseType.BF16.value
+            or self.value == SparseType.NFP8.value
         ):
             return True
         else:
@@ -193,11 +476,12 @@ class SparseType(enum.Enum):
             return QuantizationConfig()
 
 
-ELEMENT_SIZE: Dict[SparseType, int] = {
+ELEMENT_SIZE: dict[SparseType, int] = {
     SparseType.FP32: 4,
     SparseType.FP16: 2,
     SparseType.FP8: 1,
     SparseType.INT8: 1,
     SparseType.BF16: 2,
+    SparseType.NFP8: 1,
     # SparseType.INT4: 0.5,
 }

fbgemm_gpu/split_embedding_inference_converter.py

@@ -10,10 +10,11 @@
 
 import logging
 import math
-from typing import cast, Optional, Tuple
+from typing import cast, Optional
 
 import torch
 
+# fmt:skip
 from fbgemm_gpu.split_embedding_configs import (
     FP8QuantizationConfig,
     QuantizationConfig,
@@ -53,7 +54,7 @@ class SplitEmbInferenceConverter:
         return model
 
     # pyre-fixme[2]: Parameter must be annotated.
-    def _prune_by_weights_l2_norm(self, new_num_rows, weights) -> Tuple[Tensor, float]:
+    def _prune_by_weights_l2_norm(self, new_num_rows, weights) -> tuple[Tensor, float]:
         assert new_num_rows > 0
         from numpy.linalg import norm
 
@@ -75,7 +76,7 @@ class SplitEmbInferenceConverter:
         idx: int,
         num_rows: int,
         module: SplitTableBatchedEmbeddingBagsCodegen,
-    ) -> Tuple[Tensor, Optional[Tensor]]:
+    ) -> tuple[Tensor, Optional[Tensor]]:
         # TODO(yingz): Avoid DtoH / HtoD overhead.
         weights = module.split_embedding_weights()[idx].cpu()
         if self.pruning_ratio is None:
@@ -84,7 +85,7 @@ class SplitEmbInferenceConverter:
         if new_num_rows == num_rows:
             return (weights, None)
 
-        (indicators, threshold) = self._prune_by_weights_l2_norm(new_num_rows, weights)
+        indicators, threshold = self._prune_by_weights_l2_norm(new_num_rows, weights)
 
         return torch.ops.fbgemm.embedding_bag_rowwise_prune(
             weights, indicators, threshold, torch.int32
@@ -100,7 +101,7 @@ class SplitEmbInferenceConverter:
 
     def _quantize_embs(
         self, weight: Tensor, weight_ty: SparseType
-    ) -> Tuple[Tensor, Optional[Tensor]]:
+    ) -> tuple[Tensor, Optional[Tensor]]:
         fp8_quant_config = cast(FP8QuantizationConfig, self.quantization_config)
         return quantize_embs(weight, weight_ty, fp8_quant_config)
 
@@ -129,7 +130,7 @@ class SplitEmbInferenceConverter:
                 index_remapping_list = []
                 for t, (_, E, D, weight_ty) in enumerate(embedding_specs):
                     # Try to prune embeddings.
-                    (pruned_weight, index_remapping) = self._prune_embs(t, E, child)
+                    pruned_weight, index_remapping = self._prune_embs(t, E, child)
                     new_embedding_specs.append(
                         (
                             "",

fbgemm_gpu/split_embedding_optimizer_codegen/optimizer_args.py

@@ -4,6 +4,8 @@
 ## Template Source: training/python/optimizer_args.py
 ################################################################################
 
+__template_source_file__ = "training/python/optimizer_args.py"
+
 #!/usr/bin/env python3
 # Copyright (c) Meta Platforms, Inc. and affiliates.
 # All rights reserved.

fbgemm_gpu/split_embedding_optimizer_codegen/split_embedding_optimizer_rowwise_adagrad.py

@@ -4,6 +4,8 @@
 ## Template Source: training/python/split_embedding_optimizer_codegen.template
 ################################################################################
 
+__template_source_file__ = "training/python/split_embedding_optimizer_codegen.template"
+
 #!/usr/bin/env python3
 
 # Copyright (c) Meta Platforms, Inc. and affiliates.