fix

Author: xingmingyyj

python/paddle/distributed/flex_checkpoint/dcp/load_state_dict.py

@@ -475,6 +475,7 @@ def get_read_items(metadata_list, state_dict, process_group, use_dist):
             storage_state_dict_metadata[tensor_key] += local_tensor_metadata
 
     read_items = []
+    global_shape = None
     logger.debug(f"storage_state_dict_metadata:{storage_state_dict_metadata}")
     for tensor_key, val in state_dict.items():
         tensor_name = None
@@ -493,13 +494,15 @@ def get_read_items(metadata_list, state_dict, process_group, use_dist):
                     if len(val.shape) > 0
                     else ((), ())
                 )
+                global_shape = val.shape
                 if local_shape is None or global_offset is None:
                     continue
             else:
                 local_shape = tuple(val.shape)
                 global_offset = (
                     tuple([0] * len(val.shape)) if len(val.shape) > 0 else ()
                 )
+                global_shape = local_shape
             dtype = str(val.dtype).split(".")[1]
             tensor_name = tensor_key
         elif isinstance(val, ShardedWeight):
@@ -512,13 +515,14 @@ def get_read_items(metadata_list, state_dict, process_group, use_dist):
             tensor_name = (
                 tensor_key[0] if isinstance(tensor_key, tuple) else tensor_key
             )
+            global_shape = val.global_shape
         else:
             raise ValueError(
                 f"Only support paddle.Tensor., val type:{type(val)}"
             )
 
         cur_chunk_metadata = LocalTensorMetadata(
-            global_offset, local_shape, dtype
+            global_offset, local_shape, dtype, global_shape
         )
         assert tensor_name in storage_state_dict_metadata, (
             f"tensor_key:{tensor_key} not found in storage_state_dict_metadata:{storage_state_dict_metadata}."
@@ -636,7 +640,6 @@ def _handle_aoa(
     assert len(metadata_files) == 1, "Only support one metadata file now."
     metadata = paddle.load(os.path.join(path, metadata_files[0]))
     state_dict_metadata = metadata.state_dict_metadata
-
     source_state_shard_info = {
         param_name: [
             ShardedWeightDesc(
@@ -794,9 +797,9 @@ def load_state_dict(
     if not use_dist:
         load_dict = {}
         for key, val in state_dict.items():
-            assert (
-                val.local_shape == val.global_shape
-            ), f"{key} is not replicated!"
+            assert val.local_shape == val.global_shape, (
+                f"{key} is not replicated!"
+            )
             load_dict[key] = val
         load_state_dict_impl(
             load_dict,
@@ -850,16 +853,16 @@ def load_state_dict_impl(
     mw_name_compatibility: bool = True,
 ) -> None:
     with paddle.base.dygraph.guard():
-        assert isinstance(
-            state_dict, dict
-        ), "The state_dict should be a dictionary."
+        assert isinstance(state_dict, dict), (
+            "The state_dict should be a dictionary."
+        )
         first_key = next(iter(state_dict), None)
         if isinstance(first_key, tuple):
             flat_state_dict = state_dict
             mapping = {}
         else:
             flat_state_dict, mapping = flatten_state_dict(state_dict)
-            
+
         if len(flat_state_dict) > 0:
             for val in flat_state_dict.values():
                 assert isinstance(val, (paddle.Tensor, ShardedWeight)), (
@@ -998,9 +1001,9 @@ def _load_state_dict(
         idx = 0
         assert not any(
             isinstance(k, tuple) for k in copied_target_state_dict
-        ) or all(
-            isinstance(k, tuple) for k in copied_target_state_dict
-        ), "target_state_dict contains a mix of tuple and non-tuple keys. Please ensure key types are consistent."
+        ) or all(isinstance(k, tuple) for k in copied_target_state_dict), (
+            "target_state_dict contains a mix of tuple and non-tuple keys. Please ensure key types are consistent."
+        )
 
         for item in read_items:
             if any(isinstance(k, tuple) for k in copied_target_state_dict):
@@ -1055,9 +1058,9 @@ def _load_state_dict(
                     storage_chunk_tensor = storage_local_tensor
             # The read item rank need to be assigned
             if item.rank == paddle.distributed.get_rank():
-                assert (
-                    key in copied_target_state_dict
-                ), f"item:{item}, state_dict:{copied_target_state_dict}"
+                assert key in copied_target_state_dict, (
+                    f"item:{item}, state_dict:{copied_target_state_dict}"
+                )
 
                 cur_local_tensor = (
                     copied_target_state_dict[key]._local_value()

python/paddle/distributed/flex_checkpoint/dcp/save_state_dict.py

@@ -366,6 +366,7 @@ def save_state_dict_impl(
                         if len(val.shape) > 0
                         else ((), ())
                     )
+                    global_shape = val.shape
                     if local_shape is None or global_offset is None:
                         continue
                 else:
@@ -376,6 +377,7 @@ def save_state_dict_impl(
                         else ()
                     )
                     local_tensor = val
+                    global_shape = local_shape
             elif isinstance(val, ShardedWeight):
                 local_tensor = val.local_tensor
                 local_shape = val.local_shape