Skip load metadata tensor (#4830)

Summary:
X-link: meta-pytorch/torchrec#3359

Pull Request resolved: #4830

X-link: facebookresearch/FBGEMM#1856

The metadata tensor is newly added for kvzch table. Some old checkpoints may not have this fqn. Directly load old checkpoint can cause fqn missing error.

This diff try to skip init metadata tensor at load checkpoint func. Metadata tensor is not used in training, so it is okay to skip load. It will be created during saving checkpoint.

Reviewed By: steven1327, emlin

Differential Revision: D81811024

fbshipit-source-id: edd731b40c6a843b338cc0c9a7f4ffb55000b706

Author: EddyLXJ

fbgemm_gpu/fbgemm_gpu/tbe/ssd/training.py

@@ -3065,6 +3065,7 @@ def split_embedding_weights(
         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
+        skip_metadata = False
 
         table_offset = 0
         for i, (emb_height, emb_dim) in enumerate(self.embedding_specs):
@@ -3132,18 +3133,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,
-                    )
+                    skip_metadata = True
 
                     # 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)
+                if skip_metadata:
+                    metadata_splits = None
+                else:
+                    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

fbgemm_gpu/test/tbe/ssd/ssd_split_tbe_training_test.py

@@ -811,6 +811,16 @@ def test_apply_kv_state_dict(
 
         # pyre-fixme[16]: Undefined attribute: `Optional` has no attribute `__iter__`.
         emb2.local_weight_counts = [ids.numel() for ids in bucket_asc_ids_list]
+
+        (
+            _,
+            _,
+            _,
+            metadata_list,
+        ) = emb2.split_embedding_weights(no_snapshot=False, should_flush=True)
+
+        self.assertTrue(metadata_list is None)
+
         emb2.enable_load_state_dict_mode()
         self.assertIsNotNone(emb2._cached_kvzch_data)
         for i, _ in enumerate(emb.embedding_specs):
@@ -844,12 +854,14 @@ def test_apply_kv_state_dict(
             emb_state_dict_list2,
             bucket_asc_ids_list2,
             num_active_id_per_bucket_list2,
-            _,
+            metadata_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
         )
 
+        self.assertTrue(metadata_list2 is not None)
+
         for t in range(len(emb.embedding_specs)):
             sorted_ids = torch.sort(bucket_asc_ids_list[t].flatten())
             sorted_ids2 = torch.sort(bucket_asc_ids_list2[t].flatten())
@@ -881,6 +893,10 @@ def test_apply_kv_state_dict(
                 rtol=tolerance,
             )
 
+            self.assertTrue(
+                metadata_list2[t].size(0) == bucket_asc_ids_list2[t].size(0)
+            )
+
     def _check_raw_embedding_stream_call_counts(
         self,
         mock_raw_embedding_stream: unittest.mock.Mock,