chore: import upstream snapshot with attribution

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

@@ -0,0 +1,13 @@
+# Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.

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

@@ -0,0 +1,849 @@
+# Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+import abc
+import math
+from collections import defaultdict
+from copy import deepcopy
+from dataclasses import dataclass
+from enum import Enum
+from typing import (
+    TYPE_CHECKING,
+)
+
+import paddle
+
+from ..aoa.aoa_engine import SUPPORTED_DTYPES, AOAEngine
+from .resharder import (
+    ReadItem,
+)
+from .sharded_weight import (
+    ShardedWeight,
+    ShardedWeightDesc,
+)
+from .utils import (
+    assign_sharded_slice,
+    build_shard_desc,
+    merge_shard_info_list,
+    recover_shard_tensor_from_shards,
+)
+
+if TYPE_CHECKING:
+    from collections.abc import Generator, Iterable
+
+    from paddle.distributed.collective import Group
+
+    from .sharded_weight import ShardedStateDict
+
+
+INTERNAL_PADDING_TENSOR_NAME = "__internal_padding_tensor_name__"
+
+
+@dataclass(frozen=True)
+class ExtendReadItem(ReadItem):
+    target_tensor_names: tuple[str] | None = None
+    global_shape: tuple[int] | None = None
+
+
+class BaseAssembler(abc.ABC):
+    """
+    Abstract base class for assembling full parameters from sharded states.
+
+    This class encapsulates the common logic for:
+        1.  Analyzing source and destination tensor mappings (AOA).
+        2.  Creating a plan to read/communicate necessary tensor shards.
+        3.  Assembling final tensors once all their source shards are available.
+        4.  Managing memory by cleaning up consumed shards.
+
+    Subclasses must implement the `run` method, which defines the specific
+    distributed communication strategy to fetch the tensor shards.
+    """
+
+    def __init__(
+        self,
+        sharded_state_dict: ShardedStateDict,
+        aoa_config: dict[str, list[str]] | None = None,
+        num_splits: int = 1,
+        idx: int = 0,
+    ):
+        self.sharded_state_dict = sharded_state_dict
+        self.aoa_config = aoa_config or {}
+        self.num_splits = num_splits
+        self.idx = idx
+
+        self.cur_rank: int = paddle.distributed.get_rank()
+        self.world_size: int = paddle.distributed.get_world_size()
+        self.use_dist: bool = self.world_size > 1
+
+        self.filtered_sharded_state_dict = {}
+        self.aoa_engine = None
+        self.destination_sharded_weight_desc: dict[str, ShardedWeightDesc] = {}
+        self.destination_sharded_mappings = {}
+
+        self.source_to_target_names: dict[str, set[str]] = defaultdict(set)
+        self.source_consumers: dict[str, set[str]] = {}
+        self.ref_map: dict[str, set] = {}
+        self.read_items: list[ExtendReadItem] = []
+
+        self.sharded_desc_to_tensor: dict[ShardedWeightDesc, paddle.Tensor] = {}
+
+    def _prepare_metainfo(self, source_state_shard_info):
+        """Builds destination descriptions and mappings using AOAEngine."""
+        self.aoa_engine = AOAEngine(
+            aoa_config=self.aoa_config,
+            source_state_shard_info=source_state_shard_info,
+            destination_state_shard_info=None,
+        )
+
+        output_vars = self.split_output_vars()
+
+        for k, v in output_vars.items():
+            dtype = self.infer_real_dtype(v)
+            self.destination_sharded_weight_desc[k] = ShardedWeightDesc(
+                key=k,
+                local_shape=v.shape,
+                global_shape=v.shape,
+                global_offset=(0,) * len(v.shape),
+                dtype=dtype,
+            )
+
+        for k, desc in self.destination_sharded_weight_desc.items():
+            self.destination_sharded_mappings[k] = (
+                self.aoa_engine.find_shard_sources(desc)
+            )
+
+        for tgt_name, mapping in self.destination_sharded_mappings.items():
+            for m in mapping:
+                self.source_to_target_names[m.source_slice.key].add(tgt_name)
+
+        self.filtered_sharded_state_dict = {
+            k: v
+            for k, v in self.sharded_state_dict.items()
+            if k in self.source_to_target_names
+        }
+
+        self.source_consumers = deepcopy(self.source_to_target_names)
+
+    def split_output_vars(self):
+        data_dict = self.aoa_engine.output_vars
+        if self.num_splits < 1:
+            raise ValueError('num_splits must be >= 1')
+        if self.idx < 0 or self.idx >= self.num_splits:
+            raise IndexError(f'idx must be in [0,{self.num_splits - 1}]')
+
+        sorted_keys = sorted(data_dict.keys())
+        total = len(sorted_keys)
+        base = total // self.num_splits
+        extra = total % self.num_splits
+
+        if self.idx < extra:
+            start = self.idx * (base + 1)
+            end = start + (base + 1)
+        else:
+            start = extra * (base + 1) + (self.idx - extra) * base
+            end = start + base
+
+        selected_keys = sorted_keys[start:end]
+        return {k: data_dict[k] for k in selected_keys}
+
+    def _assemble_and_yield_ready_tensors(
+        self, ready_tensor_names: list[str]
+    ) -> Iterable[tuple[str, paddle.Tensor]]:
+        """
+        Assembles, yields, and cleans up tensors whose dependencies are all met.
+        This logic is shared across different communication strategies.
+        """
+        if not ready_tensor_names:
+            return
+
+        for name in ready_tensor_names:
+            target_desc = self.destination_sharded_weight_desc[name]
+            local_tensor = paddle.empty(
+                target_desc.local_shape, dtype=target_desc.dtype
+            )
+            cur_sharded_tensor = ShardedWeight(
+                key=target_desc.key,
+                local_tensor=local_tensor,
+                local_shape=target_desc.local_shape,
+                global_shape=target_desc.global_shape,
+                global_offset=target_desc.global_offset,
+            )
+
+            for mapping in self.destination_sharded_mappings[name]:
+                src_desc = mapping.source_slice
+                dst_desc = mapping.target_slice
+
+                src_shard_template = ShardedWeight(
+                    key=src_desc.key,
+                    local_tensor=paddle.zeros(
+                        src_desc.local_shape, dtype=src_desc.dtype
+                    ),
+                    local_shape=src_desc.local_shape,
+                    global_shape=src_desc.global_shape,
+                    global_offset=src_desc.global_offset,
+                )
+
+                received_shards = []
+                for desc, tensor in self.sharded_desc_to_tensor.items():
+                    if desc.key == src_desc.key:
+                        received_shards.append(
+                            ShardedWeight(
+                                key=desc.key,
+                                local_tensor=tensor,
+                                local_shape=desc.local_shape,
+                                global_shape=desc.global_shape,
+                                global_offset=desc.global_offset,
+                            )
+                        )
+
+                recover_shard_tensor_from_shards(
+                    received_shards, src_shard_template
+                )
+
+                assign_sharded_slice(
+                    src_desc=src_desc,
+                    src_shard=src_shard_template,
+                    dst_desc=dst_desc,
+                    dst_shard=cur_sharded_tensor,
+                    postprocess_list=mapping.postprocess_list,
+                )
+                src_shard_template.local_tensor._clear()
+
+            yield name, cur_sharded_tensor.local_tensor
+
+        need_clear_source_names = self._update_consumer_counts(
+            ready_tensor_names
+        )
+
+        self._cleanup_consumed_shards(need_clear_source_names)
+
+    def _update_consumer_counts(
+        self, ready_tensor_names: list[str]
+    ) -> list[str]:
+        """Decrement consumer counts and return source names that can be cleared."""
+        need_clear_source_names = []
+        del_keys = []
+        for source_name, target_names in self.source_consumers.items():
+            target_names.difference_update(ready_tensor_names)
+            if not target_names:
+                del_keys.append(source_name)
+                need_clear_source_names.append(source_name)
+
+        for k in del_keys:
+            del self.source_consumers[k]
+
+        return need_clear_source_names
+
+    def dedup_read_items(self, global_read_items):
+        group = defaultdict(list)
+        for item in global_read_items:
+            key = (item.tensor_name, item.src_global_offset, item.slice_shape)
+            group[key].append(item)
+        result = []
+        for key, items in group.items():
+            min_item = min(items, key=lambda x: x.src_rank)
+            result.append(min_item)
+        return result
+
+    def _cleanup_consumed_shards(self, source_names_to_clear: list[str]):
+        """Delete cached tensors corresponding to the given source names."""
+        if not source_names_to_clear:
+            return
+
+        to_delete_descs = []
+        for desc, tensor in self.sharded_desc_to_tensor.items():
+            if desc.key in source_names_to_clear:
+                tensor._clear()
+                to_delete_descs.append(desc)
+
+        for desc in to_delete_descs:
+            del self.sharded_desc_to_tensor[desc]
+
+    @abc.abstractmethod
+    def prepare(self):
+        """Subclasses must implement this to build their specific read plan."""
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    def run(self) -> Generator[tuple[str, paddle.Tensor], None, None]:
+        """
+        The main entry point. Subclasses must implement their communication
+        loop and yield final tensors.
+        """
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    def all_gather_fn(self, info, **kwargs):
+        raise NotImplementedError
+
+    def infer_real_dtype(self, desc) -> str:
+        found_dtypes = []
+        for slice_ref in desc.slices:
+            key, sl_src, sl_dst, pp_list = slice_ref
+            if pp_list is None or len(pp_list) == 0:
+                continue
+            last_supported = None
+            for item in reversed(pp_list):
+                if item in SUPPORTED_DTYPES:
+                    last_supported = item
+                    break
+            if last_supported:
+                found_dtypes.append(last_supported)
+        if not found_dtypes:
+            return desc.dtype
+
+        dtype_set = set(found_dtypes)
+        if len(dtype_set) > 1:
+            raise ValueError(
+                f"Found multiple different dtypes from slices: {dtype_set}"
+            )
+        return found_dtypes[0]
+
+    def build_global_state_shard_info(self, **all_gather_args):
+        state_shard_info = defaultdict(list)
+        for key, val in self.sharded_state_dict.items():
+            desc = build_shard_desc(val)
+            state_shard_info[key].append(desc)
+
+        use_dist = True if paddle.distributed.get_world_size() > 1 else False
+
+        if use_dist:
+            gathered_info = self.all_gather_fn(
+                dict(state_shard_info), **all_gather_args
+            )
+        else:
+            gathered_info = [dict(state_shard_info)]
+
+        return merge_shard_info_list(gathered_info)
+
+    def get_read_items(
+        self,
+        all_gather_args=None,
+    ):
+        current_rank = paddle.distributed.get_rank()
+        rank_vfile = f"{current_rank}.vdistcp"
+
+        local_read_plan = []
+        for tensor_name, shard_info in self.filtered_sharded_state_dict.items():
+            common_attrs = {
+                "tensor_name": tensor_name,
+                "src_rank": current_rank,
+                "src_global_offset": tuple(shard_info.global_offset),
+                "dst_global_offset": tuple(shard_info.global_offset),
+                "src_local_offset": (0,) * len(shard_info.local_shape),
+                "dst_local_offset": (0,) * len(shard_info.local_shape),
+                "slice_shape": tuple(shard_info.local_shape),
+                "global_shape": tuple(shard_info.global_shape),
+                "target_tensor_names": tuple(
+                    self.source_to_target_names[tensor_name]
+                ),
+                "file_name": rank_vfile,
+                "dtype": str(shard_info.local_tensor.dtype).split(".")[1],
+                "dst_rank": None,
+                "comm_group": None,
+            }
+            local_read_plan.append(ExtendReadItem(**common_attrs))
+        gathered_plans_per_rank = self.all_gather_fn(
+            local_read_plan, **(all_gather_args or {})
+        )
+
+        global_read_plan = [
+            item for plan in gathered_plans_per_rank for item in plan
+        ]
+
+        return self.dedup_read_items(global_read_plan)
+
+    def group_read_items_by_tensor_name(self, global_read_items):
+        groups = defaultdict(list)
+        for item in global_read_items:
+            groups[item.tensor_name].append(item)
+        return groups
+
+    def sort_groups_for_early_release(self, groups, source_to_target_names):
+        def count_fn(name):
+            return len(source_to_target_names.get(name, []))
+
+        sorted_items = sorted(groups.items(), key=lambda x: -count_fn(x[0]))
+        return dict(sorted_items)
+
+    def build_reference_map(self, groups: dict[str, set[ExtendReadItem]]):
+        ref_map = defaultdict(set)
+        for _, items in groups.items():
+            for item in items:
+                for tgt in item.target_tensor_names:
+                    ref_map[tgt].add(item)
+        return ref_map
+
+    def _build_read_plan(self, all_gather_args):
+        """Creates an optimized, sorted list of read operations."""
+        read_items = self.get_read_items(
+            all_gather_args=all_gather_args,
+        )
+        grouped = self.group_read_items_by_tensor_name(read_items)
+        grouped = self.sort_groups_for_early_release(
+            grouped, self.source_to_target_names
+        )
+        self.ref_map = self.build_reference_map(grouped)
+
+        self.read_items = [
+            item for _, items in grouped.items() for item in items
+        ]
+
+    def __iter__(self):
+        return self.run()
+
+
+class SingleCommGroupFullParamAssembler(BaseAssembler):
+    """
+    Implements the assembly logic from the original full_param function.
+    This version handles both single-card and distributed scenarios.
+    In the distributed case, it uses a broadcast-based communication strategy.
+    """
+
+    def __init__(
+        self,
+        sharded_state_dict: ShardedStateDict,
+        aoa_config: dict[str, list[str]] | None = None,
+        process_group: Group | None = None,
+        num_splits: int = 1,
+        idx: int = 0,
+    ):
+        super().__init__(sharded_state_dict, aoa_config, num_splits, idx)
+        self.process_group = process_group
+
+    def all_gather_fn(self, info, **kwargs):
+        process_group = kwargs.get('process_group', self.process_group)
+        gathered_info = []
+        paddle.distributed.all_gather_object(gathered_info, info, process_group)
+        return gathered_info
+
+    def is_identity_mapping(self, shard_mappings):
+        if len(shard_mappings) != 1:
+            return False
+        mapping = shard_mappings[0]
+        src = mapping.source_slice
+        dst = mapping.target_slice
+        return (
+            src.key == dst.key
+            and src.local_shape == dst.local_shape
+            and src.global_shape == dst.global_shape
+            and src.global_offset == dst.global_offset
+            and src.dtype == dst.dtype
+            and mapping.postprocess_list is None
+        )
+
+    def prepare(self):
+        """Prepare metadata and build the read plan."""
+        source_state_shard_info = self.build_global_state_shard_info(
+            process_group=self.process_group
+        )
+
+        self._prepare_metainfo(source_state_shard_info)
+
+        if self.use_dist:
+            self._build_read_plan(
+                all_gather_args={"process_group": self.process_group}
+            )
+
+    def run(self) -> Generator[tuple[str, paddle.Tensor], None, None]:
+        """Main execution generator."""
+        self.prepare()
+        if not self.use_dist:
+            yield from self._run_single_card()
+        else:
+            yield from self._run_distributed()
+
+    def _run_single_card(
+        self,
+    ) -> Generator[tuple[str, paddle.Tensor], None, None]:
+        """Simple assembly path for a single GPU."""
+        for k, v in self.filtered_sharded_state_dict.items():
+            assert v.local_shape == v.global_shape, (
+                "Single card params must not be sharded.But now the key is {k}, the local_shape is {v.local_shape}, the global_shape is {v.global_shape}."
+            )
+
+        for k, shard_mappings in self.destination_sharded_mappings.items():
+            if self.is_identity_mapping(shard_mappings):
+                src_key = shard_mappings[0].source_slice.key
+                yield (
+                    k,
+                    self.filtered_sharded_state_dict[
+                        src_key
+                    ].local_tensor.clone(),
+                )
+            else:
+                desc = self.destination_sharded_weight_desc[k]
+                cur_sharded_tensor = ShardedWeight(
+                    key=desc.key,
+                    local_tensor=paddle.empty(
+                        desc.local_shape, dtype=desc.dtype
+                    ),
+                    local_shape=desc.local_shape,
+                    global_shape=desc.global_shape,
+                    global_offset=desc.global_offset,
+                )
+                for mapping in shard_mappings:
+                    source_tensor = self.filtered_sharded_state_dict[
+                        mapping.source_slice.key
+                    ]
+                    assign_sharded_slice(
+                        src_desc=mapping.source_slice,
+                        src_shard=source_tensor,
+                        dst_desc=mapping.target_slice,
+                        dst_shard=cur_sharded_tensor,
+                        postprocess_list=mapping.postprocess_list,
+                    )
+                yield k, cur_sharded_tensor.local_tensor
+
+    def _run_distributed(
+        self,
+    ) -> Generator[tuple[str, paddle.Tensor], None, None]:
+        """Distributed assembly using broadcast and packed buffers."""
+        for item in self.read_items:
+            cur_src_rank = item.src_rank
+
+            if self.cur_rank == cur_src_rank:
+                local_tensor = self.filtered_sharded_state_dict[
+                    item.tensor_name
+                ].local_tensor.clone()
+            else:
+                local_tensor = paddle.empty(item.slice_shape, dtype=item.dtype)
+
+            on_cpu = local_tensor.place.is_cpu_place()
+            if on_cpu:
+                local_tensor = local_tensor.cuda()
+            paddle.distributed.broadcast(
+                local_tensor, src=cur_src_rank, group=self.process_group
+            )
+            if on_cpu:
+                local_tensor = local_tensor.cpu()
+
+            shard_desc = ShardedWeightDesc(
+                key=item.tensor_name,
+                local_shape=item.slice_shape,
+                global_shape=item.global_shape,
+                global_offset=item.src_global_offset,
+                dtype=item.dtype,
+            )
+            self.sharded_desc_to_tensor[shard_desc] = local_tensor
+
+            ready_tensor_names = []
+            for name in item.target_tensor_names:
+                self.ref_map[name].remove(item)
+                if len(self.ref_map[name]) == 0:
+                    ready_tensor_names.append(name)
+                    del self.ref_map[name]
+
+            yield from self._assemble_and_yield_ready_tensors(
+                ready_tensor_names
+            )
+
+
+class OperationType(Enum):
+    GLOBAL_BROADCAST = 1
+    BROADCAST_ALLGATHER = 2
+
+
+class HVCommGroupFullParamAssembler(BaseAssembler):
+    """
+    Implements the assembly logic using a 2D-mesh communication strategy.
+
+    This strategy involves a broadcast along the vertical axis of the process
+    mesh, followed by an all-gather along the horizontal axis.
+    """
+
+    def __init__(
+        self,
+        sharded_state_dict: ShardedStateDict,
+        horizontal_group: Group,
+        vertical_group: Group,
+        aoa_config: dict[str, list[str]] | None = None,
+        num_splits: int = 1,
+        idx: int = 0,
+        memory_growth_threshold: int = 8 * (2**30),  # 8GB
+    ):
+        super().__init__(sharded_state_dict, aoa_config, num_splits, idx)
+        self.h_group = horizontal_group
+        self.v_group = vertical_group
+        self.using_1d_comm_group = (
+            self.v_group is None or self.v_group.nranks == 1
+        )
+
+        self.topology: list[list[int]] = []
+        self.vertical_ranks: list[set[int]] = []
+        self.horizontal_index: dict[int, int] = {}
+        self.vertical_index: dict[int, int] = {}
+        self.cur_horizontal_index: int = -1
+        self.memory_growth_threshold = memory_growth_threshold
+
+    def all_gather_fn(self, info, **kwargs):
+        h_group = kwargs.get('h_group', self.h_group)
+        v_group = kwargs.get('v_group', self.v_group)
+
+        h_obj_list = []
+        paddle.distributed.all_gather_object(h_obj_list, info, h_group)
+
+        v_obj_list = []
+        if not self.using_1d_comm_group:
+            paddle.distributed.all_gather_object(
+                v_obj_list, h_obj_list, v_group
+            )
+        else:
+            v_obj_list = [h_obj_list]
+
+        gathered_info = [x for sublist in v_obj_list for x in sublist]
+        return gathered_info
+
+    def prepare(self):
+        """Build topology, prepare metadata, and build the read plan."""
+        assert self.use_dist, (
+            "FullParamAssembler only supports distributed training."
+        )
+        self._build_topology()
+
+        source_state_shard_info = self.build_global_state_shard_info(
+            h_group=self.h_group, v_group=self.v_group
+        )
+        self._prepare_metainfo(source_state_shard_info)
+        self._build_read_plan(
+            all_gather_args={'h_group': self.h_group, 'v_group': self.v_group}
+        )
+
+    def _build_topology(self):
+        h_ranks = []
+        paddle.distributed.all_gather_object(
+            h_ranks, self.cur_rank, self.h_group
+        )
+        if not self.using_1d_comm_group:
+            paddle.distributed.all_gather_object(
+                self.topology, h_ranks, self.v_group
+            )
+        else:
+            self.topology = [h_ranks]
+        self.vertical_ranks = [set(col) for col in zip(*self.topology)]
+        self.horizontal_index = {
+            rank: i
+            for i, ranks in enumerate(self.vertical_ranks)
+            for rank in ranks
+        }
+        self.vertical_index = {
+            rank: i for i, row in enumerate(self.topology) for rank in row
+        }
+        self.cur_horizontal_index = self.horizontal_index[self.cur_rank]
+
+    def run(self) -> Generator[tuple[str, paddle.Tensor], None, None]:
+        """Main execution generator using 2D-mesh communication."""
+        self.prepare()
+
+        while len(self.read_items) > 0:
+            ready_tensor_names = self._process_one_batch()
+
+            yield from self._assemble_and_yield_ready_tensors(
+                ready_tensor_names
+            )
+
+    def get_batch_read_items(self):
+        read_items = self.read_items
+        vertical_ranks = self.vertical_ranks
+        horizontal_index = self.horizontal_index
+
+        bathch_read_items = [None] * len(vertical_ranks)
+        read_item_index = [None] * len(vertical_ranks)
+        cnt = 0
+        cur_shape = None
+        cur_dtype = None
+        for i, item in enumerate(read_items):
+            src_rank = item.src_rank
+            h_index = horizontal_index[src_rank]
+            if bathch_read_items[h_index] is None and cnt == 0:
+                bathch_read_items[h_index] = item
+                read_item_index[h_index] = i
+                cnt += 1
+                cur_dtype = item.dtype
+                cur_shape = item.slice_shape
+                element_size = paddle.core.size_of_dtype(
+                    getattr(paddle, cur_dtype)
+                )
+                memory_growth = (
+                    element_size * math.prod(cur_shape) * len(vertical_ranks)
+                )
+                if memory_growth > self.memory_growth_threshold:
+                    return (
+                        bathch_read_items,
+                        read_item_index,
+                        OperationType.GLOBAL_BROADCAST,
+                    )
+                if cnt == len(vertical_ranks):
+                    return (
+                        bathch_read_items,
+                        read_item_index,
+                        OperationType.GLOBAL_BROADCAST,
+                    )
+
+            if bathch_read_items[h_index] is None and cnt != 0:
+                if item.slice_shape == cur_shape and item.dtype == cur_dtype:
+                    bathch_read_items[h_index] = item
+                    read_item_index[h_index] = i
+                    cnt += 1
+                    if cnt == len(vertical_ranks):
+                        return (
+                            bathch_read_items,
+                            read_item_index,
+                            OperationType.BROADCAST_ALLGATHER,
+                        )
+
+        assert cur_shape is not None
+        assert cur_dtype is not None
+
+        for i, item in enumerate(bathch_read_items):
+            if item is None:
+                src_rank = min(vertical_ranks[i])
+                common_attrs = {
+                    "tensor_name": INTERNAL_PADDING_TENSOR_NAME,
+                    "src_rank": src_rank,
+                    "src_global_offset": (0,) * len(cur_shape),
+                    "dst_global_offset": (0,) * len(cur_shape),
+                    "src_local_offset": (0,) * len(cur_shape),
+                    "dst_local_offset": (0,) * len(cur_shape),
+                    "slice_shape": cur_shape,
+                    "global_shape": cur_shape,
+                    "target_tensor_names": None,
+                    "file_name": "padding_vfile",
+                    "dtype": cur_dtype,
+                    "comm_group": None,
+                }
+
+                padding_read_item = ExtendReadItem(
+                    dst_rank=None, **common_attrs
+                )
+                bathch_read_items[i] = padding_read_item
+
+        return (
+            bathch_read_items,
+            read_item_index,
+            OperationType.BROADCAST_ALLGATHER,
+        )
+
+    def _process_one_batch(self) -> list[str]:
+        """Performs V-Broadcast + H-AllGather for one batch of items."""
+
+        batch_items, batch_indices, op_type = self.get_batch_read_items()
+
+        if op_type == OperationType.BROADCAST_ALLGATHER:
+            read_item = batch_items[self.cur_horizontal_index]
+        else:
+            values = [x for x in batch_items if x is not None]
+            if len(values) == 1:
+                read_item = values[0]
+            else:
+                raise ValueError(
+                    "When the comm op is GLOBAL_BROADCAST, read_items should be of length 1!"
+                )
+            batch_items = [read_item]
+
+        if self.cur_rank == read_item.src_rank:
+            buffer = (
+                paddle.empty(read_item.slice_shape, read_item.dtype)
+                if read_item.tensor_name == INTERNAL_PADDING_TENSOR_NAME
+                else self.filtered_sharded_state_dict[
+                    read_item.tensor_name
+                ].local_tensor.clone()
+            )
+        else:
+            buffer = paddle.empty(read_item.slice_shape, dtype=read_item.dtype)
+
+        if op_type == OperationType.BROADCAST_ALLGATHER:
+            if not self.using_1d_comm_group:
+                paddle.distributed.broadcast(
+                    buffer, src=read_item.src_rank, group=self.v_group
+                )
+            tensor_list = []
+            paddle.distributed.all_gather(
+                tensor_list, buffer, group=self.h_group
+            )
+        else:
+            src_rank = read_item.src_rank
+            v_ranks = sorted(
+                self.vertical_ranks[self.horizontal_index[src_rank]]
+            )
+            if self.cur_rank in v_ranks:
+                if not self.using_1d_comm_group:
+                    paddle.distributed.broadcast(
+                        buffer, src=src_rank, group=self.v_group
+                    )
+            src_rank = v_ranks[self.vertical_index[self.cur_rank]]
+            paddle.distributed.broadcast(
+                buffer, src=src_rank, group=self.h_group
+            )
+            tensor_list = [buffer]
+
+        for idx, item in enumerate(batch_items):
+            if item.tensor_name != INTERNAL_PADDING_TENSOR_NAME:
+                shard_desc = ShardedWeightDesc(
+                    key=item.tensor_name,
+                    local_shape=item.slice_shape,
+                    global_shape=item.global_shape,
+                    global_offset=item.src_global_offset,
+                    dtype=item.dtype,
+                )
+                self.sharded_desc_to_tensor[shard_desc] = tensor_list[idx]
+
+        ready_tensor_names = []
+        for item in batch_items:
+            if item.target_tensor_names:
+                for name in item.target_tensor_names:
+                    self.ref_map[name].remove(item)
+                    if not self.ref_map[name]:
+                        ready_tensor_names.append(name)
+                        del self.ref_map[name]
+
+        for index in sorted(
+            [i for i in batch_indices if i is not None], reverse=True
+        ):
+            del self.read_items[index]
+
+        return ready_tensor_names
+
+
+@paddle.no_grad()
+def full_param(
+    sharded_state_dict: ShardedStateDict,
+    aoa_config: dict[str, list[str]] | None = None,
+    **kwargs,
+):
+    h_group = kwargs.pop("h_group", None)
+    v_group = kwargs.pop("v_group", None)
+    process_group = kwargs.pop("process_group", None)
+    num_splits = kwargs.pop("num_splits", 1)
+    memory_growth_threshold = kwargs.pop("memory_growth_threshold", 8 * (2**30))
+    idx = kwargs.pop("shard_idx", 0)
+    assert (h_group and v_group) or not (h_group or v_group), (
+        "Both horizontal and vertical groups must be provided when using FullParamAssembler."
+    )
+    if h_group and v_group:
+        return HVCommGroupFullParamAssembler(
+            sharded_state_dict,
+            h_group,
+            v_group,
+            aoa_config,
+            num_splits,
+            idx,
+            memory_growth_threshold,
+        )
+    else:
+        return SingleCommGroupFullParamAssembler(
+            sharded_state_dict, aoa_config, process_group
+        )

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

