[Auto Parallel] fix save load state_dict (#66266)

xingmingyyj · zhangbo9674 · web-flow · commit 40f7c51089a9 · 2024-07-23T13:53:59.000+08:00
* fix

* fix

* fix

* fix

* fix

---------

Co-authored-by: zhangbo9674 &lt;zhangbo54@baidu.com&gt;
diff --git a/python/paddle/distributed/checkpoint/load_state_dict.py b/python/paddle/distributed/checkpoint/load_state_dict.py
@@ -32,6 +32,7 @@
 class ReadItem:
     local_tensor_index: LocalTensorIndex
     rank: int
+    dtype: str
     cur_offset: Tuple[int]
     storage_offset: Tuple[int]
     lengths: Tuple[int]
@@ -79,8 +80,21 @@ def get_rank_to_files(path, state_dict, process_group, use_dist):
             tensor_key_list.append(local_tensor_index.tensor_key)
             if local_tensor_index.tensor_key in state_dict:
                 necessary_files.append(file_name)
-    necessary_data_files_set = set(necessary_files)
-    if len(necessary_data_files_set) <= 0:
+
+    all_necessary_files = []
+    if use_dist:
+        paddle.distributed.all_gather_object(
+            all_necessary_files, necessary_files, process_group
+        )
+    else:
+        all_necessary_files.append(necessary_files)
+
+    global_necessary_files = [
+        file for files in all_necessary_files for file in files
+    ]
+
+    global_necessary_files_set = set(global_necessary_files)
+    if len(global_necessary_files_set) <= 0:
         logger.warning(
             f"No necessary data files found in the checkpoint directory:{path}. Please check the metadata_files:{metadata_files}"
         )
@@ -100,13 +114,13 @@ def get_rank_to_files(path, state_dict, process_group, use_dist):
         tmp += files
     global_data_files_set = set(tmp)
     logger.debug(
-        f"necessary_data_files_set:{necessary_data_files_set}, global_data_files_set:{global_data_files_set}"
+        f"necessary_data_files_set:{global_necessary_files_set}, global_data_files_set:{global_data_files_set}"
     )
     # check necessary files in global_data_files
     assert (
-        global_data_files_set & necessary_data_files_set
-        == necessary_data_files_set
-    ), f"The checkpoint files are not complete. Please check the checkpoint directory:{path}.global_data_files_set:{global_data_files_set}, necessary_data_files_set:{necessary_data_files_set}"
+        global_data_files_set & global_necessary_files_set
+        == global_necessary_files_set
+    ), f"The checkpoint files are not complete. Please check the checkpoint directory:{path}.global_data_files_set:{global_data_files_set}, necessary_data_files_set:{global_necessary_files_set}"
     missing_keys = set(state_dict.keys()) - set(tensor_key_list)
     if len(missing_keys) > 0:
         logger.warning(
@@ -117,7 +131,7 @@ def get_rank_to_files(path, state_dict, process_group, use_dist):
     for rank, local_files in enumerate(global_data_files):
         if len(local_files) > 0:
             local_files = [
-                f for f in local_files if f in necessary_data_files_set
+                f for f in local_files if f in all_necessary_files[rank]
             ]
             rank_to_files[rank] = local_files
     logger.debug(f"mapping rank_to_files:{rank_to_files}")
@@ -331,7 +345,9 @@ def get_read_items(path, state_dict, process_group, use_dist):
                 global_offset = (
                     tuple([0] * len(val.shape)) if len(val.shape) > 0 else ()
                 )
-            cur_chunk_metadata = LocalTensorMetadata(global_offset, local_shape)
+            cur_chunk_metadata = LocalTensorMetadata(
+                global_offset, local_shape, str(val.dtype).split(".")[1]
+            )
             assert (
                 tensor_key in storage_state_dict_metadata
             ), f"tensor_key:{tensor_key} not found in storage_state_dict_metadata:{storage_state_dict_metadata}."
@@ -353,6 +369,7 @@ def get_read_items(path, state_dict, process_group, use_dist):
                     ReadItem(
                         storage_local_tensor_index,
                         paddle.distributed.get_rank(),
+                        storage_local_tensor_metadata.dtype,
                         tuple(cur_offsets),
                         tuple(storage_offsets),
                         tuple(lengths),
@@ -437,6 +454,7 @@ def load_state_dict(
         rank_to_files, missing_keys = get_rank_to_files(
             path, flat_state_dict, process_group, use_dist
         )
+
         if len(missing_keys) > 0:
             logger.warning(
                 f"The following keys:{missing_keys} are not found in checkpoint path: {path}."
@@ -533,14 +551,13 @@ def load_state_dict(
             else:
                 cur_chunk_tensor = paddle.zeros(
                     item.lengths,
-                    dtype=flat_state_dict[
-                        item.local_tensor_index.tensor_key
-                    ].dtype,
+                    item.dtype,
                 )
 
             if src_rank == item.rank:
-                # assign value locally
-                paddle.assign(storage_chunk_tensor, cur_chunk_tensor)
+                if src_rank == paddle.distributed.get_rank():
+                    # assign value locally
+                    paddle.assign(storage_chunk_tensor, cur_chunk_tensor)
             else:
                 # assign value remotely
                 if src_rank == paddle.distributed.get_rank():
@@ -549,9 +566,11 @@ def load_state_dict(
                         storage_chunk_tensor, src=src_rank, group=process_group
                     )
                 else:
+                    tmp_tensor = paddle.assign(cur_chunk_tensor)
                     paddle.distributed.broadcast(
-                        cur_chunk_tensor, src=src_rank, group=process_group
+                        tmp_tensor, src=src_rank, group=process_group
                     )
+                    paddle.assign(tmp_tensor, cur_chunk_tensor)
 
         for k, v in flat_state_dict.items():
             if k in state_dict_in_cpu:
diff --git a/python/paddle/distributed/checkpoint/metadata.py b/python/paddle/distributed/checkpoint/metadata.py
@@ -24,6 +24,7 @@ class LocalTensorMetadata:
 
     global_offset: Tuple[int]
     local_shape: Tuple[int]
+    dtype: str
 
 
 @dataclass(frozen=True)
diff --git a/python/paddle/distributed/checkpoint/save_state_dict.py b/python/paddle/distributed/checkpoint/save_state_dict.py
@@ -25,16 +25,6 @@
 )
 
 
-def check_state_dict(state_dict, process_group):
-    local_keys = list(state_dict.keys())
-    global_keys = []
-    paddle.distributed.all_gather_object(global_keys, local_keys, process_group)
-    for keys in global_keys[1:]:
-        assert (
-            keys == global_keys[0]
-        ), f"keys:{keys} != first_keys: {global_keys[0]}"
-
-
 def check_file_name(file_name, process_group):
     all_unique_id = []
     unique_id = int(file_name.split(".")[0].split("_")[1])
@@ -160,8 +150,6 @@ def save_state_dict(
         logger.debug(f"file_name:{file_name}")
         if use_dist:
             check_file_name(file_name, process_group)
-            # the parameter_name and order in state_dict should be the same
-            check_state_dict(flat_state_dict, process_group)
         metadata = Metadata()
         local_state_dict = {}
         local_state_dict_metadata = {}
@@ -199,8 +187,9 @@ def save_state_dict(
                     )
                     local_tensor = val
                 local_state_dict[key] = local_tensor
+                local_tenosr_dtype = str(local_tensor.dtype).split('.')[1]
                 local_state_dict_metadata[key] = LocalTensorMetadata(
-                    global_offset, local_shape
+                    global_offset, local_shape, local_tenosr_dtype
                 )
                 local_storage_metadata[
                     LocalTensorIndex(key, tuple(global_offset))
