Support get/set the whole row of metaheader+weight+optimizer from backend for checkpoint saving/loading (pytorch#4429)

Summary:
Pull Request resolved: pytorch#4429

X-link: facebookresearch/FBGEMM#1495

X-link: meta-pytorch/torchrec#3148

# Context
In our current KVZCH cp loading flow, we will keep hold of weight_id, weight, optimizer tensors throughout the checkpoint loading lifecycle, and at the end when all these tensors are downloaded in hand, we will explicitly call "apply_state_dict" to actually write them by chunk to the backend to ensure id->weight and id->opt are mapped correctly. The problem is when we have large number of weights, we will be short of memory since we need to hold all 3 tensors (double memory issue). To solve this challenge, we are going to save the whole row of (metaheader + weight + opt) as the same "weight" tensor during checkpoint saving, and when downloading the checkpoint, we will be able to extract the id from the header, and directly write the weight+opt part to the backend by id. When loading cp for optimizer, we added a no-op KVTensor, so it won't need to write to backend for optimizer states again.

# This diff
* added `backend_return_whole_row` flag in KVZCH params, with validation to make sure it's only True when opt_offloading is used
* added `read_only_` flag in KVTensorWrapper to be used for checkpoint calls. When read-only=True, all write operations to this KVT will be no-op
* added metadata recalc for optimizer state dict, because we are now returning read-only KVT for opt state dict, and model store will need to correct the global metadata before creating the save plan for KVZCH opt tensors
* updated dram backend and mem pool, so it can return the metaheader + weight + optimizer_state together, as well as set them back to backend (use pointers to skip metaheader part when write weight+opt to backend)
* by default the opt offloading and return whole row is False on trunk, so should not break existing KVZCH runs

Differential Revision: D77604158

Author: Jianbo Liu

fbgemm_gpu/fbgemm_gpu/split_table_batched_embeddings_ops_common.py

@@ -106,13 +106,19 @@ class KVZCHParams(NamedTuple):
     bucket_sizes: List[int] = []
     # enable optimizer offloading or not
     enable_optimizer_offloading: bool = False
+    # when enabled, backend will return whole row(metaheader + weight + optimizer) instead of weight only
+    # can only be enabled when enable_optimizer_offloading is enabled
+    backend_return_whole_row: bool = False
     eviction_policy: Optional[EvictionPolicy] = None
 
     def validate(self) -> None:
         assert len(self.bucket_offsets) == len(self.bucket_sizes), (
             "bucket_offsets and bucket_sizes must have the same length, "
             f"actual {self.bucket_offsets} vs {self.bucket_sizes}"
         )
+        assert (
+            not self.backend_return_whole_row or self.enable_optimizer_offloading
+        ), "backend_return_whole_row can only be enabled when enable_optimizer_offloading is enabled"
 
 
 class BackendType(enum.IntEnum):

fbgemm_gpu/fbgemm_gpu/tbe/ssd/training.py

@@ -187,15 +187,30 @@ def __init__(
         self.kv_zch_params = kv_zch_params
         self.backend_type = backend_type
         self.enable_optimizer_offloading: bool = False
+        self.backend_return_whole_row: bool = False
         if self.kv_zch_params:
             self.kv_zch_params.validate()
             self.enable_optimizer_offloading = (
                 # pyre-ignore [16]
                 self.kv_zch_params.enable_optimizer_offloading
             )
+            self.backend_return_whole_row = (
+                # pyre-ignore [16]
+                self.kv_zch_params.backend_return_whole_row
+            )
 
             if self.enable_optimizer_offloading:
                 logging.info("Optimizer state offloading is enabled")
+            if self.backend_return_whole_row:
+                assert (
+                    self.backend_type == BackendType.DRAM
+                ), f"Only DRAM backend supports backend_return_whole_row, but got {self.backend_type}"
+                logging.info(
+                    "Backend will return whole row including metaheader, weight and optimizer for checkpoint"
+                )
+
+        # TODO: the metaheader is 16 bytes fixed.
+        self.metaheader_dim: int = 16 // (weights_precision.bit_rate() // 8)
 
         self.pooling_mode = pooling_mode
         self.bounds_check_mode_int: int = bounds_check_mode.value
@@ -612,13 +627,14 @@ def __init__(
             logging.info(
                 f"Logging DRAM offloading setup, tbe_unique_id:{tbe_unique_id}, l2_cache_size:{l2_cache_size}GB,"
                 f"num_shards={ssd_rocksdb_shards},num_threads={ssd_rocksdb_shards},"
-                f"max_D={self.max_D}"
+                f"max_D={self.max_D},"
                 f"uniform_init_lower={ssd_uniform_init_lower},uniform_init_upper={ssd_uniform_init_upper},"
                 f"row_storage_bitwidth={weights_precision.bit_rate()},"
                 f"self.cache_row_dim={self.cache_row_dim},"
                 f"enable_optimizer_offloading={self.enable_optimizer_offloading},"
                 f"feature_dims={self.feature_dims},"
-                f"hash_size_cumsum={self.hash_size_cumsum}"
+                f"hash_size_cumsum={self.hash_size_cumsum},"
+                f"backend_return_whole_row={self.backend_return_whole_row}"
             )
             table_dims = (
                 tensor_pad4(self.table_dims)
@@ -659,6 +675,7 @@ def __init__(
                     if self.enable_optimizer_offloading
                     else None
                 ),  # hash_size_cumsum
+                self.backend_return_whole_row,  # backend_return_whole_row
             )
         else:
             raise AssertionError(f"Invalid backend type {self.backend_type}")
@@ -2246,16 +2263,19 @@ def split_optimizer_states(
             # pyre-ignore
             bucket_size = self.kv_zch_params.bucket_sizes[t]
             row_offset = table_offset
-            if sorted_id_tensor is None or sorted_id_tensor[t].numel() == 0:
+            if not self.backend_return_whole_row and (
+                sorted_id_tensor is None or sorted_id_tensor[t].numel() == 0
+            ):
                 opt_list.append(
                     torch.empty(0, dtype=self.optimizer.dtype(), device="cpu")
                     # empty optimizer state for module initialization
+                    # which will NOT be used for cp loading
                 )
             else:
                 if not self.enable_optimizer_offloading:
                     # convert global id back to local id, then linearize with table offset
                     local_id_tensor = (
-                        sorted_id_tensor[t]
+                        sorted_id_tensor[t]  # pyre-ignore[16]
                         - bucket_id_start * bucket_size
                         + table_offset
                     )
@@ -2264,27 +2284,74 @@ def split_optimizer_states(
                     )
                 else:
                     row_offset = table_offset - (bucket_id_start * bucket_size)
-                    # using KVTensorWrapper to query backend to avoid OOM memory, since
-                    # backend will return both weight and optimizer in one tensor, read the whole tensor
-                    # out could OOM CPU memory.
-                    tensor_wrapper = torch.classes.fbgemm.KVTensorWrapper(
-                        shape=[emb_height, optimizer_dim],
-                        dtype=dtype,
-                        row_offset=row_offset,
-                        snapshot_handle=snapshot_handle,
-                        sorted_indices=sorted_id_tensor[t],
-                        width_offset=pad4(emb_dim),
-                    )
-                    (
-                        tensor_wrapper.set_embedding_rocks_dp_wrapper(self.ssd_db)
-                        if self.backend_type == BackendType.SSD
-                        else tensor_wrapper.set_dram_db_wrapper(self.ssd_db)
-                    )
-                    opt_list.append(
-                        self.get_offloaded_optimizer_states(
-                            tensor_wrapper, sorted_id_tensor[t].numel()
+                    if self.backend_return_whole_row:
+                        # When backend returns whole row, the optimizer will be returned as PMT directly
+                        if (
+                            sorted_id_tensor[t].size(0) == 0
+                            and self.local_weight_counts[t] > 0
+                        ):
+                            logging.info(
+                                f"before opt PMT loading, resetting id tensor with {self.local_weight_counts[t]}"
+                            )
+                            # pyre-ignore [16]
+                            sorted_id_tensor[t] = torch.zeros(
+                                (self.local_weight_counts[t], 1),
+                                device=torch.device("cpu"),
+                                dtype=torch.int64,
+                            )
+                        tensor_wrapper = torch.classes.fbgemm.KVTensorWrapper(
+                            shape=[
+                                (
+                                    sorted_id_tensor[t].size(0)
+                                    if sorted_id_tensor is not None
+                                    and sorted_id_tensor[t].size(0) > 0
+                                    else emb_height
+                                ),
+                                optimizer_dim,
+                            ],
+                            dtype=dtype,
+                            row_offset=row_offset,
+                            snapshot_handle=snapshot_handle,
+                            sorted_indices=sorted_id_tensor[t],
+                            width_offset=(
+                                self.metaheader_dim  # metaheader is already padded so no need for pad4
+                                + pad4(emb_dim)
+                            ),
+                            read_only=True,  # optimizer written to DB with weights, so skip write here
+                        )
+                        (
+                            tensor_wrapper.set_embedding_rocks_dp_wrapper(self.ssd_db)
+                            if self.backend_type == BackendType.SSD
+                            else tensor_wrapper.set_dram_db_wrapper(self.ssd_db)
+                        )
+                        opt_list.append(
+                            PartiallyMaterializedTensor(
+                                tensor_wrapper,
+                                True if self.kv_zch_params else False,
+                            )
+                        )
+                    else:
+                        # using KVTensorWrapper to query backend to avoid OOM memory, since
+                        # backend will return both weight and optimizer in one tensor, read the whole tensor
+                        # out could OOM CPU memory.
+                        tensor_wrapper = torch.classes.fbgemm.KVTensorWrapper(
+                            shape=[emb_height, optimizer_dim],
+                            dtype=dtype,
+                            row_offset=row_offset,
+                            snapshot_handle=snapshot_handle,
+                            sorted_indices=sorted_id_tensor[t],
+                            width_offset=pad4(emb_dim),
+                        )
+                        (
+                            tensor_wrapper.set_embedding_rocks_dp_wrapper(self.ssd_db)
+                            if self.backend_type == BackendType.SSD
+                            else tensor_wrapper.set_dram_db_wrapper(self.ssd_db)
+                        )
+                        opt_list.append(
+                            self.get_offloaded_optimizer_states(
+                                tensor_wrapper, sorted_id_tensor[t].numel()
+                            )
                         )
-                    )
             table_offset += emb_height
         logging.info(
             f"KV ZCH tables split_optimizer_states query latency: {(time.time() - start_time) * 1000} ms, "
@@ -2513,7 +2580,7 @@ def split_embedding_weights(
                     and self.local_weight_counts[i] > 0
                 ):
                     logging.info(
-                        f"resetting bucket id tensor with {self.local_weight_counts[i]}"
+                        f"before weight PMT loading, resetting id tensor with {self.local_weight_counts[i]}"
                     )
                     bucket_ascending_id_tensor = torch.zeros(
                         (self.local_weight_counts[i], 1),
@@ -2539,7 +2606,19 @@ def split_embedding_weights(
                         if bucket_ascending_id_tensor is not None
                         else emb_height
                     ),
-                    emb_dim,
+                    (
+                        (
+                            self.metaheader_dim  # metaheader is already padded
+                            + pad4(emb_dim)
+                            + pad4(
+                                self.optimizer.state_size_dim(
+                                    self.weights_precision.as_dtype()
+                                )
+                            )
+                        )
+                        if self.backend_return_whole_row
+                        else emb_dim
+                    ),
                 ],
                 dtype=dtype,
                 row_offset=row_offset,
@@ -2576,6 +2655,11 @@ def split_embedding_weights(
 
     @torch.jit.ignore
     def apply_state_dict(self) -> None:
+        if self.backend_return_whole_row:
+            logging.info(
+                "backend_return_whole_row is enabled, no need to apply_state_dict"
+            )
+            return
         # After checkpoint loading, the _cached_kvzch_data will be loaded from checkpoint.
         # Caller should call this function to apply the cached states to backend.
         if self.load_state_dict is False:
@@ -2694,6 +2778,11 @@ def streaming_write_weight_and_id_per_table(
 
     @torch.jit.ignore
     def enable_load_state_dict_mode(self) -> None:
+        if self.backend_return_whole_row:
+            logging.info(
+                "backend_return_whole_row is enabled, no need to enable load_state_dict mode"
+            )
+            return
         # Enable load state dict mode before loading checkpoint
         if self.load_state_dict:
             return