@@ -0,0 +1,751 @@
+# Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+import re
+from collections import defaultdict
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING
+
+import paddle
+from paddle.distributed.fleet.utils.log_util import logger
+
+if TYPE_CHECKING:
+    from paddle.distributed.communication.group import Group
+
+    from ..aoa.aoa_engine import AOAEngine
+    from .metadata import Metadata
+
+# ---------------------------------------------------------------------------
+# Configuration
+# ---------------------------------------------------------------------------
+
+_MAX_TOTAL_LINES = 500
+_MAX_KEYS_SHOWN = 50
+_MAX_SHAPE_MISMATCHES = 20
+_MAX_PATTERNS_SHOWN = 30
+_SRC_FOLD_THRESHOLD = 5
+_MAX_SLICE_DETAIL_KEYS = 5
+
+
+def _get_rank() -> int:
+    return paddle.distributed.get_rank()
+
+
+# ---------------------------------------------------------------------------
+# Color support (disabled by default)
+# ---------------------------------------------------------------------------
+
+
+class _C:
+    """No-op color helpers. Colors are disabled by default."""
+
+    @staticmethod
+    def green(t):
+        return t
+
+    @staticmethod
+    def yellow(t):
+        return t
+
+    @staticmethod
+    def red(t):
+        return t
+
+    @staticmethod
+    def cyan(t):
+        return t
+
+
+# ---------------------------------------------------------------------------
+# Data structures
+# ---------------------------------------------------------------------------
+
+
+@dataclass
+class ShapeMismatchInfo:
+    key: str
+    src_global_shape: tuple[int, ...]
+    dst_global_shape: tuple[int, ...]
+    src_dtype: str | None = None
+    dst_dtype: str | None = None
+
+
+@dataclass
+class KeyValidationResult:
+    missing_keys: set[str] = field(default_factory=set)
+    unexpected_keys: set[str] = field(default_factory=set)
+    shape_mismatches: list[ShapeMismatchInfo] = field(default_factory=list)
+    randomly_initialized_keys: set[str] = field(default_factory=set)
+
+
+@dataclass
+class AOASliceMapping:
+    src_key: str
+    src_slice: tuple[slice, ...]
+    dst_slice: tuple[slice, ...]
+    postprocess: list[str] | None = None
+
+
+@dataclass
+class AOAMappingEntry:
+    dst_key: str
+    dst_global_shape: tuple[int, ...]
+    slice_mappings: list[AOASliceMapping] = field(default_factory=list)
+    is_identity: bool = False
+
+
+# ---------------------------------------------------------------------------
+# Public API: Standard (non-AOA) validation
+# ---------------------------------------------------------------------------
+
+
+def validate_and_report_keys_standard(
+    metadata_list: list[Metadata],
+    state_dict_param_names: set[str],
+    process_group: Group | None,
+    use_dist: bool,
+    checkpoint_path: str,
+    state_dict: dict,
+) -> KeyValidationResult:
+    """Validate keys for the standard (non-AOA) loading path.
+
+    Gathers global dst keys across all ranks, compares with global src keys,
+    checks shape mismatches. Prints report on rank 0 only.
+    """
+    # 1. Gather global dst keys
+    if use_dist:
+        global_dst_key_list = []
+        paddle.distributed.all_gather_object(
+            global_dst_key_list, list(state_dict_param_names), process_group
+        )
+        global_dst_keys = {
+            k for sublist in global_dst_key_list for k in sublist
+        }
+    else:
+        global_dst_keys = state_dict_param_names
+
+    # 2. Collect global src keys from metadata
+    global_src_keys = set()
+    for metadata in metadata_list:
+        for local_tensor_index in metadata.storage_metadata:
+            if (
+                local_tensor_index.replica_id is not None
+                and local_tensor_index.replica_id != 0
+            ):
+                continue
+            global_src_keys.add(local_tensor_index.tensor_key)
+
+    # 3. Compute missing / unexpected
+    missing_keys = global_dst_keys - global_src_keys
+    unexpected_keys = global_src_keys - global_dst_keys
+
+    # 4. Check shape mismatches for matching keys
+    shape_mismatches = []
+    assert state_dict is not None, "state_dict must not be None"
+    # Gather dst global shapes: {key: global_shape}
+    local_dst_shapes = {}
+    for key, val in state_dict.items():
+        k = key if isinstance(key, str) else key[0]
+        if hasattr(val, "global_shape"):
+            local_dst_shapes[k] = tuple(val.global_shape)
+        else:
+            local_dst_shapes[k] = tuple(val.shape)
+
+    if use_dist:
+        all_dst_shapes_list = []
+        paddle.distributed.all_gather_object(
+            all_dst_shapes_list, local_dst_shapes, process_group
+        )
+        global_dst_shapes = {}
+        for d in all_dst_shapes_list:
+            global_dst_shapes.update(d)
+    else:
+        global_dst_shapes = local_dst_shapes
+
+    # Build src global shapes from metadata
+    src_global_shapes: dict[str, tuple[int, ...]] = {}
+    for metadata in metadata_list:
+        if not metadata.state_dict_metadata:
+            continue
+        for key, src_metas in metadata.state_dict_metadata.items():
+            if not src_metas or src_metas[0].global_shape is None:
+                continue
+            src_global_shapes[key] = tuple(src_metas[0].global_shape)
+
+    matching_keys = global_dst_keys & global_src_keys
+    for key in sorted(matching_keys):
+        src_shape = src_global_shapes.get(key)
+        dst_shape = global_dst_shapes.get(key)
+        if src_shape is None or dst_shape is None:
+            continue
+        if src_shape != dst_shape:
+            shape_mismatches.append(
+                ShapeMismatchInfo(
+                    key=key,
+                    src_global_shape=src_shape,
+                    dst_global_shape=dst_shape,
+                )
+            )
+
+    result = KeyValidationResult(
+        missing_keys=missing_keys,
+        unexpected_keys=unexpected_keys,
+        shape_mismatches=shape_mismatches,
+        randomly_initialized_keys=set(),
+    )
+
+    # 5. Print on rank 0 (or always when not using dist)
+    if not use_dist or _get_rank() == 0:
+        _print_standard_report(result, checkpoint_path, len(global_dst_keys))
+
+    return result
+
+
+# ---------------------------------------------------------------------------
+# Public API: AOA validation
+# ---------------------------------------------------------------------------
+
+
+def validate_and_report_keys_aoa(
+    aoa_engine: AOAEngine,
+    metadata: Metadata,
+    checkpoint_path: str,
+    use_dist: bool = True,
+) -> KeyValidationResult:
+    """Validate keys for the AOA loading path.
+
+    Called AFTER AOAEngine is initialized. Uses output_vars/input_vars to
+    compute truly missing/unexpected keys and builds the mapping table.
+    """
+    # 1. Covered dst keys
+    aoa_covered_dst_keys = {
+        k for k, v in aoa_engine.output_vars.items() if v is not None
+    }
+    randomly_initialized_keys = set(aoa_engine.need_add_output_vars)
+
+    # 2. Consumed src keys
+    consumed_src_keys = set()
+    for tensor_desc in aoa_engine.output_vars.values():
+        if tensor_desc is None:
+            continue
+        for src_key, _, _, _ in tensor_desc.slices:
+            consumed_src_keys.add(src_key)
+
+    # 3. Explicitly removed / all src keys
+    explicitly_removed = set(aoa_engine.need_remove_input_vars)
+    all_src_keys = set(aoa_engine.input_vars.keys())
+
+    # 4. Compute truly missing / unexpected
+    dst_state_keys = aoa_engine.context.get_all_dst_state_keys()
+    truly_missing = (
+        dst_state_keys - aoa_covered_dst_keys - randomly_initialized_keys
+    )
+    truly_unexpected = all_src_keys - consumed_src_keys - explicitly_removed
+
+    # 5. Build AOA mapping entries
+    aoa_mappings = _build_aoa_mappings(aoa_engine)
+
+    result = KeyValidationResult(
+        missing_keys=truly_missing,
+        unexpected_keys=truly_unexpected,
+        shape_mismatches=[],
+        randomly_initialized_keys=randomly_initialized_keys,
+    )
+
+    # 6. Print on rank 0 (or always when not using dist)
+    if not use_dist or _get_rank() == 0:
+        _print_aoa_report(
+            result, aoa_mappings, explicitly_removed, checkpoint_path
+        )
+
+    return result
+
+
+def _build_aoa_mappings(aoa_engine: AOAEngine) -> list[AOAMappingEntry]:
+    """Extract mapping entries from AOA engine's output_vars."""
+    entries = []
+    for dst_key, tensor_desc in sorted(aoa_engine.output_vars.items()):
+        if tensor_desc is None:
+            continue
+        shape = tuple(tensor_desc.shape)
+        slice_mappings = []
+        for src_key, src_sl, dst_sl, pp_list in tensor_desc.slices:
+            slice_mappings.append(
+                AOASliceMapping(
+                    src_key=src_key,
+                    src_slice=src_sl,
+                    dst_slice=dst_sl,
+                    postprocess=pp_list,
+                )
+            )
+        # Determine if identity
+        is_identity = (
+            len(slice_mappings) == 1
+            and slice_mappings[0].src_key == dst_key
+            and slice_mappings[0].postprocess is None
+            and _slice_covers_full(slice_mappings[0].dst_slice, shape)
+        )
+        entries.append(
+            AOAMappingEntry(
+                dst_key=dst_key,
+                dst_global_shape=shape,
+                slice_mappings=slice_mappings,
+                is_identity=is_identity,
+            )
+        )
+    return entries
+
+
+def _slice_covers_full(sl: tuple[slice, ...], shape: tuple[int, ...]) -> bool:
+    """Check if a slice tuple covers the full tensor."""
+    if len(sl) != len(shape):
+        return False
+    for s, dim in zip(sl, shape):
+        if s.start != 0 or s.stop != dim:
+            return False
+    return True
+
+
+# ---------------------------------------------------------------------------
+# Printing: Standard report
+# ---------------------------------------------------------------------------
+
+_SEP = "=" * 70
+_THIN_SEP = "-" * 70
+
+
+def _print_standard_report(
+    result: KeyValidationResult, path: str, total_keys: int
+) -> None:
+    lines = [_SEP, f"FlexCheckpoint Load Report (Checkpoint: {path})", _SEP]
+
+    if (
+        not result.missing_keys
+        and not result.unexpected_keys
+        and not result.shape_mismatches
+    ):
+        lines.append(
+            _C.green(
+                f"[OK] All {total_keys} keys matched successfully. "
+                f"(missing: 0, unexpected: 0, shape_mismatch: 0)"
+            )
+        )
+    else:
+        matched = total_keys - len(result.missing_keys)
+        lines.append(
+            f"Matched: {matched}/{total_keys} keys | "
+            f"Missing: {len(result.missing_keys)} | "
+            f"Unexpected: {len(result.unexpected_keys)} | "
+            f"Shape mismatch: {len(result.shape_mismatches)}"
+        )
+        if result.missing_keys:
+            lines.append("")
+            lines.append(
+                _C.yellow(
+                    f"[WARNING] Missing keys ({len(result.missing_keys)} total) "
+                    f"- model expects but not in checkpoint:"
+                )
+            )
+            lines.extend(_format_key_list(result.missing_keys))
+        if result.unexpected_keys:
+            lines.append("")
+            lines.append(
+                _C.yellow(
+                    f"[WARNING] Unexpected keys ({len(result.unexpected_keys)} total) "
+                    f"- in checkpoint but not used:"
+                )
+            )
+            lines.extend(_format_key_list(result.unexpected_keys))
+        if result.shape_mismatches:
+            lines.append("")
+            lines.append(
+                _C.yellow(
+                    f"[WARNING] Shape mismatches ({len(result.shape_mismatches)} total):"
+                )
+            )
+            for m in result.shape_mismatches[:_MAX_SHAPE_MISMATCHES]:
+                lines.append(
+                    f"    {m.key}: ckpt={list(m.src_global_shape)} vs model={list(m.dst_global_shape)}"
+                )
+            remaining = len(result.shape_mismatches) - _MAX_SHAPE_MISMATCHES
+            if remaining > 0:
+                lines.append(f"    ... and {remaining} more")
+
+    lines.append(_SEP)
+    _emit(lines)
+
+
+# ---------------------------------------------------------------------------
+# Printing: AOA report
+# ---------------------------------------------------------------------------
+
+
+def _print_aoa_report(
+    result: KeyValidationResult,
+    aoa_mappings: list[AOAMappingEntry],
+    explicitly_removed: set[str],
+    path: str,
+) -> None:
+    lines = [
+        _SEP,
+        f"FlexCheckpoint Load Report (Checkpoint: {path}, AOA enabled)",
+        _SEP,
+    ]
+
+    # Status
+    total_dst = (
+        len(aoa_mappings)
+        + len(result.missing_keys)
+        + len(result.randomly_initialized_keys)
+    )
+    if not result.missing_keys and not result.unexpected_keys:
+        lines.append(
+            _C.green(
+                f"[OK] All {total_dst} keys resolved via AOA mapping. "
+                f"(missing: 0, unexpected: 0)"
+            )
+        )
+    else:
+        matched = total_dst - len(result.missing_keys)
+        lines.append(
+            f"Matched: {matched}/{total_dst} keys | "
+            f"Missing: {len(result.missing_keys)} | "
+            f"Unexpected: {len(result.unexpected_keys)}"
+        )
+        if result.missing_keys:
+            lines.append("")
+            lines.append(
+                _C.yellow(
+                    f"[WARNING] Missing keys ({len(result.missing_keys)} total) "
+                    f"- no AOA source mapping:"
+                )
+            )
+            lines.extend(_format_key_list(result.missing_keys))
+        if result.unexpected_keys:
+            lines.append("")
+            lines.append(
+                _C.yellow(
+                    f"[WARNING] Unexpected keys ({len(result.unexpected_keys)} total) "
+                    f"- in checkpoint but not consumed by any AOA mapping:"
+                )
+            )
+            lines.extend(_format_key_list(result.unexpected_keys))
+
+    # AOA mapping table
+    lines.append("")
+    lines.append(_C.cyan(_THIN_SEP))
+
+    # Classify mappings
+    non_identity = [m for m in aoa_mappings if not m.is_identity]
+    rename_only, with_transform, structural = _classify_mappings(non_identity)
+
+    total_dst = len(aoa_mappings)
+    total_src = len(
+        {sm.src_key for m in aoa_mappings for sm in m.slice_mappings}
+    )
+    lines.append(
+        _C.cyan(f"AOA Key Mapping ({total_dst} dst keys, {total_src} src keys)")
+    )
+    lines.append(_C.cyan(_THIN_SEP))
+
+    # Summary
+    lines.append("Summary:")
+    lines.append(
+        f"  1-to-1 rename (same shape, no transform):  {len(rename_only)} keys (not shown)"
+    )
+    lines.append(
+        f"  1-to-1 with transform:                     {len(with_transform)} keys "
+        f"({min(len(_group_by_signature(with_transform)), _MAX_PATTERNS_SHOWN)} pattern(s) below)"
+    )
+    lines.append(
+        f"  Structural (N-to-1 / 1-to-N / reshape):    {len(structural)} keys "
+        f"({min(len(_group_by_signature(structural)), _MAX_PATTERNS_SHOWN)} pattern(s) below)"
+    )
+
+    # Print transform patterns
+    next_index = 1
+    if with_transform:
+        lines.append("")
+        result_lines, next_index = _format_pattern_groups(
+            _group_by_signature(with_transform), "1-to-1 transform", next_index
+        )
+        lines.extend(result_lines)
+
+    # Print structural patterns
+    if structural:
+        lines.append("")
+        result_lines, next_index = _format_pattern_groups(
+            _group_by_signature(structural), "structural", next_index
+        )
+        lines.extend(result_lines)
+
+    # Removed / Initialized
+    lines.append("")
+    removed_str = ", ".join(sorted(explicitly_removed)[:5])
+    if len(explicitly_removed) > 5:
+        removed_str += f" ... +{len(explicitly_removed) - 5} more"
+    lines.append(f"Removed ({len(explicitly_removed)}): {removed_str or '-'}")
+    init_keys = result.randomly_initialized_keys
+    init_str = ", ".join(sorted(init_keys)[:5])
+    if len(init_keys) > 5:
+        init_str += f" ... +{len(init_keys) - 5} more"
+    lines.append(f"Initialized ({len(init_keys)}): {init_str or '-'}")
+
+    lines.append("")
+    lines.append(_THIN_SEP)
+    lines.append(_SEP)
+    _emit(lines)
+
+
+# ---------------------------------------------------------------------------
+# Helpers: Classification & Pattern Merging
+# ---------------------------------------------------------------------------
+
+
+def _classify_mappings(
+    non_identity: list[AOAMappingEntry],
+) -> tuple[list[AOAMappingEntry], list[AOAMappingEntry], list[AOAMappingEntry]]:
+    """Classify non-identity mappings into rename_only, with_transform, structural."""
+    rename_only = []
+    with_transform = []
+    structural = []
+    for entry in non_identity:
+        if len(entry.slice_mappings) != 1:
+            structural.append(entry)
+            continue
+        sm = entry.slice_mappings[0]
+        src_norm = re.sub(r"\d+", "{N}", sm.src_key)
+        dst_norm = re.sub(r"\d+", "{N}", entry.dst_key)
+        if src_norm != dst_norm:
+            structural.append(entry)
+        elif sm.postprocess is None:
+            rename_only.append(entry)
+        else:
+            with_transform.append(entry)
+    return rename_only, with_transform, structural
+
+
+def _get_signature(entry: AOAMappingEntry) -> str:
+    """Compute a structure signature for pattern grouping."""
+    dst_norm = re.sub(r"\d+", "{N}", entry.dst_key)
+    parts = [dst_norm, str(len(entry.slice_mappings))]
+    for sm in entry.slice_mappings:
+        src_norm = re.sub(r"\d+", "{N}", sm.src_key)
+        pp = "|".join(sm.postprocess) if sm.postprocess else ""
+        parts.append(f"{src_norm}:{pp}")
+    return "@@".join(parts)
+
+
+def _group_by_signature(
+    entries: list[AOAMappingEntry],
+) -> dict[str, list[AOAMappingEntry]]:
+    """Group entries by structure signature."""
+    groups: dict[str, list[AOAMappingEntry]] = defaultdict(list)
+    for entry in entries:
+        groups[_get_signature(entry)].append(entry)
+    return groups
+
+
+def _format_pattern_groups(
+    groups: dict[str, list[AOAMappingEntry]], label: str, start_index: int = 1
+) -> tuple[list[str], int]:
+    """Format grouped patterns with box-drawing style. Returns (lines, next_index)."""
+    lines = []
+    shown = 0
+    idx = start_index
+    for _sig, entries in sorted(groups.items(), key=lambda x: -len(x[1])):
+        if shown >= _MAX_PATTERNS_SHOWN:
+            remaining = len(groups) - shown
+            lines.append(f"  ... and {remaining} more {label} pattern(s)")
+            break
+        shown += 1
+        representative = entries[0]
+        count = len(entries)
+
+        # Build pattern title
+        dst_pattern = re.sub(r"\d+", "*", representative.dst_key)
+        lines.append(f"[Pattern #{idx}] {dst_pattern}  ({count} keys, {label})")
+        lines.append("\u250c" + "\u2500" * 69)
+        # DST line
+        shape_str = list(representative.dst_global_shape)
+        lines.append(f"\u2502 DST: {representative.dst_key}  {shape_str}")
+        # SRC lines (with folding)
+        _append_src_lines(lines, representative.slice_mappings)
+        # OP line
+        ops = _describe_ops(representative)
+        if ops:
+            lines.append(f"\u2502 OP:  {ops}")
+        lines.append("\u2514" + "\u2500" * 69)
+        lines.append("")
+        idx += 1
+    return lines, idx
+
+
+def _append_src_lines(
+    lines: list[str], slice_mappings: list[AOASliceMapping]
+) -> None:
+    """Append SRC lines, folding consecutive numeric patterns."""
+    if len(slice_mappings) <= _SRC_FOLD_THRESHOLD:
+        for i, sm in enumerate(slice_mappings):
+            prefix = "\u2502 SRC:" if i == 0 else "\u2502    +"
+            slice_info = _format_slice_range(sm.src_slice, sm.dst_slice)
+            lines.append(f"{prefix} {sm.src_key}{slice_info}")
+        return
+
+    # Try to fold: find common pattern
+    src_keys = [sm.src_key for sm in slice_mappings]
+    folded = _try_fold_src_keys(src_keys)
+    if folded:
+        lines.append(f"\u2502 SRC: {folded}  (\u00d7{len(slice_mappings)})")
+    else:
+        # Show first 2 and last 1
+        lines.append(f"\u2502 SRC: {src_keys[0]}")
+        lines.append(f"\u2502    + {src_keys[1]}")
+        lines.append(f"\u2502    + ... ({len(src_keys) - 3} more)")
+        lines.append(f"\u2502    + {src_keys[-1]}")
+
+
+def _format_slice_range(
+    src_slice: tuple[slice, ...], dst_slice: tuple[slice, ...]
+) -> str:
+    """Format slice info when same src_key appears multiple times."""
+    src_str = ",".join(f"{s.start}:{s.stop}" for s in src_slice)
+    dst_str = ",".join(f"{s.start}:{s.stop}" for s in dst_slice)
+    return f"  [{src_str}] -> dst[{dst_str}]"
+
+
+def _try_fold_src_keys(keys: list[str]) -> str | None:
+    """Try to fold src keys like experts.0, experts.1, ..., experts.255 into a pattern."""
+    if len(keys) < 2:
+        return None
+    # Find varying digit segments
+    pattern = re.sub(r"\d+", "{}", keys[0])
+    for k in keys[1:]:
+        if re.sub(r"\d+", "{}", k) != pattern:
+            return None
+    # Extract the varying numbers
+    nums_per_key = [re.findall(r"\d+", k) for k in keys]
+    num_positions = len(nums_per_key[0])
+    # Find which position varies
+    varying_pos = []
+    for pos in range(num_positions):
+        vals = [int(n[pos]) for n in nums_per_key]
+        if len(set(vals)) > 1:
+            varying_pos.append(pos)
+    if len(varying_pos) != 1:
+        return None
+    vpos = varying_pos[0]
+    vals = [int(n[vpos]) for n in nums_per_key]
+    lo, hi = min(vals), max(vals)
+    # Reconstruct pattern with {lo..hi}
+    segments = re.split(r"\d+", keys[0])
+    digits = re.findall(r"\d+", keys[0])
+    result_parts = []
+    for i, seg in enumerate(segments):
+        result_parts.append(seg)
+        if i < len(digits):
+            if i == vpos:
+                result_parts.append(f"{{{lo}..{hi}}}")
+            else:
+                result_parts.append(digits[i])
+    return "".join(result_parts)
+
+
+def _describe_ops(entry: AOAMappingEntry) -> str:
+    """Describe the operations for a mapping entry."""
+    ops = []
+    if len(entry.slice_mappings) > 1:
+        ops.append("concat")
+    # Collect postprocess from first slice (representative)
+    if entry.slice_mappings:
+        pp = entry.slice_mappings[0].postprocess
+        if pp:
+            for p in pp:
+                if p.startswith("["):
+                    ops.append(f"permute({p})")
+                else:
+                    ops.append(f"cast({p})")
+    return " + ".join(ops)
+
+
+# ---------------------------------------------------------------------------
+# Helpers: Key list formatting
+# ---------------------------------------------------------------------------
+
+
+def _format_key_list(keys: set[str]) -> list[str]:
+    """Format a set of keys with prefix grouping and truncation."""
+    if not keys:
+        return []
+    sorted_keys = sorted(keys)
+    if len(sorted_keys) <= _MAX_KEYS_SHOWN:
+        return [f"    {k}" for k in sorted_keys]
+
+    # Adaptive grouping: find the prefix depth that gives reasonable group sizes
+    groups = _group_keys_adaptive(sorted_keys)
+
+    lines = []
+    groups_shown = 0
+    for prefix, group_keys in sorted(groups.items(), key=lambda x: -len(x[1])):
+        if groups_shown >= _MAX_KEYS_SHOWN:
+            remaining_groups = len(groups) - groups_shown
+            remaining_keys = sum(
+                len(v)
+                for i, v in enumerate(
+                    sorted(groups.values(), key=len, reverse=True)
+                )
+                if i >= groups_shown
+            )
+            lines.append(
+                f"    ... and {remaining_groups} more groups ({remaining_keys} keys)"
+            )
+            break
+        groups_shown += 1
+        if len(group_keys) > 3:
+            lines.append(f"    [{prefix}] ({len(group_keys)} keys):")
+            for k in group_keys[:3]:
+                lines.append(f"      {k}")
+            lines.append(f"      ... +{len(group_keys) - 3} more")
+        else:
+            for k in group_keys:
+                lines.append(f"    {k}")
+    return lines
+
+
+def _group_keys_adaptive(keys: list[str]) -> dict[str, list[str]]:
+    """Group keys by normalized pattern (digits replaced with *)."""
+    groups: dict[str, list[str]] = defaultdict(list)
+    for k in keys:
+        # Replace all digit segments with * to get the pattern
+        pattern = re.sub(r"(?<=\.)\d+(?=\.)|(?<=\.)\d+$", "*", k)
+        groups[pattern].append(k)
+    return dict(groups)
+
+
+# ---------------------------------------------------------------------------
+# Helpers: Output
+# ---------------------------------------------------------------------------
+
+
+def _emit(lines: list[str]) -> None:
+    """Output lines via logger, respecting total line limit."""
+    for i, line in enumerate(lines):
+        if i >= _MAX_TOTAL_LINES:
+            logger.info(
+                f"... output truncated ({len(lines) - i} lines omitted)"
+            )
+            break
+        logger.info(line)

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

