fbgemm-gpu-hstu-nightly 2025.6.13__cp312-cp312-manylinux_2_28_x86_64.whl → 2025.6.15__cp312-cp312-manylinux_2_28_x86_64.whl

fbgemm_gpu/docs/version.py

@@ -6,6 +6,6 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
-__version__: str = "2025.6.13"
+__version__: str = "2025.6.15"
 __target__: str = "hstu"
 __variant__: str = "cuda"

fbgemm_gpu/split_table_batched_embeddings_ops_common.py

@@ -33,6 +33,17 @@ class EmbeddingLocation(enum.IntEnum):
     HOST = 3
     MTIA = 4
 
+    @classmethod
+    # pyre-ignore[3]
+    def str_values(cls):
+        return [
+            "device",
+            "managed",
+            "managed_caching",
+            "host",
+            "mtia",
+        ]
+
     @classmethod
     # pyre-ignore[3]
     def from_str(cls, key: str):

fbgemm_gpu/tbe/bench/embedding_ops_common_config.py

@@ -99,9 +99,7 @@ class EmbeddingOpsCommonConfigLoader:
             click.option(
                 "--emb-location",
                 default="device",
-                type=click.Choice(
-                    ["device", "managed", "managed_caching"], case_sensitive=False
-                ),
+                type=click.Choice(EmbeddingLocation.str_values(), case_sensitive=False),
                 help="Memory location of the embeddings",
             ),
             click.option(

fbgemm_gpu/tbe/ssd/training.py

@@ -2804,7 +2804,12 @@ class SSDTableBatchedEmbeddingBags(nn.Module):
         """
         Create a rocksdb hard link snapshot to provide cross procs access to the underlying data
         """
-        self.ssd_db.create_rocksdb_hard_link_snapshot(self.step)
+        if self.backend_type == BackendType.SSD:
+            self.ssd_db.create_rocksdb_hard_link_snapshot(self.step)
+        else:
+            logging.warning(
+                "create_rocksdb_hard_link_snapshot is only supported for SSD backend"
+            )
 
     def prepare_inputs(
         self,

fbgemm_gpu/tbe/ssd/utils/partially_materialized_tensor.py

@@ -9,7 +9,6 @@
 from __future__ import annotations
 
 import functools
-import logging
 from typing import List, Optional, Union
 
 import torch
@@ -26,58 +25,6 @@ def implements(torch_function):
     return decorator
 
 
-class KVTensorMetadata:
-    """
-    Class that is used to represent a KVTensor as a Serialized Metadata in python
-    This object is used to reconstruct the KVTensor in the publish component
-    """
-
-    checkpoint_paths: List[str]
-    tbe_uuid: str
-    rdb_num_shards: int
-    rdb_num_threads: int
-    max_D: int
-    table_offset: int
-    table_shape: List[int]
-    dtype: int
-    checkpoint_uuid: str
-
-    def __init__(
-        self,
-        checkpoint_paths: List[str],
-        tbe_uuid: str,
-        rdb_num_shards: int,
-        rdb_num_threads: int,
-        max_D: int,
-        table_offset: int,
-        table_shape: List[int],
-        dtype: int,
-        checkpoint_uuid: str,
-    ) -> None:
-        """
-        Ensure caller loads the module before creating this object.
-
-        ```
-        load_torch_module(
-            "//deeplearning/fbgemm/fbgemm_gpu:ssd_split_table_batched_embeddings"
-        )
-        ```
-
-        Args:
-
-            wrapped: torch.classes.fbgemm.KVTensorWrapper
-        """
-        self.checkpoint_paths = checkpoint_paths
-        self.tbe_uuid = tbe_uuid
-        self.rdb_num_shards = rdb_num_shards
-        self.rdb_num_threads = rdb_num_threads
-        self.max_D = max_D
-        self.table_offset = table_offset
-        self.table_shape = table_shape
-        self.checkpoint_uuid = checkpoint_uuid
-        self.dtype = dtype
-
-
 class PartiallyMaterializedTensor:
     """
     A tensor-like object that represents a partially materialized tensor in memory.
@@ -104,55 +51,6 @@ class PartiallyMaterializedTensor:
         self._is_virtual = is_virtual
         self._requires_grad = False
 
-    @property
-    def generate_kvtensor_metadata(self) -> KVTensorMetadata:
-        serialized_metadata = self.wrapped.get_kvtensor_serializable_metadata()
-        try:
-            metadata_itr = 0
-            num_rdb_ckpts = int(serialized_metadata[0])
-            metadata_itr += 1
-            checkpoint_paths: List[str] = []
-            for i in range(num_rdb_ckpts):
-                checkpoint_paths.append(serialized_metadata[i + metadata_itr])
-            metadata_itr += num_rdb_ckpts
-            tbe_uuid = serialized_metadata[metadata_itr]
-            metadata_itr += 1
-            rdb_num_shards = int(serialized_metadata[metadata_itr])
-            metadata_itr += 1
-            rdb_num_threads = int(serialized_metadata[metadata_itr])
-            metadata_itr += 1
-            max_D = int(serialized_metadata[metadata_itr])
-            metadata_itr += 1
-            table_offset = int(serialized_metadata[metadata_itr])
-            metadata_itr += 1
-            table_shape: List[int] = []
-            table_shape.append(int(serialized_metadata[metadata_itr]))
-            metadata_itr += 1
-            table_shape.append(int(serialized_metadata[metadata_itr]))
-            metadata_itr += 1
-            dtype = int(serialized_metadata[metadata_itr])
-            metadata_itr += 1
-            checkpoint_uuid = serialized_metadata[metadata_itr]
-            metadata_itr += 1
-            res = KVTensorMetadata(
-                checkpoint_paths,
-                tbe_uuid,
-                rdb_num_shards,
-                rdb_num_threads,
-                max_D,
-                table_offset,
-                table_shape,
-                dtype,
-                checkpoint_uuid,
-            )
-
-            return res
-        except Exception as e:
-            logging.error(
-                f"Failed to parse metadata: {e}, here is metadata: {serialized_metadata}"
-            )
-            raise e
-
     @property
     def wrapped(self):
         """

{fbgemm_gpu_hstu_nightly-2025.6.13.dist-info → fbgemm_gpu_hstu_nightly-2025.6.15.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: fbgemm_gpu_hstu_nightly
-Version: 2025.6.13
+Version: 2025.6.15
 Home-page: https://github.com/pytorch/fbgemm
 Author: FBGEMM Team
 Author-email: packages@pytorch.org

{fbgemm_gpu_hstu_nightly-2025.6.13.dist-info → fbgemm_gpu_hstu_nightly-2025.6.15.dist-info}/RECORD

@@ -15,7 +15,7 @@ fbgemm_gpu/split_embedding_inference_converter.py,sha256=ilVVowkTiY0WDpOYorj917T
 fbgemm_gpu/split_embedding_optimizer_ops.py,sha256=wXuGazClBMk62yL_r9udUIKaPgQP7SlkSb5ugB75wrQ,711
 fbgemm_gpu/split_embedding_utils.py,sha256=Gb40ZKeATxIKEKI3aVQMgDDBanNpKMc53Z43mnzdR_I,851
 fbgemm_gpu/split_table_batched_embeddings_ops.py,sha256=_MIp6uHYHLn4GxGdrGsfddfSsZ2Z9mjsYIrih3ncI1I,2339
-fbgemm_gpu/split_table_batched_embeddings_ops_common.py,sha256=jRCuFRUPchgqL9Z9RTEIjZGl9R6YJe1RlU_Su7AgCMw,8030
+fbgemm_gpu/split_table_batched_embeddings_ops_common.py,sha256=qglNRKKuHkrKiTw90ACjZpMzcjHKXKV7ME3a8QHfQt4,8237
 fbgemm_gpu/split_table_batched_embeddings_ops_inference.py,sha256=bUDWa6IR0vGLDThgB3nmD1yfYa8_HD34B0dtLnd7thw,81692
 fbgemm_gpu/split_table_batched_embeddings_ops_training.py,sha256=YCLPSW9CXrRwMN5KEU6x0ESbutdhzKTaNOO8oN5kX7I,163875
 fbgemm_gpu/split_table_batched_embeddings_ops_training_common.py,sha256=ktC10-nakOBpcmJNCOGQsxuBCP8XTwXJ2WeEgIg91tc,5455
@@ -32,10 +32,10 @@ fbgemm_gpu/docs/merge_pooled_embedding_ops.py,sha256=oJLgSgZQmhsyGLbTmZTxNgQrk65
 fbgemm_gpu/docs/permute_pooled_embedding_ops.py,sha256=tZUqLVXlk5O6VAKKDA-OEMx2fCu5QPOOeoAPZA9_nLY,4454
 fbgemm_gpu/docs/quantize_ops.py,sha256=xTtOaVK1P02ymreE_i21YiyYDZCqhoZY9eWp_mEIRlo,1297
 fbgemm_gpu/docs/sparse_ops.py,sha256=NTcTm0q9h8W2B8PKPoic2fHsAaCbCYunSa_EYK0LtHQ,21382
-fbgemm_gpu/docs/version.py,sha256=9kytZ-RJTwv5Z4Lajmjf4ELvL28p4wC2CGW1Zv5ZfSg,315
+fbgemm_gpu/docs/version.py,sha256=Wfdofi1dDfFLcvCYRMDlfYaAgz8J7SIMZPIBXTmCxcA,315
 fbgemm_gpu/experimental/hstu/__init__.py,sha256=KNisP6qDMwgjgxkGlqUZRNjJ_8o8R-cTmm3HxF7pSqI,1564
 fbgemm_gpu/experimental/hstu/cuda_hstu_attention.py,sha256=5425GRjJuzpXQC-TowgQOCFjZmOwv_EK0lKbURhHBTQ,9920
-fbgemm_gpu/experimental/hstu/fbgemm_gpu_experimental_hstu.so,sha256=WxrpQwaZYcuLBp7wHQlsF5It8r5J5VSavTl6X1vWRSU,352696288
+fbgemm_gpu/experimental/hstu/fbgemm_gpu_experimental_hstu.so,sha256=AdQzuwtVkIlxjhXCuWHjla0czOJXEuqM6CAWbiBmocU,352696288
 fbgemm_gpu/quantize/__init__.py,sha256=pftciXHE7csekDFkl7Ui1AWglVMMnSrOO04mREnUdb0,921
 fbgemm_gpu/quantize/quantize_ops.py,sha256=25AIOv9n2UoxamMUaI6EK1Ur4gSHxbZIReHBtgOjjCs,2228
 fbgemm_gpu/sll/__init__.py,sha256=rgXh35-OFUE54E9gGBq3NGxouGvgMv2ccY2bWUTxONY,4191
@@ -61,7 +61,7 @@ fbgemm_gpu/tbe/bench/__init__.py,sha256=Uhk3IzAs9f_2kVYcMuFP3KiwwkqqYxaZFivPaBIN
 fbgemm_gpu/tbe/bench/bench_config.py,sha256=GYxVotMi-JK7AEm96sOfB2uOlmRtXpOTr-tkPEB1Q48,4241
 fbgemm_gpu/tbe/bench/bench_runs.py,sha256=hmoldOLh2UsjPAeb0zdfm5aTbs_0iEnlCxsXCHuzAMg,18343
 fbgemm_gpu/tbe/bench/eeg_cli.py,sha256=T8Wa1PeRyFZ0Ge-SErHQEYDY8LvHVoCV_qQlE_6kER0,3600
-fbgemm_gpu/tbe/bench/embedding_ops_common_config.py,sha256=yBk0zaZJ5b3ffPG91SoqD088shrLYdyCFGfn6DsXK9A,5021
+fbgemm_gpu/tbe/bench/embedding_ops_common_config.py,sha256=mdG3JZwgclp6DiVwQSKl5jrirLSId4OuM64knj9TkEk,4973
 fbgemm_gpu/tbe/bench/eval_compression.py,sha256=bINVERk42VJDSdenQHKWApmRMrW8rhkevOgE0hDR-S8,3499
 fbgemm_gpu/tbe/bench/reporter.py,sha256=ZK5RFolUmZEcsEaife270_iOdXAQD5EjTUkuxctnAbY,804
 fbgemm_gpu/tbe/bench/tbe_data_config.py,sha256=uvepBrbzERALBB-RPZVGFra4a8ALCqsOe9X6iWpqAyU,9413
@@ -73,9 +73,9 @@ fbgemm_gpu/tbe/cache/split_embeddings_cache_ops.py,sha256=mQkCl0xN8xUu5bjEWcOOFN
 fbgemm_gpu/tbe/ssd/__init__.py,sha256=wzfMT10cp_dqK2lrebC449hOdexBnizcf_98lA1NyHs,483
 fbgemm_gpu/tbe/ssd/common.py,sha256=1J8K7sTQswgCYWaVwF-ZdCJj7mNN6O9GI70AaZWzJGE,1044
 fbgemm_gpu/tbe/ssd/inference.py,sha256=DTjwj3f6JaUMcecWoRNkZpRgXDJ-eE3grtixYwKb5DI,22829
-fbgemm_gpu/tbe/ssd/training.py,sha256=gCvMY8fGP2JNosMyzW7uiQUuUsDIdUULB_PG51wOpRk,131519
+fbgemm_gpu/tbe/ssd/training.py,sha256=Dx-rJqjrD1A4U4MEVaP3OJl3CZz0VRSTWcukx5557Jw,131715
 fbgemm_gpu/tbe/ssd/utils/__init__.py,sha256=5DgmR2HA6NtmYh2ddkUgpDsZ6a7hF0DPedA1gMpdh18,250
-fbgemm_gpu/tbe/ssd/utils/partially_materialized_tensor.py,sha256=ePF2b3vCFB_ZzPaMneaXe5gODFVRJmXKbagT2uNKAAQ,10934
+fbgemm_gpu/tbe/ssd/utils/partially_materialized_tensor.py,sha256=uwwEdUiaVlnWZ_rQax2z28VYROfivdMqIdWLy8IZ6cE,7646
 fbgemm_gpu/tbe/stats/__init__.py,sha256=on29iDtq7cVNh90JR9aeFNG-K9DDoYq0JryzoplL49I,322
 fbgemm_gpu/tbe/stats/bench_params_reporter.py,sha256=7XIWVObJOxSVUG73xsd_lVSuCFUQkMEGSWW--BoyCH0,7358
 fbgemm_gpu/tbe/utils/__init__.py,sha256=rlXFm-kTByFZO4SS5C5zMzANRiQmM1NT__eWBayncYg,549
@@ -93,7 +93,7 @@ fbgemm_gpu/utils/__init__.py,sha256=JQQNdcTTaEU6ptK-OW-ZQBwTFxEZZpWOtBXWwEZm39o,
 fbgemm_gpu/utils/filestore.py,sha256=Zshw1dA03m9aHMMAtETdq4bgOLocyLhzlkAUoG8VkdM,4743
 fbgemm_gpu/utils/loader.py,sha256=1hCEhNvkflniH46fGcrguLeP1z-6uyOu2QFwqKU5CIM,990
 fbgemm_gpu/utils/torch_library.py,sha256=dQcHv1qgpu5QYlJjxjd6oeHjtxnmmXzx3PL6vjCmxL4,4199
-fbgemm_gpu_hstu_nightly-2025.6.13.dist-info/METADATA,sha256=G6aLh_w9SYySzT7xRYMUj8OW6l873R2B71m-tgD6GWs,2794
-fbgemm_gpu_hstu_nightly-2025.6.13.dist-info/WHEEL,sha256=vUT1hK8fT5m5CAs5kDyQ_ABrvCmtd0TCp5-4vN9tR5A,108
-fbgemm_gpu_hstu_nightly-2025.6.13.dist-info/top_level.txt,sha256=2tlbTWLkPjhqvLF_6BbqKzkcPluSE-oPRVjI8axK76I,11
-fbgemm_gpu_hstu_nightly-2025.6.13.dist-info/RECORD,,
+fbgemm_gpu_hstu_nightly-2025.6.15.dist-info/METADATA,sha256=SxliCSqrubFUeLb4Kc97w7fTVCFw0mTOAK4PtWGAw6U,2794
+fbgemm_gpu_hstu_nightly-2025.6.15.dist-info/WHEEL,sha256=vUT1hK8fT5m5CAs5kDyQ_ABrvCmtd0TCp5-4vN9tR5A,108
+fbgemm_gpu_hstu_nightly-2025.6.15.dist-info/top_level.txt,sha256=2tlbTWLkPjhqvLF_6BbqKzkcPluSE-oPRVjI8axK76I,11
+fbgemm_gpu_hstu_nightly-2025.6.15.dist-info/RECORD,,