【FlexCheckpoint】fix safetensors ckpt integrity (PaddlePaddle#78208)

* fix safetensors ckpt integrity

* add test for coverage

Author: xingmingyyj

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

@@ -85,6 +85,8 @@ def get_checkpoint_files(
     # if unique_id is None, all file ends with .metadata and .distcp is returned
     if unique_id is None:
         unique_id = ''
+    if use_dist is None:
+        use_dist = paddle.distributed.get_world_size() > 1
     global PATH_TO_CHECKPOINT_FILES
     if use_cache and path in PATH_TO_CHECKPOINT_FILES:
         return PATH_TO_CHECKPOINT_FILES[path]
@@ -106,26 +108,44 @@ def get_checkpoint_files(
             with open(index_file_path, "r") as f:
                 index_data = json.load(f)
             if "weight_map" in index_data:
-                from safetensors.numpy import safe_open
-
                 mapping_key_to_safetensors_file = index_data["weight_map"]
-                global_safetensors_files = []
+                # All files referenced in the index
+                expected_files_in_index = set(
+                    mapping_key_to_safetensors_file.values()
+                )
+
+                # Gather safetensors files visible on each rank, then take union
+                global_safetensors_files_list = []
                 if use_dist:
                     paddle.distributed.all_gather_object(
-                        global_safetensors_files,
+                        global_safetensors_files_list,
                         safetensors_files,
                         process_group,
                     )
-                    global_safetensors_files = list(
-                        {
-                            file
-                            for files in global_safetensors_files
-                            for file in files
-                        }
-                    )
+                    global_safetensors_files = {
+                        file
+                        for files in global_safetensors_files_list
+                        for file in files
+                    }
                 else:
-                    global_safetensors_files = safetensors_files
-                for file in global_safetensors_files:
+                    global_safetensors_files = set(safetensors_files)
+
+                # Check that every file referenced in the index is visible on at least one rank
+                missing_files = (
+                    expected_files_in_index - global_safetensors_files
+                )
+                assert len(missing_files) == 0, (
+                    f"The following safetensors files are referenced in '{index_file_name}' "
+                    f"but not found on any rank: {sorted(missing_files)}"
+                )
+
+                # Check local files: every key in a locally accessible file must be
+                # consistent with the index mapping (only open files this rank can see).
+                # Meanwhile, collect actual keys present locally for the reverse check below.
+                from safetensors.numpy import safe_open
+
+                local_actual_keys = set()
+                for file in safetensors_files:
                     if file.endswith(".safetensors"):
                         file_path = os.path.join(path, file)
                         with safe_open(file_path, framework="np") as f:
@@ -139,6 +159,30 @@ def get_checkpoint_files(
                                 assert expected_file == file, (
                                     f"Key '{key}' is mapped to file '{expected_file}' in index, but found in file '{file}'"
                                 )
+                                local_actual_keys.add(key)
+
+                # Reverse check: every key declared in index must actually exist in
+                # some safetensors file across the cluster.
+                expected_keys_in_index = set(
+                    mapping_key_to_safetensors_file.keys()
+                )
+                global_actual_keys_list = []
+                if use_dist:
+                    paddle.distributed.all_gather_object(
+                        global_actual_keys_list,
+                        local_actual_keys,
+                        process_group,
+                    )
+                    global_actual_keys = set().union(*global_actual_keys_list)
+                else:
+                    global_actual_keys = local_actual_keys
+                missing_keys_in_files = (
+                    expected_keys_in_index - global_actual_keys
+                )
+                assert len(missing_keys_in_files) == 0, (
+                    f"The following keys are declared in '{index_file_name}' weight_map "
+                    f"but not found in any safetensors file: {sorted(missing_keys_in_files)}"
+                )
 
         if len(metadata_files) == 0:
             logger.info(

test/auto_parallel/hybrid_strategy/save_safetensors_load_fc.py

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import json
 import os
 
 import paddle
@@ -153,5 +154,70 @@ def test_save_safetensors_load_fc():
     dist.barrier()
 
 
+def create_index_json(ckpt_path):
+    """Create model.safetensors.index.json that maps keys to their safetensors files."""
+    index_data = {
+        "weight_map": {
+            "tensor1": "tensor1.safetensors",
+            "tensor2": "tensor2.safetensors",
+        }
+    }
+    index_file_path = os.path.join(ckpt_path, "model.safetensors.index.json")
+    if dist.get_rank() == 0:
+        with open(index_file_path, "w") as f:
+            json.dump(index_data, f)
+    dist.barrier()
+
+
+def test_save_safetensors_load_fc_with_index():
+    """Test saving safetensors and loading with flex checkpoint when model.safetensors.index.json exists."""
+    ckpt_path = os.getenv("ckpt_path")
+    dist.init_parallel_env()
+
+    save_safetensors_to_ranks(ckpt_path)
+
+    # Create index json to exercise the integrity check branch
+    create_index_json(ckpt_path)
+
+    sharded_state_dict = create_sharded_state_dict_for_loading()
+
+    from paddle.distributed.flex_checkpoint.dcp.load_state_dict import (
+        load_state_dict,
+    )
+
+    load_state_dict(sharded_state_dict, ckpt_path, safetensors=True)
+
+    loaded_tensor1 = sharded_state_dict["tensor1"].local_tensor
+    loaded_tensor2 = sharded_state_dict["tensor2"].local_tensor
+
+    if dist.get_rank() == 0:
+        expected_tensor1 = paddle.to_tensor([[0], [2]], dtype='float32')
+        expected_tensor2 = paddle.to_tensor([[4], [6]], dtype='float32')
+
+        assert paddle.allclose(loaded_tensor1, expected_tensor1), (
+            f"Rank 0 tensor1 mismatch: got {loaded_tensor1}, expected {expected_tensor1}"
+        )
+        assert paddle.allclose(loaded_tensor2, expected_tensor2), (
+            f"Rank 0 tensor2 mismatch: got {loaded_tensor2}, expected {expected_tensor2}"
+        )
+
+    elif dist.get_rank() == 1:
+        expected_tensor1 = paddle.to_tensor([[1], [3]], dtype='float32')
+        expected_tensor2 = paddle.to_tensor([[5], [7]], dtype='float32')
+
+        assert paddle.allclose(loaded_tensor1, expected_tensor1), (
+            f"Rank 1 tensor1 mismatch: got {loaded_tensor1}, expected {expected_tensor1}"
+        )
+        assert paddle.allclose(loaded_tensor2, expected_tensor2), (
+            f"Rank 1 tensor2 mismatch: got {loaded_tensor2}, expected {expected_tensor2}"
+        )
+
+    dist.barrier()
+
+
 if __name__ == "__main__":
-    test_save_safetensors_load_fc()
+    test_func = os.getenv("test_func", "test_save_safetensors_load_fc")
+    if test_func == "test_save_safetensors_load_fc_with_index":
+        test_save_safetensors_load_fc_with_index()
+    else:
+        test_save_safetensors_load_fc()

test/auto_parallel/hybrid_strategy/test_save_load_state_dict.py

@@ -132,6 +132,20 @@ def test_save_safetensors_load_fc(self):
         )
         ckpt_path.cleanup()
 
+    def test_save_safetensors_load_fc_with_index(self):
+        """Test saving safetensors files and loading with flex checkpoint when model.safetensors.index.json exists."""
+        ckpt_path = tempfile.TemporaryDirectory()
+        super().setUp(num_of_devices=2, timeout=120, nnode=1)
+        self.run_test_case(
+            "save_safetensors_load_fc.py",
+            user_defined_envs={
+                "device_num": "2",
+                "ckpt_path": ckpt_path.name,
+                "test_func": "test_save_safetensors_load_fc_with_index",
+            },
+        )
+        ckpt_path.cleanup()
+
     def test_save_load_state_dict_with_aoa_config_reverse(self):
         """Test saving state dict and loading with flex checkpoint."""
         ckpt_path = tempfile.TemporaryDirectory()