@@ -0,0 +1,51 @@
+# Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+
+@dataclass(frozen=True)
+class LocalTensorMetadata:
+    """
+    The location of a local tensor in the global tensor.
+    """
+
+    global_offset: tuple[int]
+    local_shape: tuple[int]
+    dtype: str
+    global_shape: tuple[int] | None = None
+    is_flattened: bool = False
+    flattened_range: tuple[int] | None = None
+
+
+@dataclass(frozen=True)
+class LocalTensorIndex:
+    """
+    The identifier of a local tensor.
+    """
+
+    tensor_key: str
+    global_offset: tuple[int]
+    is_flattened: bool = False
+    flattened_range: tuple[int] | None = None
+    replica_id: int | None = None
+    local_shape: tuple[int] | None = None
+
+
+@dataclass
+class Metadata:
+    state_dict_metadata: dict[str, list[LocalTensorMetadata]] = None
+    storage_metadata: dict[LocalTensorIndex, str] = None
+    flat_mapping: dict[str, tuple[str]] = None

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

@@ -0,0 +1,126 @@
+# Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from collections import defaultdict
+
+from .metadata import LocalTensorIndex, LocalTensorMetadata, Metadata
+
+TensorLocation = tuple[str, str]
+
+
+class MetadataManager:
+    def __init__(self):
+        self._metadata_list: list[Metadata] = []
+        self.local_tensor_metadata: dict[
+            TensorLocation, LocalTensorMetadata
+        ] = {}
+        self.has_flattened_tensors: bool = False
+        self.file_storage_info: defaultdict[str, set[LocalTensorIndex]] = (
+            defaultdict(set)
+        )
+
+    def set_metadata_list(self, metadata_list: list[Metadata]):
+        assert len(metadata_list) == 1, "Only support single metadata list"
+        self.clear()
+
+        self.local_tensor_metadata = {}
+        self.has_flattened_tensors = False
+
+        self._metadata_list = metadata_list
+        self._extract_local_tensor_metadata()
+        self._extract_file_storage_info()
+
+    def get_metadata_list(self) -> list[Metadata]:
+        return self._metadata_list
+
+    def is_metadata_list_empty(self) -> bool:
+        return not self._metadata_list
+
+    def get_flat_mapping(self) -> dict:
+        if self.is_metadata_list_empty():
+            raise ValueError(
+                "Cannot get flat mapping because metadata list is empty."
+            )
+        return self._metadata_list[0].flat_mapping
+
+    def get_file_storage_info(self) -> defaultdict:
+        if self.is_metadata_list_empty():
+            raise ValueError(
+                "Cannot get file_storage_info because metadata list is empty."
+            )
+        return self.file_storage_info
+
+    def _extract_local_tensor_metadata(self):
+        if self.is_metadata_list_empty():
+            return
+
+        metadata = self._metadata_list[0]
+        state_dict_metadata = metadata.state_dict_metadata
+        storage_metadata = metadata.storage_metadata
+
+        storage_metadata_split_replica_id = {}
+        for local_tensor_index, file_name in storage_metadata.items():
+            local_tensor_index = LocalTensorIndex(
+                tensor_key=local_tensor_index.tensor_key,
+                global_offset=local_tensor_index.global_offset,
+                is_flattened=local_tensor_index.is_flattened,
+                flattened_range=local_tensor_index.flattened_range,
+                local_shape=local_tensor_index.local_shape,
+            )
+            replica_id = local_tensor_index.replica_id
+            storage_metadata_split_replica_id[local_tensor_index] = (
+                file_name,
+                replica_id,
+            )
+
+        for k, local_tensor_meta_list in state_dict_metadata.items():
+            for local_tensor_meta in local_tensor_meta_list:
+                local_tensor_index = LocalTensorIndex(
+                    tensor_key=k,
+                    global_offset=local_tensor_meta.global_offset,
+                    is_flattened=local_tensor_meta.is_flattened,
+                    flattened_range=local_tensor_meta.flattened_range,
+                    local_shape=local_tensor_meta.local_shape,
+                )
+
+                if local_tensor_meta.is_flattened:
+                    self.has_flattened_tensors = True
+
+                if local_tensor_index not in storage_metadata_split_replica_id:
+                    continue
+
+                file_name, replica_id = storage_metadata_split_replica_id[
+                    local_tensor_index
+                ]
+                if replica_id is not None and replica_id > 0:
+                    continue
+
+                location_key: TensorLocation = (k, file_name)
+
+                self.local_tensor_metadata[location_key] = local_tensor_meta
+
+    def _extract_file_storage_info(self):
+        if self.is_metadata_list_empty():
+            return
+
+        metadata = self._metadata_list[0]
+        storage_metadata = metadata.storage_metadata
+        for local_tensor_index, file_name in storage_metadata.items():
+            self.file_storage_info[file_name].add(local_tensor_index)
+
+    def clear(self):
+        self._metadata_list = []
+        self.local_tensor_metadata = {}
+        self.has_flattened_tensors = False
+        self.file_storage_info = defaultdict(set)

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

