Adding eviction metadata tensor fqn (pytorch#4611)

Summary:
X-link: pytorch/torchrec#3247

Pull Request resolved: pytorch#4611

X-link: facebookresearch/FBGEMM#1646

Adding a new metadata fqn in kvzch ckpt, which is needed for eviction filter in publishing.

Differential Revision: D78768842

Author: EddyLXJ

fbgemm_gpu/fbgemm_gpu/tbe/ssd/training.py

@@ -2783,6 +2783,7 @@ def split_embedding_weights(
         Union[List[PartiallyMaterializedTensor], List[torch.Tensor]],
         Optional[List[torch.Tensor]],
         Optional[List[torch.Tensor]],
+        Optional[List[torch.Tensor]],
     ]:
         """
         This method is intended to be used by the checkpointing engine
@@ -2802,6 +2803,7 @@ def split_embedding_weights(
             2nd arg: input id sorted in bucket id ascending order
             3rd arg: active id count per bucket id, tensor size is [bucket_id_end - bucket_id_start]
                     where for the i th element, we have i + bucket_id_start = global bucket id
+            4th arg: kvzch eviction metadata for each input id sorted in bucket id ascending order
         """
         snapshot_handle, checkpoint_handle = self._may_create_snapshot_for_state_dict(
             no_snapshot=no_snapshot,
@@ -2818,16 +2820,19 @@ def split_embedding_weights(
                 self._cached_kvzch_data.cached_weight_tensor_per_table,
                 self._cached_kvzch_data.cached_id_tensor_per_table,
                 self._cached_kvzch_data.cached_bucket_splits,
+                [],  # metadata tensor is not needed for checkpointing loading
             )
         start_time = time.time()
         pmt_splits = []
         bucket_sorted_id_splits = [] if self.kv_zch_params else None
         active_id_cnt_per_bucket_split = [] if self.kv_zch_params else None
+        metadata_splits = [] if self.kv_zch_params else None
 
         table_offset = 0
         for i, (emb_height, emb_dim) in enumerate(self.embedding_specs):
             bucket_ascending_id_tensor = None
             bucket_t = None
+            metadata_tensor = None
             row_offset = table_offset
             metaheader_dim = 0
             if self.kv_zch_params:
@@ -2859,6 +2864,12 @@ def split_embedding_weights(
                         bucket_size,
                     )
                 )
+                metadata_tensor = self._ssd_db.get_kv_zch_eviction_metadata_by_snapshot(
+                    bucket_ascending_id_tensor,
+                    torch.as_tensor(bucket_ascending_id_tensor.size(0)),
+                    snapshot_handle,
+                )
+
                 # 3. convert local id back to global id
                 bucket_ascending_id_tensor.add_(bucket_id_start * bucket_size)
 
@@ -2874,11 +2885,17 @@ def split_embedding_weights(
                         device=torch.device("cpu"),
                         dtype=torch.int64,
                     )
+                    metadata_tensor = torch.zeros(
+                        (self.local_weight_counts[i], 1),
+                        device=torch.device("cpu"),
+                        dtype=torch.int64,
+                    )
                     # self.local_weight_counts[i] = 0  # Reset the count
 
                 # pyre-ignore [16] bucket_sorted_id_splits is not None
                 bucket_sorted_id_splits.append(bucket_ascending_id_tensor)
                 active_id_cnt_per_bucket_split.append(bucket_t)
+                metadata_splits.append(metadata_tensor)
 
                 # for KV ZCH tbe, the sorted_indices is global id for checkpointing and publishing
                 # but in backend, local id is used during training, so the KVTensorWrapper need to convert global id to local id
@@ -2934,7 +2951,12 @@ def split_embedding_weights(
                 f"num ids list: {[ids.numel() for ids in bucket_sorted_id_splits]}"
             )
 
-        return (pmt_splits, bucket_sorted_id_splits, active_id_cnt_per_bucket_split)
+        return (
+            pmt_splits,
+            bucket_sorted_id_splits,
+            active_id_cnt_per_bucket_split,
+            metadata_splits,
+        )
 
     @torch.jit.ignore
     def _apply_state_dict_w_offloading(self) -> None:

fbgemm_gpu/test/tbe/ssd/ssd_split_tbe_training_test.py

@@ -788,7 +788,7 @@ def execute_ssd_backward_(
     def split_optimizer_states_(
         self, emb: SSDTableBatchedEmbeddingBags
     ) -> List[List[torch.Tensor]]:
-        _, bucket_asc_ids_list, _ = emb.split_embedding_weights(
+        _, bucket_asc_ids_list, _, _ = emb.split_embedding_weights(
             no_snapshot=False, should_flush=True
         )
 
@@ -1113,7 +1113,7 @@ def test_ssd_emb_state_dict(
         split_optimizer_states = self.split_optimizer_states_(emb)
 
         # Compare emb state dict with expected values from nn.EmbeddingBag
-        emb_state_dict, _, _ = emb.split_embedding_weights(no_snapshot=False)
+        emb_state_dict, _, _, _ = emb.split_embedding_weights(no_snapshot=False)
         for feature_index, table_index in self.get_physical_table_arg_indices_(
             emb.feature_table_map
         ):
@@ -1728,9 +1728,12 @@ def test_kv_emb_state_dict(
         split_optimizer_states = []
 
         # Compare emb state dict with expected values from nn.EmbeddingBag
-        emb_state_dict_list, bucket_asc_ids_list, num_active_id_per_bucket_list = (
-            emb.split_embedding_weights(no_snapshot=False, should_flush=True)
-        )
+        (
+            emb_state_dict_list,
+            bucket_asc_ids_list,
+            num_active_id_per_bucket_list,
+            metadata_list,
+        ) = emb.split_embedding_weights(no_snapshot=False, should_flush=True)
 
         for s in emb.split_optimizer_states(
             bucket_asc_ids_list, no_snapshot=False, should_flush=True
@@ -1797,6 +1800,7 @@ def test_kv_emb_state_dict(
             )
             self.assertLess(table_index, len(emb_state_dict_list))
             assert len(split_optimizer_states[table_index][0]) == num_ids
+            assert len(metadata_list[table_index]) == num_ids
             # NOTE: The [0] index is a hack since the test is fixed to use
             # EXACT_ROWWISE_ADAGRAD optimizer.  The test in general should
             # be upgraded in the future to support multiple optimizers
@@ -1943,7 +1947,7 @@ def test_kv_opt_state_w_offloading(
         )
 
         # Compare emb state dict with expected values from nn.EmbeddingBag
-        emb_state_dict_list, bucket_asc_ids_list, num_active_id_per_bucket_list = (
+        emb_state_dict_list, bucket_asc_ids_list, num_active_id_per_bucket_list, _ = (
             emb.split_embedding_weights(no_snapshot=False, should_flush=True)
         )
         split_optimizer_states = emb.split_optimizer_states(
@@ -2172,7 +2176,7 @@ def test_kv_state_dict_w_backend_return_whole_row(
         )
 
         # Compare emb state dict with expected values from nn.EmbeddingBag
-        emb_state_dict_list, bucket_asc_ids_list, num_active_id_per_bucket_list = (
+        emb_state_dict_list, bucket_asc_ids_list, num_active_id_per_bucket_list, _ = (
             emb.split_embedding_weights(no_snapshot=False, should_flush=True)
         )
         split_optimizer_states = emb.split_optimizer_states(
@@ -2440,7 +2444,7 @@ def test_apply_kv_state_dict(
         )
 
         # Compare emb state dict with expected values from nn.EmbeddingBag
-        emb_state_dict_list, bucket_asc_ids_list, num_active_id_per_bucket_list = (
+        emb_state_dict_list, bucket_asc_ids_list, num_active_id_per_bucket_list, _ = (
             emb.split_embedding_weights(no_snapshot=False, should_flush=True)
         )
         split_optimizer_states = emb.split_optimizer_states(
@@ -2508,6 +2512,7 @@ def test_apply_kv_state_dict(
             emb_state_dict_list2,
             bucket_asc_ids_list2,
             num_active_id_per_bucket_list2,
+            _,
         ) = emb2.split_embedding_weights(no_snapshot=False, should_flush=True)
         split_optimizer_states2 = emb2.split_optimizer_states(
             bucket_asc_ids_list2, no_snapshot=False, should_flush=True
@@ -2963,7 +2968,7 @@ def copy_opt_states_hook(
         emb.flush()
 
         # Compare emb state dict with expected values from nn.EmbeddingBag
-        _emb_state_dict_list, bucket_asc_ids_list, _num_active_id_per_bucket_list = (
+        _emb_state_dict_list, bucket_asc_ids_list, _num_active_id_per_bucket_list, _ = (
             emb.split_embedding_weights(no_snapshot=False, should_flush=True)
         )
         assert bucket_asc_ids_list is not None