@@ -0,0 +1,721 @@
+# Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import types
+from abc import ABC, abstractmethod
+from collections import defaultdict
+from dataclasses import replace
+
+import paddle
+import paddle.distributed as dist
+from paddle.distributed.collective import Group
+from paddle.distributed.fleet.utils.log_util import logger
+
+from .resharder import ReadItem
+from .utils import (
+    get_target_tensor,
+    slice_tensor,
+)
+
+GROUPED_BATCH_SIZE = 10
+
+
+class CommunicatorFactory:
+    registry = {}
+
+    @classmethod
+    def register(cls, method, creator):
+        cls.registry[method] = creator
+
+    @classmethod
+    def create(cls, comm_method, **kwargs):
+        if comm_method not in cls.registry:
+            raise ValueError(
+                f"Unknown communication method '{comm_method}'. "
+                f"Available: {list(cls.registry.keys())}"
+            )
+        return cls.registry[comm_method](**kwargs)
+
+
+class AbstractCommunicator(ABC):
+    @staticmethod
+    def schedule_read_items(
+        read_items: list[ReadItem],
+    ) -> dict[str, list[ReadItem]]:
+        order_rules = lambda read_item: (
+            read_item.tensor_name,
+            read_item.src_rank,
+            read_item.src_global_offset,
+            read_item.dst_rank,
+            read_item.dst_local_offset,
+            read_item.dst_global_offset
+            if read_item.dst_global_offset is not None
+            else (),
+            read_item.src_local_offset,
+            read_item.slice_shape,
+            read_item.file_name,
+            read_item.dtype,
+        )
+        # Step 1: Group by tensor_name
+        tensor_groups = defaultdict(list)
+        for item in read_items:
+            tensor_groups[item.tensor_name].append(item)
+
+        scheduled_items = defaultdict(list)
+
+        # Step 2: For each tensor_name group, further group by all attributes except dst_rank
+        for tensor_name, items in tensor_groups.items():
+            grouped_items = defaultdict(list)
+            for item in items:
+                key = (
+                    item.src_global_offset,
+                    item.dst_global_offset,
+                    item.src_rank,
+                    item.dst_local_offset,
+                    item.src_local_offset,
+                    item.slice_shape,
+                    item.file_name,
+                    item.dtype,
+                )
+                grouped_items[key].append(item)
+
+            # Step 3: Combine items with the same key into a single ReadItem with all dst_ranks
+            for key, grouped_item in grouped_items.items():
+                combined_dst_rank = []
+                for item in grouped_item:
+                    combined_dst_rank.extend(item.dst_rank)
+                combined_dst_rank = sorted(
+                    set(combined_dst_rank)
+                )  # Remove duplicates
+
+                # Create a new ReadItem with combined dst_ranks
+                scheduled_item = ReadItem(
+                    tensor_name=tensor_name,
+                    src_global_offset=key[0],
+                    dst_global_offset=key[1],
+                    dst_rank=tuple(combined_dst_rank),
+                    src_rank=key[2],
+                    dst_local_offset=key[3],
+                    src_local_offset=key[4],
+                    slice_shape=key[5],
+                    file_name=key[6],
+                    dtype=key[7],
+                )
+                scheduled_items[tensor_name].append(scheduled_item)
+        for key, items in scheduled_items.items():
+            scheduled_items[key] = sorted(items, key=order_rules)
+
+        return dict(sorted(scheduled_items.items()))
+
+    @staticmethod
+    def split_read_items(
+        read_items: list[ReadItem],
+    ) -> (list[ReadItem], list[ReadItem]):
+        local_read_items = []
+        comm_read_items = []
+
+        for item in read_items:
+            assert len(item.dst_rank) == 1, (
+                "Before read_items is split, each ReadItem describes a communication task between one rank and another."
+            )
+            if item.src_rank == item.dst_rank[0]:
+                local_read_items.append(item)
+            else:
+                comm_read_items.append(item)
+
+        return local_read_items, comm_read_items
+
+    @staticmethod
+    def process_local_copy_tasks(
+        local_tasks, cur_rank, source_state_dict, target_state_dict
+    ):
+        """
+        Complete local copy tasks.
+        """
+        logger.debug(
+            f"Rank {cur_rank} starting local copy for {len(local_tasks)} tasks."
+        )
+        for task in local_tasks:
+            if task.src_rank != cur_rank:
+                continue
+
+            src_tensor = source_state_dict[task.file_name][task.tensor_name]
+            dst_tensor = get_target_tensor(target_state_dict, task)
+            src_chunk_tensor = slice_tensor(
+                src_tensor, task.src_local_offset, task.slice_shape
+            )
+
+            dst_chunk_tensor = slice_tensor(
+                dst_tensor, task.dst_local_offset, task.slice_shape
+            )
+            if src_chunk_tensor.place == dst_chunk_tensor.place:
+                paddle.assign(src_chunk_tensor, dst_chunk_tensor)
+                logger.debug(f"Local copy (same device) for task {task}.")
+            else:
+                tmp = (
+                    src_chunk_tensor.cuda()
+                    if dst_chunk_tensor.place.is_gpu_place()
+                    else src_chunk_tensor.cpu()
+                )
+                paddle.assign(tmp, dst_chunk_tensor)
+                del tmp
+                logger.debug(f"Local copy (cross device) for task {task}.")
+
+    @abstractmethod
+    def communicate(self, read_items, state, context):
+        pass
+
+
+class BroadcastCommunicator(AbstractCommunicator):
+    """
+    Communicator that uses broadcast operation for data transfer.
+    """
+
+    def communicate(self, read_items, state, context):
+        cur_rank = context['rank']
+        process_group = context['process_group']
+
+        source_state_dict = state['source_state_dict']
+        target_state_dict = state['target_state_dict']
+
+        local_read_items, comm_read_items = (
+            BroadcastCommunicator.split_read_items(read_items)
+        )
+
+        logger.info(f"Generated {len(comm_read_items)} communication tasks.")
+        logger.info(f"Generated {len(local_read_items)} local tasks.")
+
+        BroadcastCommunicator.process_local_copy_tasks(
+            local_read_items,
+            cur_rank,
+            source_state_dict,
+            target_state_dict,
+        )
+
+        logger.info(
+            f"Rank {cur_rank} finished local copy and entered communication phase."
+        )
+
+        comm_tasks = BroadcastCommunicator.schedule_read_items(comm_read_items)
+        cnt = 0
+        total_task_len = len(comm_tasks)
+        for tensor_name, read_items in comm_tasks.items():
+            cnt += 1
+            if cnt % 500 == 0 or cnt == total_task_len:
+                logger.info(
+                    f"{cnt}/{total_task_len} tasks have been sent/received successfully!"
+                )
+
+            source_tensors = {}
+            destination_tensors = {}
+            for item in read_items:
+                logger.debug(f"Beginning to send/recv task {item}.")
+                if item.src_rank == cur_rank:
+                    src_tensor = source_state_dict[item.file_name][
+                        item.tensor_name
+                    ]
+                    if not src_tensor.place.is_gpu_place():
+                        src_tensor = src_tensor.cuda()
+                    source_tensors[(tensor_name, item.file_name)] = src_tensor
+                elif cur_rank in item.dst_rank:
+                    dst_tensor = get_target_tensor(target_state_dict, item)
+                    if not dst_tensor.place.is_gpu_place():
+                        gpu_dst_tensor = dst_tensor.cuda()
+                        gpu_dst_tensor.need_cross_device_copy = True
+                        gpu_dst_tensor.target_tensor = dst_tensor
+                        destination_tensors[
+                            (tensor_name, cur_rank, item.dst_global_offset)
+                        ] = gpu_dst_tensor
+                    else:
+                        gpu_dst_tensor = dst_tensor
+                        gpu_dst_tensor.target_tensor = dst_tensor
+                        destination_tensors[
+                            (tensor_name, cur_rank, item.dst_global_offset)
+                        ] = dst_tensor
+
+            for item in read_items:
+                logger.debug(f"Beginning to send/recv task {item}.")
+                if item.src_rank == cur_rank:
+                    src_tensor = source_tensors[(tensor_name, item.file_name)]
+                    src_chunk_tensor = slice_tensor(
+                        src_tensor, item.src_local_offset, item.slice_shape
+                    )
+                    buffer_tensor = src_chunk_tensor.contiguous()
+                elif cur_rank in item.dst_rank:
+                    dst_tensor = destination_tensors[
+                        (tensor_name, cur_rank, item.dst_global_offset)
+                    ]
+                    dst_chunk_tensor = slice_tensor(
+                        dst_tensor, item.dst_local_offset, item.slice_shape
+                    )
+                    buffer_tensor = paddle.zeros_like(dst_chunk_tensor)
+                    paddle.assign(dst_chunk_tensor, buffer_tensor)
+
+                else:
+                    buffer_tensor = paddle.zeros(item.slice_shape, item.dtype)
+                paddle.distributed.broadcast(
+                    buffer_tensor, src=item.src_rank, group=process_group
+                )
+                if cur_rank in item.dst_rank:
+                    paddle.assign(buffer_tensor, dst_chunk_tensor)
+                del buffer_tensor
+
+            for dst_tensor in destination_tensors.values():
+                if getattr(dst_tensor, 'need_cross_device_copy', False):
+                    target_tensor = dst_tensor.target_tensor
+                    delattr(dst_tensor, "target_tensor")
+                    target_tensor.copy_(dst_tensor)
+                else:
+                    target_tensor = dst_tensor.target_tensor
+                    delattr(dst_tensor, "target_tensor")
+                    paddle.assign(dst_tensor, target_tensor)
+                del dst_tensor
+
+            del source_tensors
+            paddle.distributed.barrier(process_group)
+
+        logger.info("All communication tasks completed.")
+
+
+class MultiGroupBroadcastCommunicator(AbstractCommunicator):
+    """
+    Communicator that uses broadcast for data transfer across multiple communication groups.
+    """
+
+    def __init__(self, worker_groups):
+        if worker_groups is None:
+            raise ValueError(
+                "worker_groups must be specified when using multi_group_broadcast."
+            )
+        self.worker_groups = worker_groups
+
+    @staticmethod
+    def schedule_read_items(
+        comm_read_items: list[ReadItem],
+        worker_groups: list[Group],
+    ) -> list[list[ReadItem]]:
+        group_members = {}
+        name_to_groups = {}
+        read_items = []
+
+        order_rules = lambda read_item: (
+            read_item.tensor_name,
+            read_item.src_rank,
+            read_item.src_global_offset,
+            read_item.dst_rank,
+            read_item.dst_local_offset,
+            read_item.dst_global_offset
+            if read_item.dst_global_offset is not None
+            else (),
+            read_item.src_local_offset,
+            read_item.slice_shape,
+            read_item.file_name,
+            read_item.dtype,
+        )
+
+        def _find_min_group(need_ranks, group_members, name_to_groups):
+            min_group = None
+            min_size = None
+            for name, ranks in group_members.items():
+                if need_ranks <= ranks:
+                    if (min_size is None) or (len(ranks) < min_size):
+                        min_size = len(ranks)
+                        min_group = name_to_groups[name]
+            assert min_group is not None, f"No group found for {need_ranks}!"
+            return min_group
+
+        for group in worker_groups:
+            if len(group.ranks) <= 1:
+                continue
+            group_members[group.name] = set(group.ranks)
+            name_to_groups[group.name] = group
+
+        for read_item in comm_read_items:
+            need_ranks = need_ranks = {*read_item.dst_rank, read_item.src_rank}
+            group = _find_min_group(
+                need_ranks,
+                group_members,
+                name_to_groups,
+            )
+            read_items.append(replace(read_item, comm_group=group))
+
+        read_items = sorted(read_items, key=order_rules)
+
+        def _build_group_conflict(group_members: dict[str, set]):
+            member_to_groups = defaultdict(set)
+            for g, members in group_members.items():
+                for m in members:
+                    member_to_groups[m].add(g)
+            group_conflict = defaultdict(set)
+            for group_set in member_to_groups.values():
+                for g1 in group_set:
+                    for g2 in group_set:
+                        if g1 != g2:
+                            group_conflict[g1].add(g2)
+            return group_conflict
+
+        def _dsatur_coloring(group_conflict: dict[str, set]) -> dict[str, int]:
+            import heapq
+
+            all_groups = sorted(group_conflict.keys())
+            sorted_conflict = {g: sorted(group_conflict[g]) for g in all_groups}
+
+            color_map = {}
+            neighbor_colors = {g: set() for g in all_groups}
+            uncolored = set(all_groups)
+
+            degree = {g: len(sorted_conflict[g]) for g in all_groups}
+
+            heap = []
+            for g in all_groups:
+                heapq.heappush(heap, (0, -degree[g], g))
+            saturation = dict.fromkeys(all_groups, 0)
+
+            while uncolored:
+                while True:
+                    _, _, node = heapq.heappop(heap)
+                    if node in uncolored:
+                        break
+                used = neighbor_colors[node]
+                color = 0
+                while color in used:
+                    color += 1
+                color_map[node] = color
+                uncolored.remove(node)
+                for neighbor in sorted_conflict[node]:
+                    if neighbor in uncolored:
+                        if color not in neighbor_colors[neighbor]:
+                            neighbor_colors[neighbor].add(color)
+                            saturation[neighbor] += 1
+                            heapq.heappush(
+                                heap,
+                                (
+                                    -saturation[neighbor],
+                                    -degree[neighbor],
+                                    neighbor,
+                                ),
+                            )
+            return color_map
+
+        def _assign_batches(tasks, group_color_map):
+            batches = defaultdict(list)
+            for t in tasks:
+                g = t.comm_group.name
+                batches[group_color_map[g]].append(t)
+            return [
+                sorted(batches[c], key=order_rules) for c in sorted(batches)
+            ]
+
+        group_conflict = _build_group_conflict(group_members)
+        group_color_map = _dsatur_coloring(group_conflict)
+        results = _assign_batches(read_items, group_color_map)
+        return results
+
+    def communicate(self, read_items, state, context):
+        cur_rank = context['rank']
+        process_group = context['process_group']
+        worker_groups = self.worker_groups
+
+        source_state_dict = state['source_state_dict']
+        target_state_dict = state['target_state_dict']
+
+        local_read_items, comm_read_items = (
+            MultiGroupBroadcastCommunicator.split_read_items(read_items)
+        )
+
+        logger.info(f"Generated {len(comm_read_items)} communication tasks.")
+        logger.info(f"Generated {len(local_read_items)} local tasks.")
+
+        MultiGroupBroadcastCommunicator.process_local_copy_tasks(
+            local_read_items,
+            cur_rank,
+            source_state_dict,
+            target_state_dict,
+        )
+        results = MultiGroupBroadcastCommunicator.schedule_read_items(
+            comm_read_items, worker_groups
+        )
+        logger.info(
+            f"Communication task scheduling completed, {len(results)}  batches in total."
+        )
+        for read_items in results:
+            source_tensors = {}
+            destination_tensors = {}
+            for item in read_items:
+                tensor_name = item.tensor_name
+                if item.src_rank == cur_rank:
+                    src_tensor = source_state_dict[item.file_name][tensor_name]
+                    if not src_tensor.place.is_gpu_place():
+                        src_tensor = src_tensor.cuda()
+                    source_tensors[(tensor_name, item.file_name)] = src_tensor
+                elif cur_rank in item.dst_rank:
+                    dst_tensor = get_target_tensor(target_state_dict, item)
+                    if not dst_tensor.place.is_gpu_place():
+                        gpu_dst_tensor = dst_tensor.cuda()
+                        gpu_dst_tensor.need_cross_device_copy = True
+                        gpu_dst_tensor.target_tensor = dst_tensor
+                        destination_tensors[
+                            (tensor_name, cur_rank, item.dst_global_offset)
+                        ] = gpu_dst_tensor
+                    else:
+                        gpu_dst_tensor = dst_tensor
+                        gpu_dst_tensor.target_tensor = dst_tensor
+                        destination_tensors[
+                            (tensor_name, cur_rank, item.dst_global_offset)
+                        ] = dst_tensor
+
+            for item in read_items:
+                logger.debug(f"Beginning to send/recv task {item}.")
+                tensor_name = item.tensor_name
+                if item.src_rank == cur_rank:
+                    src_tensor = source_tensors[(tensor_name, item.file_name)]
+                    src_chunk_tensor = slice_tensor(
+                        src_tensor, item.src_local_offset, item.slice_shape
+                    )
+                    buffer_tensor = src_chunk_tensor.contiguous()
+                elif cur_rank in item.dst_rank:
+                    dst_tensor = destination_tensors[
+                        (tensor_name, cur_rank, item.dst_global_offset)
+                    ]
+                    dst_chunk_tensor = slice_tensor(
+                        dst_tensor, item.dst_local_offset, item.slice_shape
+                    )
+                    buffer_tensor = paddle.zeros_like(dst_chunk_tensor)
+                    paddle.assign(dst_chunk_tensor, buffer_tensor)
+
+                elif cur_rank in item.comm_group.ranks:
+                    buffer_tensor = paddle.zeros(item.slice_shape, item.dtype)
+                else:
+                    buffer_tensor = None
+
+                if cur_rank in item.comm_group.ranks:
+                    paddle.distributed.broadcast(
+                        buffer_tensor, src=item.src_rank, group=item.comm_group
+                    )
+
+                if cur_rank in item.dst_rank:
+                    paddle.assign(buffer_tensor, dst_chunk_tensor)
+                del buffer_tensor
+
+            for dst_tensor in destination_tensors.values():
+                if getattr(dst_tensor, 'need_cross_device_copy', False):
+                    target_tensor = dst_tensor.target_tensor
+                    delattr(dst_tensor, "target_tensor")
+                    target_tensor.copy_(dst_tensor)
+                else:
+                    target_tensor = dst_tensor.target_tensor
+                    delattr(dst_tensor, "target_tensor")
+                    paddle.assign(dst_tensor, target_tensor)
+                del dst_tensor
+
+            del source_tensors
+
+        paddle.distributed.barrier(process_group)
+        logger.info("All communication tasks completed.")
+
+
+class SendRecvCommunicator(AbstractCommunicator):
+    """
+    Communicator that uses send/recv operations for data transfer.
+
+    The process is broken down into batches to manage memory and communication overhead.
+    """
+
+    def __init__(self, use_group):
+        self.use_group = use_group
+
+    @staticmethod
+    def schedule_read_items(
+        read_items: list[ReadItem],
+    ) -> dict[str, list[ReadItem]]:
+        order_rules = lambda read_item: (
+            read_item.tensor_name,
+            read_item.src_rank,
+            read_item.src_global_offset,
+            read_item.dst_rank,
+            read_item.dst_local_offset,
+            read_item.dst_global_offset
+            if read_item.dst_global_offset is not None
+            else (),
+            read_item.src_local_offset,
+            read_item.slice_shape,
+            read_item.file_name,
+            read_item.dtype,
+        )
+
+        tensor_groups = defaultdict(list)
+        for item in read_items:
+            tensor_groups[item.tensor_name].append(item)
+
+        return dict(sorted(tensor_groups.items()))
+
+    def communicate(self, read_items, state, context):
+        comm_tasks = SendRecvCommunicator.schedule_read_items(read_items)
+        cur_rank = context['rank']
+        process_group = context['process_group']
+
+        source_state_dict = state['source_state_dict']
+        target_state_dict = state['target_state_dict']
+
+        total_items = sum(len(items) for items in comm_tasks.values())
+        processed_items = 0
+
+        for batch_data in self._process_batches(
+            comm_tasks, cur_rank, source_state_dict
+        ):
+            received_slices = {}
+            self._execute_p2p_ops(
+                batch_data, cur_rank, use_group=self.use_group
+            )
+
+            for item, tensor in batch_data.source_slices.items():
+                if item not in batch_data.local_copy_tasks:
+                    tensor._clear()
+
+            received_slices.update(batch_data.target_slices)
+
+            processed_items += len(batch_data.read_items)
+            progress = processed_items / total_items * 100
+            logger.info(
+                f"Batch communication completed. Progress: {processed_items}/{total_items} ({progress:.1f}%)."
+            )
+
+            self._assign_received_data(received_slices, target_state_dict)
+
+            for received_slice in received_slices.values():
+                received_slice._clear()
+
+            del received_slices
+
+        if self.use_group:
+            paddle.distributed.barrier(process_group)
+        logger.info("All communication tasks completed successfully.")
+
+    def _process_batches(self, comm_tasks, cur_rank, source_state_dict):
+        total_items = sum(len(items) for items in comm_tasks.values())
+        item_count = 0
+
+        batch_read_items = []
+        batch_source_slices = {}
+        batch_target_slices = {}
+        batch_local_copy_tasks = set()
+
+        for tensor_name, read_items in comm_tasks.items():
+            tensors_to_clear = set()
+            for item in read_items:
+                item_count += 1
+                batch_read_items.append(item)
+                if cur_rank == item.src_rank:
+                    src_tensor = source_state_dict[item.file_name][
+                        item.tensor_name
+                    ]
+                    src_slice = (
+                        slice_tensor(
+                            src_tensor, item.src_local_offset, item.slice_shape
+                        )
+                        .cuda()
+                        .clone()
+                    )
+                    batch_source_slices[item] = src_slice
+                    tensors_to_clear.add(src_tensor)
+                if cur_rank in item.dst_rank:
+                    if cur_rank == item.src_rank:
+                        batch_local_copy_tasks.add(item)
+                        batch_target_slices[item] = batch_source_slices[item]
+                    else:
+                        dst_slice = paddle.zeros(
+                            item.slice_shape, dtype=item.dtype
+                        )
+                        batch_target_slices[item] = dst_slice
+
+                if ((item_count % GROUPED_BATCH_SIZE) == 0) or (
+                    item_count == total_items
+                ):
+                    batch_data = types.SimpleNamespace(
+                        read_items=batch_read_items,
+                        source_slices=batch_source_slices,
+                        target_slices=batch_target_slices,
+                        local_copy_tasks=batch_local_copy_tasks,
+                    )
+                    yield batch_data
+                    batch_read_items = []
+                    batch_source_slices = {}
+                    batch_target_slices = {}
+                    batch_local_copy_tasks = set()
+
+            for tensor in tensors_to_clear:
+                tensor._clear_to_zero_allocation()
+
+    def _execute_p2p_ops(self, batch_data, cur_rank, use_group):
+        p2p_ops = []
+        for item in batch_data.read_items:
+            if item.src_rank == cur_rank:
+                for rank in item.dst_rank:
+                    if rank != cur_rank:
+                        send_tensor = batch_data.source_slices[item]
+                        if use_group:
+                            p2p_ops.append(
+                                dist.P2POp(dist.isend, send_tensor, rank)
+                            )
+                        else:
+                            dist.send(send_tensor, rank)
+
+            if cur_rank in item.dst_rank and item.src_rank != cur_rank:
+                recv_tensor = batch_data.target_slices[item]
+                if use_group:
+                    p2p_ops.append(
+                        dist.P2POp(dist.irecv, recv_tensor, item.src_rank)
+                    )
+                else:
+                    dist.recv(recv_tensor, item.src_rank)
+
+        if use_group and p2p_ops:
+            logger.info(
+                f"Starting batched send/recv for {len(p2p_ops)} P2P operations."
+            )
+            reqs = dist.batch_isend_irecv(p2p_ops)
+            for req in reqs:
+                req.wait()
+            logger.info("Batched send/recv finished.")
+
+    def _assign_received_data(self, received_slices, target_state_dict):
+        for item, received_slice in received_slices.items():
+            dest_tensor = get_target_tensor(target_state_dict, item)
+            if not dest_tensor._is_initialized():
+                buffer = paddle.zeros_like(dest_tensor)
+                buffer._share_buffer_to(dest_tensor)
+
+            dest_slice = slice_tensor(
+                dest_tensor, item.dst_local_offset, item.slice_shape
+            )
+
+            if dest_slice.place != received_slice.place:
+                received_slice = received_slice.to(dest_slice.place)
+
+            paddle.assign(received_slice, dest_slice)
+
+
+CommunicatorFactory.register(
+    "multi_group_broadcast",
+    lambda worker_groups: MultiGroupBroadcastCommunicator(worker_groups),
+)
+CommunicatorFactory.register(
+    "send_recv", lambda **kwargs: SendRecvCommunicator(use_group=False)
+)
+CommunicatorFactory.register(
+    "grouped_send_recv", lambda **kwargs: SendRecvCommunicator(use_group=True)
+)
+CommunicatorFactory.register(
+    "broadcast", lambda **kwargs: BroadcastCommunicator()
+)

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

@@ -0,0 +1,323 @@
+# Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from __future__ import annotations
+
+import multiprocessing
+import os
+import time
+from collections import defaultdict
+from dataclasses import replace
+from typing import TYPE_CHECKING
+
+import paddle
+from paddle.distributed.communication.group import is_initialized
+from paddle.distributed.fleet.utils.log_util import logger
+
+from .metadata import LocalTensorIndex, Metadata
+from .sharded_weight import (
+    ShardedWeight,
+)
+from .utils import (
+    check_unique_id,
+    extract_tensor_metadata,
+    flatten_state_dict,
+    get_max_id,
+    merge_state_dict_metadata,
+    write_to_file_if_empty,
+)
+
+if TYPE_CHECKING:
+    from paddle import Tensor
+    from paddle.distributed.collective import Group
+async_save_queue = []
+
+
+def check_exitcode(task):
+    exitcode = task.exitcode
+    if exitcode != 0:
+        logger.error(
+            f"Error: save ckpt process failed with exitcode {exitcode}!!!"
+        )
+
+
+def clear_async_save_task_queue():
+    """
+    wait until all async save task to be done.
+    """
+    while len(async_save_queue) > 0:
+        task = async_save_queue.pop()
+        if task and task.is_alive():
+            task.join(timeout=60)
+            if task.is_alive():
+                logger.error("Error: save ckpt process timeout!!!")
+                async_save_queue.append(task)
+            else:
+                check_exitcode(task)
+        else:
+            check_exitcode(task)
+
+
+def copy_dict_to_cpu(nested_dict):
+    """
+    Copy the paddle.Tensor objects in the nested dictionary to the CPU and return a new dict.
+    """
+    new_dict = {}
+    for key, value in nested_dict.items():
+        if isinstance(value, paddle.Tensor):
+            new_dict[key] = value.cpu()
+            paddle.device.synchronize()
+        elif isinstance(value, dict):
+            new_dict[key] = copy_dict_to_cpu(value)
+        else:
+            new_dict[key] = value
+    return new_dict
+
+
+def dedup_key_in_dict(global_storage_metadata):
+    out = {}
+    for storage_metadata in global_storage_metadata:
+        for key, val in storage_metadata.items():
+            if key in out:
+                continue
+            out[key] = val
+    return out
+
+
+def balanced_dedup_key_in_dict(global_storage_metadata, save_replicas=False):
+    lti_to_files = defaultdict(set)
+    for storage_metadata in global_storage_metadata:
+        for lti, fname in storage_metadata.items():
+            lti_to_files[lti].add(fname)
+
+    file_load = defaultdict(int)
+    out = {}
+    for lti, file_candidates in lti_to_files.items():
+        candidates = sorted(file_candidates)
+        selected_main_file = min(candidates, key=lambda f: file_load[f])
+        file_load[selected_main_file] += 1
+
+        if save_replicas:
+            lti_main = replace(lti, replica_id=0)
+            out[lti_main] = selected_main_file
+            replica_id = 1
+            for fname in candidates:
+                if fname == selected_main_file:
+                    continue
+                lti_replica = replace(lti, replica_id=replica_id)
+                out[lti_replica] = fname
+                replica_id += 1
+        else:
+            out[lti] = selected_main_file
+
+    return out
+
+
+def dedup_tensor(
+    local_state_dict, local_storage_metadata, global_storage_metadata
+):
+    """
+    Dedup the replicated tensor in local state_dict.
+
+    Args:
+        local_state_dict(Dict[str, paddle.Tensor]): The state_dict of current rank.
+        local_storage_metadata(Dict[LocalTensorIndex, str]): The storage metadata of current rank.
+        global_storage_metadata(Dict[LocalTensorIndex, str]): The final storage metadata of all ranks.
+
+    Examples:
+        In rank0, local_state_dict:{"w1": t1_0, "w2": t2}, local_storage_metadata:{LocalTensorIndex("w1", (0,0)): "0_0.distcp", LocalTensorIndex("w2", (0,0)): "0_0.distcp"},
+        in rank1, local_state_dict:{"w1": t1_1, "w2": t2}, local_storage_metadata:{LocalTensorIndex("w1", (1,0)): "1_0.distcp", LocalTensorIndex("w2", (0,0)): "1_0.distcp"},
+        global_storage_metadata:{LocalTensorIndex("w1", (0,0)): "0_0.distcp", LocalTensorIndex("w1", (1,0)): "1_0.distcp", LocalTensorIndex("w2", (0, 0)): "0_0.distcp"}.
+        w2 is replicated in rank0 and rank1. We save it in rank0 as default thus need to remove it in other ranks.
+        Finally, the local_state_dict:{"w1": t1_1, "w2": t2} in rank1 update to {"w1": t1_1}.
+    """
+
+    for tensor_index, file_name in global_storage_metadata.items():
+        rank = int(file_name.split(".")[0].split("_")[0])
+        if (
+            tensor_index in local_storage_metadata
+            and rank != paddle.distributed.get_rank()
+        ):
+            local_state_dict.pop(tensor_index.tensor_key)
+
+
+def save_state_dict(
+    state_dict: dict[str, Tensor] | dict[str, ShardedWeight],
+    path: str,
+    process_group: Group | None = None,
+    coordinator_rank: int = 0,
+    unique_id: int | None = None,
+    async_save: bool = False,
+    safetensors: bool = False,
+    save_replicas: bool = False,
+) -> None:
+    r"""
+    Save the state_dict of model to path.
+
+    Args:
+        state_dict(Dict[str, paddle.Tensor]): The state_dict to save.
+        path(str): The directory to save state_dict.
+        process_group(paddle.distributed.collective.Group): ProcessGroup to be used for cross-rank synchronization. Use the default process group which contains all cards.
+        coordinator_rank(int): The rank used to save non distributed values. Rank 0 is used by default.
+        unique_id(int): The unique id of checkpoint, used to distinguish between different checkpoint versions. Default is None, in which case the id 0 when save for the first time and increased by 1 each time when calling save_state_dict in the same path. If unique_id is given and there is already checkpoint with the same unique_id, it will be overrited.
+        async_save(bool): Async save the state_dict, default is False.
+        safetensors(bool): Whether to save using safetensors format. Default is False.
+        save_replicas (bool): Whether to save all tensor replicas (e.g., from different ranks) instead of only one deduplicated copy per tensor. Default is False.
+    Examples:
+        .. code-block:: pycon
+
+            >>> # doctest: +SKIP('run in distributed mode')
+            >>> import paddle
+            >>> import paddle.distributed as dist
+            >>> w1 = paddle.arange(32).reshape([4, 8])
+            >>> mesh = dist.ProcessMesh([0, 1])
+            >>> sharded_w1 = dist.shard_tensor(w1, mesh, [dist.Shard(0), dist.Replicate()])
+            >>> state_dict = {"w1": sharded_w1}
+            >>> dist.save_state_dict(state_dict, "./checkpoint")
+            >>> # doctest: -SKIP
+    """
+    with paddle.base.dygraph.guard():
+        assert isinstance(state_dict, dict), (
+            f"The state_dict should be a dictionary.But now the type is {type(state_dict)}."
+        )
+        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)), (
+                    f"The value of state_dict should be a paddle.Tensor or ShardedWeight, but got: {val}."
+                )
+
+        if not os.path.exists(path):
+            os.makedirs(path, exist_ok=True)
+
+        use_dist = True if paddle.distributed.get_world_size() > 1 else False
+
+        if use_dist and process_group is None and not is_initialized():
+            # Init the default global process group
+            paddle.distributed.init_parallel_env()
+
+        if unique_id is None:
+            max_unique_id = get_max_id(path)
+            logger.debug(f"Max unique id: {max_unique_id}")
+            if max_unique_id is None:
+                unique_id = 0
+            else:
+                unique_id = max_unique_id
+        else:
+            assert unique_id >= 0, f'{unique_id} should be >= 0'
+        if use_dist:
+            check_unique_id(unique_id, process_group)
+        file_suffix = "distcp" if not safetensors else "safetensors"
+        file_name = f"{paddle.distributed.get_rank()}_{unique_id}.{file_suffix}"
+        logger.debug(f"The checkpoint is saved to file_name:{file_name}")
+
+        metadata = Metadata()
+        local_state_dict = {}
+        local_state_dict_metadata = {}
+        local_storage_metadata = {}
+        global_shape = None
+        for key, val in flat_state_dict.items():
+            local_tensor, local_tensor_metadata = extract_tensor_metadata(val)
+            if local_tensor is None and local_tensor_metadata is None:
+                continue
+
+            local_state_dict[key] = local_tensor
+            local_state_dict_metadata[key] = local_tensor_metadata
+            global_offset = local_tensor_metadata.global_offset
+            is_flattened = local_tensor_metadata.is_flattened
+            flattened_range = local_tensor_metadata.flattened_range
+            local_shape = local_tensor_metadata.local_shape
+
+            local_storage_metadata[
+                LocalTensorIndex(
+                    tensor_key=key,
+                    global_offset=global_offset,
+                    is_flattened=is_flattened,
+                    flattened_range=flattened_range,
+                    local_shape=local_shape,
+                )
+            ] = file_name
+
+        global_state_dict_metadata = []
+        global_storage_metadata = []
+        global_flatten_mapping = []
+        if use_dist:
+            paddle.distributed.all_gather_object(
+                global_state_dict_metadata,
+                local_state_dict_metadata,
+                process_group,
+            )
+            paddle.distributed.all_gather_object(
+                global_storage_metadata, local_storage_metadata, process_group
+            )
+            paddle.distributed.all_gather_object(
+                global_flatten_mapping, mapping, process_group
+            )
+        else:
+            global_state_dict_metadata.append(local_state_dict_metadata)
+            global_storage_metadata.append(local_storage_metadata)
+            global_flatten_mapping.append(mapping)
+
+        metadata.state_dict_metadata = merge_state_dict_metadata(
+            global_state_dict_metadata
+        )
+        metadata.storage_metadata = balanced_dedup_key_in_dict(
+            global_storage_metadata, save_replicas=save_replicas
+        )
+        metadata.flat_mapping = dedup_key_in_dict(global_flatten_mapping)
+
+        logger.debug(f"metadata:{metadata}")
+        write_to_file_if_empty(
+            metadata, os.path.join(path, f"{unique_id}.metadata")
+        )
+
+        if not save_replicas:
+            dedup_tensor(
+                local_state_dict,
+                local_storage_metadata,
+                metadata.storage_metadata,
+            )
+
+        if async_save:
+            cpu_state_dict = copy_dict_to_cpu(local_state_dict)
+            clear_async_save_task_queue()
+
+            attempt = 0
+            ctx = multiprocessing.get_context("spawn")
+
+            def start_process():
+                nonlocal attempt
+                try:
+                    p = ctx.Process(
+                        target=paddle.save,
+                        args=(cpu_state_dict, os.path.join(path, file_name)),
+                        kwargs={'safetensors': safetensors},
+                    )
+                    p.start()
+                    return p
+                except Exception as e:
+                    logger.error(
+                        f"Attempt {attempt + 1} failed with error: {e}"
+                    )
+                    attempt += 1
+                    time.sleep(1)
+                    return start_process()
+
+            p = start_process()
+            async_save_queue.append(p)
+        else:
+            paddle.save(
+                local_state_dict,
+                os.path.join(path, file_name),
+                safetensors=safetensors,
+            )

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

@@ -0,0 +1,271 @@
+# Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+from __future__ import annotations
+
+from collections import OrderedDict
+from copy import deepcopy
+from dataclasses import dataclass
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from paddle import Tensor
+    from paddle.distributed.communication.group import Group
+
+
+@dataclass(frozen=True)
+class ShardedWeightDesc:
+    key: str
+    local_shape: tuple[int, ...]
+    global_shape: tuple[int, ...]
+    global_offset: tuple[int, ...]
+    dtype: str | None = None
+
+
+class ShardedWeight:
+    """
+    Represents a local shard of a distributed tensor parameter.
+
+    Args:
+        key (str): The name of the parameter.
+        local_tensor (Tensor): The local shard of the parameter.
+        local_shape (Tuple[int, ...]): The shape of the local shard.
+        global_shape (Tuple[int, ...]): The global logical shape of the parameter.
+        global_offset (Tuple[int, ...]): The offset of the local shard in the global parameter.
+        is_flattened (bool, optional): Whether the parameter has been flattened (used in sharding_v2 scenarios). Default is False.
+        flattened_range (slice, optional): If the parameter is flattened, this indicates the index range of the actual local shard within the local_tensor.
+    """
+
+    def __init__(
+        self,
+        key: str,
+        local_tensor: Tensor,
+        local_shape: tuple[int, ...],
+        global_shape: tuple[int, ...],
+        global_offset: tuple[int, ...],
+        is_flattened: bool = False,
+        flattened_range: slice | None = None,
+    ) -> None:
+        self.key = key
+        if local_tensor.is_dist():
+            self.local_tensor = local_tensor._local_value()
+            # Note: The local_tensor must keep the same name with the original tensor. Otherwise, the static_to_struct_mapping will be wrong.
+            self.local_tensor.name = local_tensor.name
+            self.local_shape = local_tensor._local_shape
+        else:
+            self.local_tensor = local_tensor
+            self.local_shape = tuple(local_shape)
+        self.global_shape = global_shape
+        self.global_offset = global_offset
+        self.is_flattened = is_flattened
+        self.flattened_range = flattened_range
+
+    def __str__(self) -> str:
+        """Returns a formatted string representation of the sharded tensor."""
+        return (
+            f"ShardedWeight(\n"
+            f"  key={self.key},\n"
+            f"  local_tensor={type(self.local_tensor).__name__}(shape={self.local_tensor.shape}),\n"
+            f"  local_shape={self.local_shape},\n"
+            f"  global_shape={self.global_shape},\n"
+            f"  global_offset={self.global_offset},\n"
+            f"  flattened_range={self.flattened_range}\n"
+            f")"
+        )
+
+
+ShardedStateDict = dict[str, ShardedWeight] | OrderedDict[str, ShardedWeight]
+
+
+def shard_weight(
+    key: str,
+    weight: Tensor,
+    axis: int,
+    group: Group,
+) -> ShardedWeight:
+    """Creates a ShardedWeight by splitting the input tensor along a specified axis.
+
+    Args:
+        key: Unique identifier for the tensor.
+        weight: The input tensor to be sharded.
+        axis: The axis along which to shard the tensor.
+        group: The process group used for distributed communication.
+
+    Returns:
+        A ShardedWeight representing the local portion of the global tensor.
+    """
+    if axis < 0 or axis >= len(weight.shape):
+        raise ValueError(
+            f"Shard axis {axis} is invalid for tensor with shape {weight.shape}"
+        )
+
+    # Get hybrid communication group and rank information
+    current_rank = group.rank
+    world_size = group.nranks
+
+    # Calculate shapes and offsets
+    local_shape = weight.shape
+    global_shape = deepcopy(local_shape)
+    global_shape[axis] = local_shape[axis] * world_size
+    global_shape = tuple(global_shape)
+    local_shape = tuple(local_shape)
+    global_offset = [0] * len(global_shape)
+    if world_size > 1:
+        global_offset[axis] = current_rank * local_shape[axis]
+    global_offset = tuple(global_offset)
+
+    return ShardedWeight(
+        key=key,
+        local_tensor=weight,
+        local_shape=local_shape,
+        global_shape=global_shape,
+        global_offset=global_offset,
+    )
+
+
+def make_tp_sharded_weight_for_checkpoint(
+    key: str,
+    tensor: Tensor,
+    tensor_parallel_axis: int = 0,
+) -> ShardedWeight:
+    """Creates a tensor-parallel sharded tensor for checkpointing purposes.
+
+    Args:
+        key: Unique identifier for the tensor in the checkpoint.
+        tensor: The local tensor portion to be sharded.
+        tensor_parallel_axis: The axis along which tensor parallelism is applied.
+                            Defaults to 0 (first dimension).
+
+    Returns:
+        A ShardedWeight configured for tensor parallel checkpointing.
+    """
+    from paddle.distributed.fleet import get_hybrid_communicate_group
+
+    hcg = get_hybrid_communicate_group()
+    tensor_parallel_group = hcg.get_model_parallel_group()
+
+    return shard_weight(
+        key=key,
+        weight=tensor,
+        axis=tensor_parallel_axis,
+        group=tensor_parallel_group,
+    )
+
+
+def make_replicated_sharded_weight(
+    key: str,
+    tensor: Tensor,
+) -> ShardedWeight:
+    """
+    Creates a ShardedWeight that represents a fully replicated tensor (each process holds a full copy).
+
+    Args:
+        key: Unique identifier for the tensor in the checkpoint.
+        tensor: The local tensor (full copy).
+
+    Returns:
+        ShardedWeight: A ShardedWeight instance representing the replicated tensor.
+    """
+    zero_offset = tuple(0 for _ in tensor.shape)
+    return ShardedWeight(
+        key=key,
+        local_tensor=tensor,
+        local_shape=tuple(tensor.shape),
+        global_shape=tuple(tensor.shape),
+        global_offset=zero_offset,
+    )
+
+
+def build_sharded_state_dict(
+    state_dict: dict[str, Tensor],
+    shard_rules: dict[str, int] | None = None,
+    prefix: str = "",
+) -> dict[str, ShardedWeight]:
+    """Converts a regular state dict to a sharded state dict based on sharding rules.
+
+    Args:
+        state_dict: The original state dictionary containing tensors
+        shard_rules: Dictionary mapping tensor names to their sharding axes.
+                    If None, treated as empty dict (no tensor parallelism).
+        prefix: Optional prefix to prepend to all tensor keys
+
+    Returns:
+        Dictionary with the same keys as input but values converted to ShardedWeight
+        or regular Tensor based on sharding rules.
+
+    Note:
+        Tensors not in shard_rules will be wrapped as non-sharded ShardedWeights.
+    """
+    shard_rules = shard_rules or {}
+    sharded_state_dict = {}
+
+    for key, tensor in state_dict.items():
+        full_key = f"{prefix}{key}" if prefix else key
+
+        if key in shard_rules:
+            # Apply tensor parallelism sharding
+            sharded_state_dict[full_key] = (
+                make_tp_sharded_weight_for_checkpoint(
+                    key=full_key,
+                    tensor=tensor,
+                    tensor_parallel_axis=shard_rules[key],
+                )
+            )
+        else:
+            # Create regular sharded tensor (non-tensor-parallel)
+            sharded_state_dict[full_key] = make_replicated_sharded_weight(
+                key=full_key,
+                tensor=tensor,
+            )
+
+    return sharded_state_dict
+
+
+def create_sharded_weight_with_new_local(
+    new_key: str,
+    new_local_tensor: Tensor,
+    reference_tensor: ShardedWeight,
+) -> ShardedWeight:
+    """
+    Creates a new ShardedWeight with a new local tensor while preserving the metadata from a reference ShardedWeight.
+
+    Args:
+        new_key (str): The new key for the ShardedWeight.
+        new_local_tensor (Tensor): The new local tensor to use (must match reference_tensor.local_shape).
+        reference_tensor (ShardedWeight): The reference ShardedWeight to copy metadata from.
+
+    Returns:
+        ShardedWeight: A new ShardedWeight with the new local tensor and copied metadata.
+
+    """
+    # Copy metadata from the reference tensor
+    global_shape = deepcopy(reference_tensor.global_shape)
+    local_shape = deepcopy(reference_tensor.local_shape)
+    global_offset = deepcopy(reference_tensor.global_offset)
+
+    # Input validation: Check if new_local_tensor's shape matches local_shape
+    if tuple(new_local_tensor.shape) != tuple(local_shape):
+        raise ValueError(
+            f"Shape mismatch: new_local_tensor has shape {new_local_tensor.shape}, "
+            f"but expected shape {local_shape} (from reference_tensor.local_shape)."
+        )
+
+    return ShardedWeight(
+        key=new_key,
+        local_tensor=new_local_tensor,
+        local_shape=tuple(local_shape),
+        global_shape=tuple(global_shape),
+        global_offset=tuple(global_offset),
+    )

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

@@ -0,0 +1,755 @@
+# Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from __future__ import annotations
+
+import ast
+import copy
+import os
+import re
+from collections import defaultdict
+from dataclasses import replace
+from typing import TYPE_CHECKING
+
+import numpy as np
+from safetensors.numpy import safe_open
+
+import paddle
+from paddle.distributed.fleet.utils.log_util import logger
+
+from ..aoa.aoa_engine import (
+    postprocess_transpose,
+)
+from .metadata import (
+    LocalTensorIndex,
+    LocalTensorMetadata,
+    Metadata,
+)
+from .sharded_weight import (
+    ShardedWeight,
+    ShardedWeightDesc,
+)
+
+if TYPE_CHECKING:
+    from paddle.framework import core
+
+
+def get_coordinator(mesh: np.array | list[list[int]], rank: int):
+    mesh = paddle.to_tensor(mesh)
+    rand_coordinator = (mesh == rank).nonzero()
+    assert rand_coordinator.shape[0] in (
+        0,
+        1,
+    ), f"rand_coordinator.shape: {rand_coordinator.shape}"
+    return (
+        rand_coordinator[0].tolist() if rand_coordinator.shape[0] > 0 else None
+    )
+
+
+# NOTE(zhangbo): Refer to the BalancedSplit function in the reshard_utils.cc file.
+def balanced_split(total_nums, num_of_pieces):
+    has_remainder = total_nums % num_of_pieces != 0
+    result = [(total_nums + num_of_pieces - 1) // num_of_pieces] * num_of_pieces
+    if has_remainder:
+        last_value = result[-1]
+        result[-1] = last_value - (last_value * num_of_pieces - total_nums)
+    return result
+
+
+def compute_local_shape_and_global_offset(
+    global_shape: list[int],
+    process_mesh: core.ProcessMesh,
+    placements: list[core.Placement],
+) -> tuple[tuple[int], tuple[int]]:
+    from paddle.distributed.auto_parallel.placement_type import (
+        placemetns_to_dist_status,
+    )
+
+    mesh = np.array(process_mesh.process_ids).reshape(process_mesh.shape)
+    # deal with cross mesh case
+    if paddle.distributed.get_rank() not in mesh:
+        return (None, None)
+    rank_coordinator = get_coordinator(mesh, paddle.distributed.get_rank())
+    local_shape = copy.copy(global_shape)
+    global_offset = [0 for _ in global_shape]
+
+    dims_mapping, _ = placemetns_to_dist_status(placements, len(global_shape))
+    for tensor_dim, mesh_dims in enumerate(dims_mapping):
+        if len(mesh_dims) == 0:
+            continue
+        local_offset = [0] * len(global_shape)
+        for mesh_dim in mesh_dims:
+            chunk_idx = rank_coordinator[mesh_dim]
+            chunks = balanced_split(
+                local_shape[tensor_dim], process_mesh.shape[mesh_dim]
+            )
+            local_shape[tensor_dim] = chunks[chunk_idx]
+            local_offset[tensor_dim] = sum(chunks[:chunk_idx])
+
+            if global_offset[tensor_dim] <= local_offset[tensor_dim]:
+                global_offset[tensor_dim] = local_offset[tensor_dim]
+            else:
+                global_offset[tensor_dim] += local_offset[tensor_dim]
+
+    return tuple(local_shape), tuple(global_offset)
+
+
+def flatten_state_dict(state_dict):
+    """
+    Flatten the nested dict to a flat dict.
+    {"model": {"w0": xxx}} -> {model.w0: xxx}
+    """
+    flatten_state_dict = {}
+    mapping = {}
+
+    def _flatten(key, value):
+        nonlocal _flatten
+        if isinstance(value, dict):
+            for k, v in value.items():
+                assert isinstance(k, str), f"The key should be str, but is {k}"
+                _flatten((*key, k), v)
+        elif isinstance(value, (paddle.Tensor, ShardedWeight)):
+            flatten_key_str = ".".join(key)
+            flatten_state_dict[flatten_key_str] = value
+            mapping[flatten_key_str] = key
+        else:
+            raise ValueError(
+                f"The value should be dict or paddle.Tensor, but is {value}"
+            )
+
+    _flatten((), state_dict)
+    del _flatten  # force python gc of recursive closure
+
+    return flatten_state_dict, mapping
+
+
+def unflatten_state_dict(flat_state_dict, mapping):
+    """
+    Unflatten the flat dict to a nested dict.
+    {model.w0: xxx} -> {"model": {"w0": xxx}}
+    """
+    state_dict = {}
+    for key, value in flat_state_dict.items():
+        key_tuple = mapping[key]
+        assert isinstance(key_tuple, tuple), (
+            f"The key should be tuple, but is {key_tuple}"
+        )
+        tmp = state_dict
+        for i in range(len(key_tuple) - 1):
+            key = key_tuple[i]
+            tmp = tmp.setdefault(key, {})
+        tmp[key_tuple[-1]] = value
+
+    return state_dict
+
+
+def get_max_id(path):
+    numbers = [0]
+    pattern = re.compile(r"^(\d+)_(\d+)\.distcp$")
+    files = os.listdir(path)
+    for file in files:
+        match = pattern.match(file)
+        if match:
+            numbers.append(int(match.group(2)))
+    return max(numbers) if numbers else None
+
+
+def check_unique_id(unique_id, process_group):
+    all_unique_id = []
+    paddle.distributed.all_gather_object(
+        all_unique_id, unique_id, process_group
+    )
+    for id in all_unique_id[1:]:
+        assert id == all_unique_id[0], f"id:{id} !=  all_unique_id[0]"
+
+
+def ravel_index(indices, shape):
+    idx = 0
+    for i, dim in zip(indices, shape):
+        idx = idx * dim + i
+    return idx
+
+
+def unravel_index(idx, shape):
+    indices = []
+    for dim in reversed(shape):
+        indices.append(idx % dim)
+        idx //= dim
+    return tuple(reversed(indices))
+
+
+def minimal_nd_slice(shape, flat_start, flat_end):
+    start_idx = unravel_index(flat_start, shape)
+    end_idx = unravel_index(flat_end - 1, shape)
+    min_slices = []
+    for axis in range(len(shape)):
+        if axis == 0:
+            s = start_idx[axis]
+            e = end_idx[axis] + 1
+        else:
+            if start_idx[axis - 1] == end_idx[axis - 1]:
+                s = min(start_idx[axis], end_idx[axis])
+                e = max(start_idx[axis], end_idx[axis]) + 1
+            else:
+                s = 0
+                e = shape[axis]
+        min_slices.append((s, e))
+    return min_slices, start_idx, end_idx
+
+
+def flat_range_in_min_slice(shape, min_slices, flat_start, flat_end):
+    min_starts = tuple(s[0] for s in min_slices)
+    min_flat_start = ravel_index(min_starts, shape)
+    return flat_start - min_flat_start, flat_end - min_flat_start
+
+
+def is_sharded_state_dict(state_dict, use_dist=True, process_group=None):
+    values = list(state_dict.values())
+    is_all_sharded = all(isinstance(v, ShardedWeight) for v in values)
+    has_sharded = any(isinstance(v, ShardedWeight) for v in values)
+
+    if has_sharded and not is_all_sharded:
+        raise TypeError(
+            "All values must be ShardedWeight if any value is ShardedWeight."
+        )
+
+    if not use_dist:
+        return is_all_sharded
+
+    if is_all_sharded:
+        flag = 1
+    elif len(values) == 0:
+        flag = 0
+    else:
+        flag = -1
+
+    all_flags = []
+    paddle.distributed.all_gather_object(all_flags, flag, process_group)
+
+    assert all(f >= 0 for f in all_flags) or all(f <= 0 for f in all_flags), (
+        "Not support mixed type of ShardedWeight and non-ShardedWeight in the same state_dict!"
+    )
+    return all(f >= 0 for f in all_flags)
+
+
+def get_overlap_region(desc_offset, desc_shape, shard_offset, shard_shape):
+    ndim = len(desc_offset)
+    overlap_offset = []
+    overlap_shape = []
+    desc_starts = []
+    shard_starts = []
+    for i in range(ndim):
+        desc_lo = desc_offset[i]
+        desc_hi = desc_offset[i] + desc_shape[i]
+        shard_lo = shard_offset[i]
+        shard_hi = shard_offset[i] + shard_shape[i]
+        # overlap
+        lo = max(desc_lo, shard_lo)
+        hi = min(desc_hi, shard_hi)
+        if lo >= hi:
+            return False, None, None, None, None
+        overlap_offset.append(lo)
+        overlap_shape.append(hi - lo)
+        desc_starts.append(lo - desc_lo)
+        shard_starts.append(lo - shard_lo)
+    return True, overlap_offset, overlap_shape, desc_starts, shard_starts
+
+
+def assign_sharded_slice(
+    src_desc, src_shard, dst_desc, dst_shard, postprocess_list=None
+):
+    src_has, _, overlap_shape, src_desc_starts, src_shard_starts = (
+        get_overlap_region(
+            src_desc.global_offset,
+            src_desc.local_shape,
+            src_shard.global_offset,
+            src_shard.local_shape,
+        )
+    )
+
+    dst_has, _, overlap_shape2, dst_desc_starts, dst_shard_starts = (
+        get_overlap_region(
+            dst_desc.global_offset,
+            dst_desc.local_shape,
+            dst_shard.global_offset,
+            dst_shard.local_shape,
+        )
+    )
+
+    assert src_has or dst_has, "no overlap!"
+    if overlap_shape != overlap_shape2:
+        assert postprocess_list is not None, (
+            "only post transpose operation could make overlap shape mismatch"
+        )
+        transposed_src_overlap_shape = postprocess_transpose(
+            overlap_shape, postprocess_list
+        )
+
+        assert transposed_src_overlap_shape == overlap_shape2, (
+            f"overlap shape mismatch: {transposed_src_overlap_shape} vs {overlap_shape2}"
+        )
+        axes = list(range(len(transposed_src_overlap_shape)))
+
+        src_tensor_slice = paddle.slice(
+            src_shard.local_tensor,
+            axes=axes,
+            starts=src_shard_starts,
+            ends=[s + o for s, o in zip(src_shard_starts, overlap_shape)],
+        )
+
+        dst_tensor_slice = paddle.slice(
+            dst_shard.local_tensor,
+            axes=axes,
+            starts=dst_shard_starts,
+            ends=[s + o for s, o in zip(dst_shard_starts, overlap_shape2)],
+        )
+
+    else:
+        axes = list(range(len(overlap_shape)))
+
+        src_tensor_slice = paddle.slice(
+            src_shard.local_tensor,
+            axes=axes,
+            starts=src_shard_starts,
+            ends=[s + o for s, o in zip(src_shard_starts, overlap_shape)],
+        )
+
+        dst_tensor_slice = paddle.slice(
+            dst_shard.local_tensor,
+            axes=axes,
+            starts=dst_shard_starts,
+            ends=[s + o for s, o in zip(dst_shard_starts, overlap_shape)],
+        )
+
+    if postprocess_list is not None:
+        for ps in postprocess_list:
+            is_list, result = is_list_string(ps)
+            if is_list:
+                src_tensor_slice = paddle.transpose(src_tensor_slice, result)
+            else:
+                if isinstance(ps, str):
+                    src_tensor_slice = paddle.cast(src_tensor_slice, ps)
+
+    paddle.assign(src_tensor_slice, dst_tensor_slice)
+
+
+def merge_shard_info_list(list_of_dicts):
+    merged = defaultdict(list)
+    for info in list_of_dicts:
+        for k, v in info.items():
+            merged[k].extend(v)
+    return dict(merged)
+
+
+def build_shard_desc(val):
+    return ShardedWeightDesc(
+        key=val.key,
+        local_shape=tuple(val.local_shape),
+        global_shape=tuple(val.global_shape),
+        global_offset=tuple(val.global_offset),
+        dtype=str(val.local_tensor.dtype).split(".")[-1],
+    )
+
+
+def is_list_string(s):
+    try:
+        result = ast.literal_eval(s)
+        return (True, result) if isinstance(result, list) else (False, None)
+    except:
+        return False, None
+
+
+def write_to_file_if_empty(data, path):
+    lock_path = f"{path}.lock"
+    try:
+        fd = os.open(lock_path, os.O_CREAT | os.O_EXCL | os.O_WRONLY)
+        os.close(fd)
+        try:
+            if os.path.exists(path) and os.path.getsize(path) > 0:
+                logger.info(
+                    f"Process {os.getpid()} found the metadata file already written."
+                )
+                return
+            paddle.save(data, path)
+            logger.info(
+                f"Process {os.getpid()} successfully wrote the metadata to the file."
+            )
+        finally:
+            if os.path.exists(lock_path):
+                os.remove(lock_path)
+    except FileExistsError:
+        logger.info(
+            f"Process {os.getpid()} could not acquire the lock; another process is writing or has written the metadata."
+        )
+
+
+def build_global_state_shard_info(sharded_state_dict, process_group):
+    state_shard_info = defaultdict(list)
+    for key, val in sharded_state_dict.items():
+        desc = build_shard_desc(val)
+        state_shard_info[key].append(desc)
+
+    gathered_info = []
+
+    use_dist = True if paddle.distributed.get_world_size() > 1 else False
+    if use_dist:
+        paddle.distributed.all_gather_object(
+            gathered_info, dict(state_shard_info), process_group
+        )
+    else:
+        gathered_info = [dict(state_shard_info)]
+
+    return merge_shard_info_list(gathered_info)
+
+
+def merge_state_dict_metadata(global_state_dict_metadata):
+    assert isinstance(global_state_dict_metadata, list), (
+        "The global_state_dict should be a list."
+    )
+    out = {}
+    for state_dict in global_state_dict_metadata:
+        for key, val in state_dict.items():
+            if key not in out:
+                out[key] = []
+
+            if isinstance(val, list):
+                for item in val:
+                    if item not in out[key]:
+                        out[key].append(item)
+            else:
+                if val not in out[key]:
+                    out[key].append(val)
+
+    return out
+
+
+def recover_shard_tensor_from_shards(sharded_weights: list, sw):
+    def _assign_slice(dst_tensor, dst_starts, dst_ends, src_tensor):
+        axes = list(range(len(dst_starts)))
+        view = paddle.slice(
+            dst_tensor, axes=axes, starts=dst_starts, ends=dst_ends
+        )
+        paddle.assign(src_tensor, output=view)
+        return dst_tensor
+
+    dims = len(sw.global_offset)
+    sw_glo_start = sw.global_offset
+    sw_glo_end = [sw.global_offset[i] + sw.local_shape[i] for i in range(dims)]
+    sw_shape = sw.local_shape
+
+    for s in sharded_weights:
+        s_glo_start = s.global_offset
+        s_glo_end = [s.global_offset[i] + s.local_shape[i] for i in range(dims)]
+
+        overlap = []
+        for i in range(dims):
+            ol_start = max(s_glo_start[i], sw_glo_start[i])
+            ol_end = min(s_glo_end[i], sw_glo_end[i])
+            if ol_start >= ol_end:
+                break
+            overlap.append((ol_start, ol_end))
+        else:
+            s_starts = [ol[0] - s_glo_start[i] for i, ol in enumerate(overlap)]
+            s_ends = [ol[1] - s_glo_start[i] for i, ol in enumerate(overlap)]
+            sw_starts = [
+                ol[0] - sw_glo_start[i] for i, ol in enumerate(overlap)
+            ]
+            sw_ends = [ol[1] - sw_glo_start[i] for i, ol in enumerate(overlap)]
+
+            axes = list(range(len(s_starts)))
+            src = paddle.slice(
+                s.local_tensor, axes=axes, starts=s_starts, ends=s_ends
+            )
+            _assign_slice(sw.local_tensor, sw_starts, sw_ends, src)
+
+    return sw
+
+
+def create_hf_ckpt_metadata(
+    ckpt_path: str,
+    process_group=None,
+):
+    dtype_mapping = {
+        'U16': 'bfloat16',
+        'U8': 'uint8',
+        'I8': 'int8',
+        'I16': 'int16',
+        'BOOL': 'bool',
+        'F16': 'float16',
+        'F32': 'float32',
+        'F64': 'float64',
+        'BF16': 'bfloat16',
+        'I64': 'int64',
+    }
+
+    use_dist = paddle.distributed.get_world_size() > 1
+    cur_rank = paddle.distributed.get_rank() if use_dist else 0
+
+    accessible_files = os.listdir(ckpt_path)
+    safetensors_files = [
+        file for file in accessible_files if file.endswith(".safetensors")
+    ]
+    if use_dist:
+        rank_visible_files = []
+        local_files = {cur_rank: safetensors_files}
+        paddle.distributed.all_gather_object(
+            rank_visible_files, local_files, process_group
+        )
+        rank_visible_files = {
+            rank: files for d in rank_visible_files for rank, files in d.items()
+        }
+    else:
+        rank_visible_files = {0: safetensors_files}
+
+    def assign_files(
+        rank_visible_files: dict[int, list[str]],
+    ) -> dict[int, list[str]]:
+        all_files = set()
+        for files in rank_visible_files.values():
+            all_files.update(files)
+        all_files = list(all_files)
+
+        file2ranks = defaultdict(list)
+        for rank, files in rank_visible_files.items():
+            for f in files:
+                file2ranks[f].append(rank)
+
+        result = defaultdict(list)
+
+        all_files.sort(key=lambda f: (len(file2ranks[f]), f))
+
+        rank_load = dict.fromkeys(rank_visible_files, 0)
+
+        for f in all_files:
+            candidates = file2ranks[f]
+            min_rank = min(candidates, key=lambda r: (rank_load[r], r))
+            result[min_rank].append(f)
+            rank_load[min_rank] += 1
+
+        return {rank: result.get(rank, []) for rank in rank_visible_files}
+
+    rank2file = assign_files(rank_visible_files)
+    need_handle_files = rank2file[cur_rank]
+
+    local_state_dict_metadata = defaultdict(set)
+    local_storage_metadata = {}
+    for file_name in need_handle_files:
+        file_path = os.path.join(ckpt_path, file_name)
+        with safe_open(file_path, framework="np") as f:
+            for key in f.keys():
+                t_s = f.get_slice(key)
+                shape = tuple(t_s.get_shape())
+                dtype = t_s.get_dtype()
+                assert dtype in dtype_mapping, f"{dtype} is not supported yet."
+                dtype = dtype_mapping[dtype]
+                ltm = LocalTensorMetadata(
+                    global_offset=(0,) * len(shape),
+                    local_shape=shape,
+                    dtype=dtype,
+                    global_shape=shape,
+                    is_flattened=False,
+                )
+                lti = LocalTensorIndex(
+                    tensor_key=key,
+                    global_offset=(0,) * len(shape),
+                    is_flattened=False,
+                    local_shape=shape,
+                )
+                local_state_dict_metadata[key].add(ltm)
+                local_storage_metadata[lti] = file_name
+
+    if use_dist:
+        global_state_dict_metadata = []
+        global_storage_metadata = []
+        paddle.distributed.all_gather_object(
+            global_state_dict_metadata,
+            dict(local_state_dict_metadata),
+            process_group,
+        )
+        paddle.distributed.all_gather_object(
+            global_storage_metadata, local_storage_metadata, process_group
+        )
+    else:
+        global_state_dict_metadata = [dict(local_state_dict_metadata)]
+        global_storage_metadata = [local_storage_metadata]
+
+    state_dict_metadata = defaultdict(set)
+    for md in global_state_dict_metadata:
+        for k, v in md.items():
+            state_dict_metadata[k].update(v)
+    state_dict_metadata = {k: list(v) for k, v in state_dict_metadata.items()}
+
+    storage_metadata = {}
+    for md in global_storage_metadata:
+        storage_metadata.update(md)
+
+    metadata = Metadata(
+        state_dict_metadata=state_dict_metadata,
+        storage_metadata=storage_metadata,
+    )
+
+    METADATA_FILE_NAME = "flex-ckpt.auto_generated.metadata"
+    write_to_file_if_empty(
+        metadata, os.path.join(ckpt_path, METADATA_FILE_NAME)
+    )
+
+    if use_dist:
+        paddle.distributed.barrier(process_group)
+
+
+def get_target_tensor(target_state_dict, read_item):
+    use_dist = paddle.distributed.get_world_size() > 1
+    if any(isinstance(k, tuple) for k in target_state_dict):
+        key = (read_item.tensor_name, read_item.dst_global_offset)
+    else:
+        key = read_item.tensor_name
+
+    tensor = target_state_dict[key]
+    return tensor._local_value() if use_dist and tensor.is_dist() else tensor
+
+
+def slice_tensor(tensor, slice_begin, slice_shape):
+    if not slice_shape:
+        assert not tensor.shape, (
+            "Only 0-dimensional tensor supports empty slice_shape."
+        )
+        return tensor
+
+    slice_end = [
+        start + length for start, length in zip(slice_begin, slice_shape)
+    ]
+    axes = list(range(tensor.ndim))
+    return paddle.slice(tensor, axes=axes, starts=slice_begin, ends=slice_end)
+
+
+def extract_tensor_metadata(val):
+    if isinstance(val, paddle.Tensor):
+        # Case1: not initialized means this tensor is placed in another mesh which do not contain this rank
+        if not val._is_initialized():
+            return None, None
+        if val.is_dist():
+            local_tensor = val._local_value()
+            # Note: The local_tensor must keep the same name with the original tensor. Otherwise, the StructuredToParameterName@@ mapping will be wrong.
+            local_tensor.name = val.name
+            # when val is scalar, the shape is []
+            (
+                local_shape,
+                global_offset,
+            ) = (
+                compute_local_shape_and_global_offset(
+                    val.shape,
+                    val.process_mesh,
+                    val.placements,
+                )
+                if len(val.shape) > 0
+                else ((), ())
+            )
+            global_shape = val.shape
+            if local_shape is None or global_offset is None:
+                return None, None
+        else:
+            local_shape = tuple(val.shape)
+            global_offset = (
+                tuple([0] * len(val.shape)) if len(val.shape) > 0 else ()
+            )
+            global_shape = local_shape
+            local_tensor = val
+        is_flattened = False
+        flattened_range = None
+    elif isinstance(val, ShardedWeight):
+        local_tensor = val.local_tensor
+        local_shape = val.local_shape
+        global_offset = val.global_offset
+        global_shape = val.global_shape
+        is_flattened = val.is_flattened
+        flattened_range = val.flattened_range
+    else:
+        raise ValueError(
+            f"The value of state_dict should be a paddle.Tensor, but got: {val}"
+        )
+
+    local_tensor_dtype = str(local_tensor.dtype).split('.')[1]
+    if flattened_range is not None:
+        flattened_range = (flattened_range.start, flattened_range.stop)
+    else:
+        flattened_range = None
+    local_tensor_metadata = LocalTensorMetadata(
+        tuple(global_offset),
+        tuple(local_shape),
+        local_tensor_dtype,
+        tuple(global_shape),
+        is_flattened,
+        flattened_range,
+    )
+    assert (local_tensor is None) == (local_tensor_metadata is None), (
+        "local_tensor and local_tensor_metadata must both be None or both not None!"
+    )
+    return local_tensor, local_tensor_metadata
+
+
+def check_resumable_locally(
+    path, state_dict, metadata_manager, use_dist, process_group
+):
+    local_load = True
+    rank = paddle.distributed.get_rank() if use_dist else 0
+    checkpoint_file = f"{rank}_0.distcp"
+    file_path = os.path.join(path, checkpoint_file)
+
+    if not os.path.isfile(file_path):
+        local_load = False
+
+    state_dict_metadata = {}
+    for key, value in state_dict.items():
+        _, local_tensor_metadata = extract_tensor_metadata(value)
+        if local_tensor_metadata is not None:
+            state_dict_metadata[key] = local_tensor_metadata
+
+    if local_load:
+        file_storage_info = metadata_manager.get_file_storage_info()
+        cur_file_storage = {
+            replace(index, replica_id=None)
+            for index in file_storage_info.get(checkpoint_file, [])
+        }
+
+        for key, local_tensor_metadata in state_dict_metadata.items():
+            local_tensor_index = LocalTensorIndex(
+                tensor_key=key,
+                global_offset=local_tensor_metadata.global_offset,
+                is_flattened=local_tensor_metadata.is_flattened,
+                flattened_range=local_tensor_metadata.flattened_range,
+                local_shape=local_tensor_metadata.local_shape,
+                replica_id=None,
+            )
+            if local_tensor_index not in cur_file_storage:
+                local_load = False
+                break
+
+    if use_dist:
+        global_local_loads = []
+        paddle.distributed.all_gather_object(
+            global_local_loads, local_load, process_group
+        )
+        return all(global_local_loads)
+    else:
+        return local_load
+
+
+def need_transpose(postprocess_list):
+    if postprocess_list is None:
+        return False
+
+    for pp in postprocess_list:
+        if "[" in pp:
+            return True
+    else:
+        return False