chore: import upstream snapshot with attribution

python/paddle/distributed/flex_checkpoint/__init__.py

@@ -0,0 +1,13 @@
+# 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.

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

@@ -0,0 +1,13 @@
+# 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.

python/paddle/distributed/flex_checkpoint/aoa/aoa_engine.py

@@ -0,0 +1,944 @@
+# 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 ast
+import logging
+import re
+from dataclasses import dataclass
+
+import numpy as np
+
+logger = logging.getLogger(__name__)
+
+from ..dcp.sharded_weight import ShardedWeightDesc
+from .lexer import Lexer
+from .parser import Parser
+from .traceback import AOATraceback
+
+_ShardInfo = dict[str, list[ShardedWeightDesc]]
+
+# SliceRef := (key, src_slice, dst_slice, postprocess_list)
+SliceRef = tuple[str, tuple[slice, ...], tuple[slice, ...], list[str] | None]
+SUPPORTED_DTYPES = ['float16', 'float32', 'bfloat16']
+
+
+class TensorDesc:
+    def __init__(
+        self,
+        slices: list[SliceRef],
+        shape: tuple[int],
+        in_degree: int = 0,
+        out_degree: int = 0,
+        dtype: str | None = None,
+    ):
+        self.slices = slices
+        self.shape = shape
+        self.in_degree = in_degree
+        self.out_degree = out_degree
+        self.dtype = dtype
+
+    def __repr__(self):
+        s = []
+        for key, sl_src, sl_dst, pp_list in self.slices:
+            s.append(
+                f"{key}{sl_src} -> self{sl_dst}, postprocess_list={pp_list}"
+            )
+        return f"Tensor(shape={self.shape}, slices={s}, in_degree={self.in_degree}, out_degree={self.out_degree}, dtype={self.dtype})"
+
+
+@dataclass(frozen=True)
+class ShardMappingEntry:
+    target_slice: ShardedWeightDesc
+    source_slice: ShardedWeightDesc
+    postprocess_list: list[str] | None = None
+
+
+ShardMapping = list[ShardMappingEntry]
+
+OPTIMIZER_STATE_NAME = [
+    ".w_0",
+    ".moment1_0",
+    ".moment2_0",
+    ".beta1_pow_acc_0",
+    ".beta2_pow_acc_0",
+]
+
+
+def split_optimizer_state_key(key: str) -> tuple[str, str]:
+    for opt_state_name in OPTIMIZER_STATE_NAME:
+        if key.endswith(opt_state_name):
+            return key[: -len(opt_state_name)], opt_state_name
+    return key, None
+
+
+class AOAShardInfoContext:
+    def __init__(
+        self,
+        source_state_shard_info: _ShardInfo,
+        destination_state_shard_info: _ShardInfo,
+        aoa_config_reverse: bool = False,
+    ) -> None:
+        self.source_state_shard_info = source_state_shard_info
+        self.destination_state_shard_info = destination_state_shard_info
+        self.aoa_config_reverse = aoa_config_reverse
+        self.left_var_to_right_var_mapping = {}
+        self.right_var_from_left_var_mapping = {}
+        self.src_state_keys = set()
+        self.dst_state_keys = set()
+        self.init_src_state_keys()
+        self.init_dst_state_keys()
+
+    def init_src_state_keys(self):
+        for k in self.source_state_shard_info.keys():
+            model_state_key, _ = split_optimizer_state_key(k)
+            self.src_state_keys.add(model_state_key)
+
+    def init_dst_state_keys(self):
+        if self.destination_state_shard_info is None:
+            return
+        for k in self.destination_state_shard_info.keys():
+            model_state_key, _ = split_optimizer_state_key(k)
+            self.dst_state_keys.add(model_state_key)
+
+    def get_all_dst_state_keys(self):
+        return self.dst_state_keys
+
+    def get_all_src_state_keys(self):
+        return self.src_state_keys
+
+    def get_num_hidden_layers(
+        self,
+        name_with_layer_id: str,
+        layer_id_macro_tag: str,
+    ) -> int:
+        if layer_id_macro_tag not in name_with_layer_id:
+            raise ValueError(
+                f"layer_id_macro_tag '{layer_id_macro_tag}' not in name_with_layer_id '{name_with_layer_id}'"
+            )
+        prefix, suffix = name_with_layer_id.split(layer_id_macro_tag, 1)
+        pattern = re.compile(rf"{re.escape(prefix)}(\d+){re.escape(suffix)}")
+        match_layer_id = set()
+        for key in self.get_all_src_state_keys():
+            match = pattern.fullmatch(key)
+            if match:
+                layer_num = int(match.group(1))
+                match_layer_id.add(layer_num)
+        return match_layer_id
+
+    def get_src_state_shard_num(self, src_state_key: str) -> int:
+        model_state_key, opt_state_name = split_optimizer_state_key(
+            src_state_key
+        )
+
+        assert opt_state_name is None, (
+            "AOA notions apply only to the model state, but are automatically propagated to the optimizer state.Now the src_state_key is {src_state_key}, which is a optimizer state key."
+        )
+        reverse = True
+        if self.aoa_config_reverse:
+            reverse = False
+        # Only need to parse the model state key for optimizer state shard num, because the optimizer state slice info is completely consistent with the model state slice info.
+        resolved_model_state_key = self.resolve_mapping_chain(
+            model_state_key, reverse=reverse
+        )
+
+        state_keys = [
+            resolved_model_state_key,
+            f"{resolved_model_state_key}.w_0",
+            f"{resolved_model_state_key}.moment1_0",
+            f"{resolved_model_state_key}.moment2_0",
+        ]
+
+        shard_nums = {
+            len(
+                {
+                    shard_info.global_offset
+                    for shard_info in self.source_state_shard_info[key]
+                }
+            )
+            for key in state_keys
+            if key in self.source_state_shard_info
+        }
+
+        if not shard_nums:
+            logger.warning(
+                f"No shard information found for any of the keys: {state_keys}, return 1."
+            )
+            return 1
+        if len(shard_nums) > 1:
+            raise AssertionError(
+                f"Inconsistent shard numbers among keys in source_sharded_state_dict for the key {src_state_key}: shard_nums={shard_nums}."
+            )
+        return shard_nums.pop()
+
+    def get_dst_state_shard_num(self, dst_state_key: str) -> int:
+        if self.destination_state_shard_info is None:
+            # Default `dst_state_shard_num=1` if `destination_state_shard_info` is missing.
+            return 1
+        model_state_key, opt_state_name = split_optimizer_state_key(
+            dst_state_key
+        )
+
+        assert opt_state_name is None, (
+            "AOA notions apply only to the model state, but are automatically propagated to the optimizer state.Now the dst_state_key is {dst_state_key}, which is a optimizer state key."
+        )
+        reverse = False
+        if self.aoa_config_reverse:
+            reverse = True
+        # Only need to parse the model state key for optimizer state shard num, because the optimizer state slice info is completely consistent with the model state slice info.
+        resolved_model_state_key = self.resolve_mapping_chain(
+            model_state_key, reverse=reverse
+        )
+
+        state_keys = [
+            resolved_model_state_key,
+            f"{resolved_model_state_key}.w_0",
+            f"{resolved_model_state_key}.moment1_0",
+            f"{resolved_model_state_key}.moment2_0",
+        ]
+
+        shard_nums = {
+            len(
+                {
+                    shard_info.global_offset
+                    for shard_info in self.destination_state_shard_info[key]
+                }
+            )
+            for key in state_keys
+            if key in self.destination_state_shard_info
+        }
+
+        if not shard_nums:
+            logger.warning(
+                f"No shard information found for any of the keys: {state_keys}, return 1."
+            )
+            return 1
+        if len(shard_nums) > 1:
+            raise AssertionError(
+                f"Inconsistent shard numbers among keys in destination_state_shard_info for the key {dst_state_key}: shard_nums={shard_nums}."
+            )
+        return shard_nums.pop()
+
+    def resolve_mapping_chain(self, key: str, reverse: bool = False) -> str:
+        """
+        Recursively resolve the mapping chain, find the final leaf node
+
+        Args:
+            key: The key to be resolved
+            reverse: False use left_var_to_right_var_mapping，True use right_var_from_left_var_mapping
+
+        For example:
+        - reverse=False: temp_var -> dst_key
+        - reverse=True: temp_var -> src_key
+        """
+        visited = set()  # avoid infinite loop
+        current_key = key
+
+        if reverse:
+            mapping_dict = self.right_var_from_left_var_mapping
+        else:
+            mapping_dict = self.left_var_to_right_var_mapping
+
+        while current_key in mapping_dict:
+            assert current_key not in visited, (
+                f"Infinite loop detected in resolve_mapping_chain, which means the start key is not src_key or the end key is not dst_key, the aoa_config is error. current_key={current_key}, the loop is: {'->'.join(visited)}->{current_key}"
+            )
+            visited.add(current_key)
+            if reverse and current_key in self.get_all_src_state_keys():
+                break
+            elif not reverse and current_key in self.get_all_dst_state_keys():
+                break
+
+            mapped_vars = mapping_dict[current_key]
+            if mapped_vars and len(mapped_vars) > 0:
+                assert len(mapped_vars) == 1, (
+                    f"Reference chain resolution failed: "
+                    f"Unable to determine which leaf node the intermediate node '{key}' is directly associated with, "
+                    f"because a many-to-one mapping was found in the mapping relationship. "
+                    f"The many-to-one mapping is {current_key} : {mapped_vars}."
+                )
+                current_key = mapped_vars[0]
+            else:
+                break
+
+        return current_key
+
+
+class AOAEngine:
+    def __init__(
+        self,
+        aoa_config: dict[str, list[str]],
+        source_state_shard_info: _ShardInfo,
+        destination_state_shard_info: _ShardInfo,
+    ):
+        self.aoa_config = aoa_config
+        self.source_state_shard_info = source_state_shard_info
+        self.destination_state_shard_info = destination_state_shard_info
+        self.aoa_config_reverse = self.aoa_config.get(
+            "aoa_config_reverse", False
+        )
+        enable_traceback = self.aoa_config.get("enable_traceback", True)
+        self.traceback = AOATraceback() if enable_traceback else None
+        self.context = AOAShardInfoContext(
+            source_state_shard_info,
+            destination_state_shard_info,
+            self.aoa_config_reverse,
+        )
+        self.lexer = Lexer(self.context, traceback=self.traceback)
+        tokens = self.lexer.all_tokens(
+            self.aoa_config.get("aoa_statements", [])
+        )
+        self.parser = Parser(tokens)
+        self.statements = self.parser.parse_program()
+
+        if self.traceback and getattr(self.lexer, "final_expressions", None):
+            final_exprs = self.lexer.final_expressions
+            if len(final_exprs) == len(self.statements):
+                for expr, stmt in zip(final_exprs, self.statements):
+                    self.traceback.record_children(
+                        expr, [repr(stmt)], macro_name="parser"
+                    )
+
+        if self.aoa_config_reverse:
+            self.statements = list(reversed(self.statements))
+        self.input_vars = self.build_input_vars()
+        self.output_vars = {}
+        self.intermediate_vars = {}
+        self.need_remove_input_vars = set()
+        self.need_add_output_vars = set()
+
+        self.shape_propagation()
+
+    def make_input_tensor(
+        self, key: str, shape: tuple[int], dtype: str
+    ) -> TensorDesc:
+        base_slice = tuple([slice(0, s) for s in shape])
+        return TensorDesc(
+            [(key, base_slice, base_slice, None)],
+            shape,
+            in_degree=0,
+            out_degree=0,
+            dtype=dtype,
+        )
+
+    def build_input_vars(self):
+        input_vars = {}
+        dtype = None
+        for key, shards in sorted(self.source_state_shard_info.items()):
+            global_shape = shards[0].global_shape
+            model_state_key, opt_state_name = split_optimizer_state_key(key)
+            if opt_state_name is None:
+                dtype = shards[0].dtype
+            if model_state_key in input_vars.keys() or opt_state_name in [
+                ".beta1_pow_acc_0",
+                ".beta2_pow_acc_0",
+            ]:
+                continue
+            input_vars[model_state_key] = self.make_input_tensor(
+                model_state_key, global_shape, dtype
+            )
+        return input_vars
+
+    def split(
+        self, tensor: TensorDesc, axis: int, sizes: list[int]
+    ) -> list[TensorDesc]:
+        results = []
+        start = 0
+        tensor.out_degree += len(sizes)
+        dtype = tensor.dtype
+        for sz in sizes:
+            sub_dst_slice = [slice(None)] * len(tensor.shape)
+            sub_dst_slice[axis] = slice(0, sz)
+            sub_slices = []
+            for aidx, src_sl, dst_sl, pp_list in tensor.slices:
+                if pp_list is not None:
+                    src_sl = postprocess_transpose(list(src_sl), pp_list)
+
+                dst_start = (
+                    dst_sl[axis].start if dst_sl[axis].start is not None else 0
+                )
+                dst_stop = (
+                    dst_sl[axis].stop
+                    if dst_sl[axis].stop is not None
+                    else tensor.shape[axis]
+                )
+                inter_begin = max(start, dst_start)
+                inter_end = min(start + sz, dst_stop)
+                if inter_begin < inter_end:
+                    src_axis_start = (
+                        src_sl[axis].start
+                        if src_sl[axis].start is not None
+                        else 0
+                    )
+                    sub_src_sl = list(src_sl)
+                    sub_dst_sl = list(dst_sl)
+                    offset = inter_begin - dst_start
+                    length = inter_end - inter_begin
+                    sub_src_sl[axis] = slice(
+                        src_axis_start + offset,
+                        src_axis_start + offset + length,
+                    )
+                    sub_dst_sl[axis] = slice(
+                        inter_begin - start, inter_begin - start + length
+                    )
+                    if pp_list is not None:
+                        sub_src_sl = postprocess_transpose(
+                            list(sub_src_sl), pp_list, reverse=True
+                        )
+                        sub_slices.append(
+                            (
+                                aidx,
+                                tuple(sub_src_sl),
+                                tuple(sub_dst_sl),
+                                pp_list.copy(),
+                            )
+                        )
+                    else:
+                        sub_slices.append(
+                            (aidx, tuple(sub_src_sl), tuple(sub_dst_sl), None)
+                        )
+            new_shape = list(tensor.shape)
+            new_shape[axis] = sz
+            results.append(
+                TensorDesc(
+                    sub_slices,
+                    tuple(new_shape),
+                    in_degree=1,
+                    out_degree=0,
+                    dtype=dtype,
+                )
+            )
+            start += sz
+        return results
+
+    def concat(self, tensors: list[TensorDesc], axis: int) -> TensorDesc:
+        slices = []
+        assert len(tensors) >= 1, (
+            "When concatenating multiple tensors, there should be at least one!"
+        )
+        shape = list(tensors[0].shape)
+        ndim = len(shape)
+        assert 0 <= axis < ndim, (
+            f"when concat, the axis {axis} is out of range for tensors "
+            f"with shape {shape} (valid range: {0} to {ndim - 1})."
+        )
+        shape[axis] = sum(t.shape[axis] for t in tensors)
+        dtype = tensors[0].dtype
+        assert all(t.dtype == dtype for t in tensors), (
+            f"All tensors must have the same dtype when concatenating multiple tensors!But the tensors {tensors} have different dtypes: {[t.dtype for t in tensors]}."
+        )
+        curr = 0
+        for t in tensors:
+            t.out_degree += 1
+            for aidx, src_sl, dst_sl, pp_list in t.slices:
+                new_dst_sl = list(dst_sl)
+                dst_start = (
+                    dst_sl[axis].start if dst_sl[axis].start is not None else 0
+                )
+                dst_stop = (
+                    dst_sl[axis].stop
+                    if dst_sl[axis].stop is not None
+                    else t.shape[axis]
+                )
+                length = dst_stop - dst_start
+                new_dst_sl[axis] = slice(
+                    dst_start + curr, dst_start + curr + length
+                )
+                if pp_list is not None:
+                    slices.append(
+                        (aidx, src_sl, tuple(new_dst_sl), pp_list.copy())
+                    )
+                else:
+                    slices.append((aidx, src_sl, tuple(new_dst_sl), None))
+            curr += t.shape[axis]
+        return TensorDesc(
+            slices,
+            tuple(shape),
+            in_degree=len(tensors),
+            out_degree=0,
+            dtype=dtype,
+        )
+
+    def transpose(self, tensor: TensorDesc, permutation: str) -> TensorDesc:
+        slices = []
+        tensor.out_degree += 1
+        tensor_shape = transpose_list(
+            tensor.shape, ast.literal_eval(permutation)
+        )
+        dtype = tensor.dtype
+        for aidx, src_sl, dst_sl, pp_list in tensor.slices:
+            trans_dst_sl = transpose_list(dst_sl, ast.literal_eval(permutation))
+            if pp_list is not None:
+                new_pp_list = pp_list.copy()
+                new_pp_list.append(permutation)
+                slices.append((aidx, src_sl, trans_dst_sl, new_pp_list))
+            else:
+                slices.append((aidx, src_sl, trans_dst_sl, [permutation]))
+        return TensorDesc(
+            slices, tensor_shape, in_degree=1, out_degree=0, dtype=dtype
+        )
+
+    def cast(self, tensor: TensorDesc, dtype: str) -> TensorDesc:
+        slices = []
+        tensor.out_degree += 1
+        for aidx, src_sl, dst_sl, pp_list in tensor.slices:
+            if pp_list is not None:
+                new_pp_list = pp_list.copy()
+                new_pp_list.append(dtype)
+                slices.append((aidx, src_sl, dst_sl, new_pp_list))
+            else:
+                slices.append((aidx, src_sl, dst_sl, [dtype]))
+        # For the cast operation, post_process is required. Therefore, the returned
+        # Tensor's dtype here is the same as the input tensor's dtype, rather than the casted dtype.
+        return TensorDesc(
+            slices, tensor.shape, in_degree=1, out_degree=0, dtype=tensor.dtype
+        )
+
+    def identity(self, tensor: TensorDesc) -> TensorDesc:
+        tensor.out_degree += 1
+        return TensorDesc(
+            tensor.slices,
+            tensor.shape,
+            in_degree=1,
+            out_degree=0,
+            dtype=tensor.dtype,
+        )
+
+    def shape_propagation(self):
+        def _get_var_ref(var):
+            if var.name in self.intermediate_vars:
+                return self.intermediate_vars[var.name]
+            elif var.name in self.input_vars:
+                return self.input_vars[var.name]
+            else:
+                raise ValueError(f"{var.name} should be assigned before!")
+
+        for stmt in self.statements:
+            stmt_repr = repr(stmt)
+            left_vars = stmt.left_vars
+            right_vars = stmt.right_vars
+            if self.aoa_config_reverse:
+                left_vars, right_vars = right_vars, left_vars
+            attrs = stmt.attrs
+
+            try:
+                if len(left_vars) > 1 or len(right_vars) > 1:
+                    if not (len(attrs) == 1 and attrs[0].key == "axis"):
+                        raise ValueError(
+                            f"When split/concat, only support one attr named `axis`, but got {attrs}."
+                        )
+                    axis = attrs[0].value
+
+                    if len(left_vars) == 1:
+                        in_name = left_vars[0].name
+                        in_ref = _get_var_ref(left_vars[0])
+                        ndim = len(in_ref.shape)
+                        assert 0 <= axis < ndim, (
+                            f"when split, the axis {axis} is out of range for tensor {in_name} "
+                            f"with shape {in_ref.shape} (valid range: {0} to {ndim - 1})."
+                        )
+                        assert in_ref.shape[axis] % len(right_vars) == 0, (
+                            f"when split, the shape of the input tensor {in_name} is {in_ref.shape}, the axis is {axis}, the number of right_vars is {len(right_vars)}, but the shape of the input tensor {in_name} is not divisible by the number of right_vars."
+                        )
+                        sizes = [
+                            in_ref.shape[axis] // len(right_vars)
+                            for var in right_vars
+                        ]
+                        result = self.split(in_ref, axis, sizes)
+                        for out_var, out_ref in zip(right_vars, result):
+                            self.intermediate_vars[out_var.name] = out_ref
+                            if (
+                                out_var.name
+                                in self.context.get_all_dst_state_keys()
+                            ):
+                                self.output_vars[out_var.name] = out_ref
+
+                    elif len(right_vars) == 1:
+                        left_refs = [_get_var_ref(var) for var in left_vars]
+                        result = self.concat(left_refs, axis)
+                        out_name = right_vars[0].name
+                        self.intermediate_vars[out_name] = result
+                        if out_name in self.context.get_all_dst_state_keys():
+                            self.output_vars[out_name] = result
+
+                    else:
+                        raise SyntaxError(
+                            f'Unexpected split/concat statement: {stmt}'
+                        )
+
+                elif len(left_vars) == 1 and len(right_vars) == 1:
+                    lvar, rvar = left_vars[0], right_vars[0]
+                    if rvar.name == "_":
+                        self.need_remove_input_vars.add(lvar.name)
+                    elif lvar.name == "_":
+                        self.need_add_output_vars.add(rvar.name)
+                    else:
+                        if len(attrs) > 0:
+                            assert len(attrs) == 1 or (
+                                len(attrs) == 2
+                                and {attr.key for attr in attrs}
+                                == {"src_dtype", "dst_dtype"}
+                            ), (
+                                "Only support:\n"
+                                " - One operator, OR\n"
+                                " - Two operators with keys {'src_dtype', 'dst_dtype'}."
+                            )
+                            attr = attrs[0]
+                            in_ref = _get_var_ref(lvar)
+                            if attr.key == "permute":
+                                if attr.value == "[]":
+                                    ndim = len(in_ref.shape)
+                                    perm = str(list(range(ndim - 1, -1, -1)))
+                                else:
+                                    perm = attr.value
+                                    if self.aoa_config_reverse:
+                                        perm = str(
+                                            invert_permutation(
+                                                ast.literal_eval(perm)
+                                            )
+                                        )
+                                result = self.transpose(in_ref, perm)
+                            elif attr.key == "dtype":
+                                assert not self.aoa_config_reverse, (
+                                    "When `aoa_config_reverse=True`, the dtype must be specified as "
+                                    "'src_dtype=...,dst_dtype=...'. Formats like 'dtype=xxx' are not supported."
+                                )
+                                assert attr.value in SUPPORTED_DTYPES, (
+                                    f"Unsupported cast dtype: {attr.value}"
+                                )
+                                result = self.cast(in_ref, attr.value)
+                            elif (
+                                attrs[0].key == "src_dtype"
+                                and attrs[1].key == "dst_dtype"
+                            ):
+                                src_dtype, dst_dtype = (
+                                    attrs[0].value,
+                                    attrs[1].value,
+                                )
+                                assert src_dtype in SUPPORTED_DTYPES, (
+                                    f"Unsupported cast dtype: {src_dtype}"
+                                )
+                                assert dst_dtype in SUPPORTED_DTYPES, (
+                                    f"Unsupported cast dtype: {dst_dtype}"
+                                )
+                                if self.aoa_config_reverse:
+                                    src_dtype, dst_dtype = dst_dtype, src_dtype
+                                result = self.cast(in_ref, dst_dtype)
+                            elif attr.key == "axis":
+                                result = in_ref
+                            else:
+                                raise ValueError(
+                                    f"Unsupported attribute: {attr}"
+                                )
+
+                            self.intermediate_vars[rvar.name] = result
+                            if (
+                                rvar.name
+                                in self.context.get_all_dst_state_keys()
+                            ):
+                                self.output_vars[rvar.name] = result
+                        else:
+                            # rename operation
+                            in_ref = _get_var_ref(lvar)
+                            result = self.identity(in_ref)
+                            self.intermediate_vars[rvar.name] = result
+                            if (
+                                rvar.name
+                                in self.context.get_all_dst_state_keys()
+                            ):
+                                self.output_vars[rvar.name] = result
+                else:
+                    raise SyntaxError(f'Unexpected statement: {stmt}')
+            except (
+                AssertionError,
+                ValueError,
+                KeyError,
+                SyntaxError,
+                RuntimeError,
+            ) as e:
+                if self.traceback:
+                    chain = self.traceback.build_chain(stmt_repr)
+                    self.traceback.add_error(
+                        error_message=str(e),
+                        stage="shape_propagation",
+                        chain=chain,
+                        error_type=type(e).__name__,
+                    )
+                    self.traceback.print()
+                raise
+        if self.destination_state_shard_info is not None:
+            for name in self.destination_state_shard_info:
+                model_state_key, _ = split_optimizer_state_key(name)
+                if model_state_key not in self.output_vars:
+                    if model_state_key in self.need_add_output_vars:
+                        self.output_vars[model_state_key] = None
+                    else:
+                        assert model_state_key in self.input_vars, (
+                            f"{model_state_key} needs to be loaded, "
+                            f"but not found in checkpoint. "
+                            f"If the key exists in the current model but not in the loaded checkpoint, please use the add primitive in aoa_statements: "
+                            f"_ -> {model_state_key}, and {model_state_key} will be randomly initialized."
+                        )
+                        self.output_vars[model_state_key] = self.input_vars[
+                            model_state_key
+                        ]
+        else:
+            # When destination_state_shard_info is not provided, the AOAEngine automatically derives it
+            # from source_state_shard_info and aha_statements. In this case, all destination_states
+            # remain unsharded (not partitioned).
+            for name, ref_t in self.input_vars.items():
+                if (
+                    name not in self.output_vars
+                    and ref_t.out_degree == 0
+                    and name not in self.need_remove_input_vars
+                ):
+                    self.output_vars[name] = self.identity(ref_t)
+            for name, ref_t in self.intermediate_vars.items():
+                if name not in self.output_vars and ref_t.out_degree == 0:
+                    self.output_vars[name] = self.identity(ref_t)
+
+    def find_source_slices(
+        self, key: str, local_slice: tuple[slice, ...]
+    ) -> list[SliceRef]:
+        assert key in self.output_vars, (
+            f"The key {key} is not in the output_vars (which is built during load_state_dict)."
+        )
+        tensor = self.output_vars[key]
+        if tensor is None:
+            return []
+        results = []
+        assert len(local_slice) == len(tensor.shape), (
+            f"For the key {key}, the target_tensor has {len(local_slice)} dimensions, "
+            f"but the tensor in output_vars has {len(tensor.shape)} dimensions (shape={tensor.shape}). "
+        )
+        ndim = len(tensor.shape)
+
+        def slice_intersect(a: slice, b: slice):
+            start = max(a.start, b.start)
+            stop = min(a.stop, b.stop)
+            if start >= stop:
+                return None
+            return slice(start, stop, 1)
+
+        for src_key, sl_src, sl_dst, pp_list in tensor.slices:
+            intersection = []
+            for i in range(ndim):
+                inter = slice_intersect(local_slice[i], sl_dst[i])
+                if inter is None:
+                    break
+                intersection.append(inter)
+            else:
+                # Compute corresponding src_slice for the intersection
+                if pp_list is not None:
+                    sl_src = postprocess_transpose(list(sl_src), pp_list)
+                src_slice = []
+                for i in range(ndim):
+                    dst = sl_dst[i]
+                    src = sl_src[i]
+                    dst_start = dst.start
+                    src_start = src.start
+                    inter_start, inter_stop = (
+                        intersection[i].start,
+                        intersection[i].stop,
+                    )
+                    offset = inter_start - dst_start
+                    src_inter_start = src_start + offset
+                    src_inter_stop = src_inter_start + (
+                        inter_stop - inter_start
+                    )
+                    src_slice.append(slice(src_inter_start, src_inter_stop, 1))
+                if pp_list is not None:
+                    src_slice = postprocess_transpose(
+                        list(src_slice), pp_list, reverse=True
+                    )
+                    results.append(
+                        (
+                            src_key,
+                            tuple(src_slice),
+                            tuple(intersection),
+                            pp_list.copy(),
+                        ),
+                    )
+                else:
+                    results.append(
+                        (src_key, tuple(src_slice), tuple(intersection), None)
+                    )
+        return results
+
+    def find_shard_sources(
+        self,
+        target: ShardedWeightDesc,
+    ) -> ShardMapping:
+        target_key, opt_state_name = split_optimizer_state_key(target.key)
+        target_local_shape = target.local_shape
+        target_global_offset = target.global_offset
+        target_global_shape = target.global_shape
+
+        if opt_state_name in [".beta1_pow_acc_0", ".beta2_pow_acc_0"]:
+            assert target_key in self.output_vars, (
+                f"The key {target_key} is not in the output_vars (which is built during load_state_dict)."
+            )
+            tensor = self.output_vars[target_key]
+            target_local_shape = tensor.shape
+            target_global_offset = (0,) * len(target_local_shape)
+            target_global_shape = target_local_shape
+
+        slices = tuple(
+            slice(offset, offset + size, 1)
+            for offset, size in zip(target_global_offset, target_local_shape)
+        )
+
+        results = self.find_source_slices(target_key, slices)
+
+        shard_mappings = []
+
+        target_key = (
+            target_key + opt_state_name
+            if opt_state_name is not None
+            else target_key
+        )
+
+        src_keys = {
+            result[0]
+            for result in results
+            if result[0] not in self.need_remove_input_vars
+        }
+        if opt_state_name in [".beta1_pow_acc_0", ".beta2_pow_acc_0"]:
+            if len(src_keys) == 0:
+                return shard_mappings
+            elif len(src_keys) > 1:
+                logger.warning(
+                    f"{target_key} has multiple sources: {src_keys} (e.g., .beta1_pow_acc_0). Returning one arbitrarily."
+                )
+                src_key = next(iter(src_keys))
+            else:
+                src_key = next(iter(src_keys))
+            return [
+                ShardMappingEntry(
+                    target,
+                    ShardedWeightDesc(
+                        src_key + opt_state_name,
+                        target.local_shape,
+                        target.global_shape,
+                        target.global_offset,
+                        target.dtype,
+                    ),
+                    None,
+                )
+            ]
+
+        for src_key, src_slices, local_slices, pp_list in results:
+            src_var = self.input_vars[src_key]
+            target_model_state_key, target_opt_state_name = (
+                split_optimizer_state_key(target.key)
+            )
+            if target_opt_state_name is None:
+                if src_var.dtype != target.dtype:
+                    assert pp_list is not None and target.dtype in str(
+                        pp_list
+                    ), (
+                        "Direct assignment of Tensors with different types is prohibited in AOA. "
+                        f"If you want to achieve this functionality, please use the cast semantics provided by AOA. "
+                        f"Now the src_var.dtype is {src_var.dtype}, the target.dtype is {target.dtype}, the pp_list is {pp_list}."
+                        f"The src_key is {src_key}, the target_key is {target.key}."
+                    )
+            else:
+                src_var.dtype = target.dtype
+
+            src_global_shape = src_var.shape
+
+            src_local_shape = tuple(slc.stop - slc.start for slc in src_slices)
+            src_global_offset = tuple(slc.start for slc in src_slices)
+
+            tgt_local_shape = tuple(
+                slc.stop - slc.start for slc in local_slices
+            )
+            tgt_global_offset = tuple(slc.start for slc in local_slices)
+
+            new_src_key = (
+                src_key + opt_state_name
+                if opt_state_name is not None
+                else src_key
+            )
+
+            source_sharded_weight = ShardedWeightDesc(
+                new_src_key,
+                src_local_shape,
+                tuple(src_global_shape),
+                src_global_offset,
+                src_var.dtype,
+            )
+            target_sharded_weight = ShardedWeightDesc(
+                target_key,
+                tgt_local_shape,
+                tuple(target_global_shape),
+                tgt_global_offset,
+                target.dtype,
+            )
+
+            if src_key in self.need_remove_input_vars:
+                mapping_entry = ShardMappingEntry(
+                    target_sharded_weight,
+                    source_sharded_weight,
+                    [],
+                )
+                continue
+
+            shard_mappings.append(
+                ShardMappingEntry(
+                    target_sharded_weight,
+                    source_sharded_weight,
+                    pp_list,
+                )
+            )
+
+        return shard_mappings
+
+
+def postprocess_transpose(
+    li: list[tuple[slice, ...]] | tuple[tuple[slice, ...]],
+    postprocess_list: list[str],
+    reverse: bool = False,
+) -> list[tuple[slice, ...]] | tuple[tuple[slice, ...]]:
+    result = li
+    if reverse:
+        for pp in list(reversed(postprocess_list)):
+            if pp.startswith("["):
+                reversed_transpose = np.argsort(ast.literal_eval(pp)).tolist()
+                result = transpose_list(result, reversed_transpose)
+    else:
+        for pp in postprocess_list:
+            if pp.startswith("["):
+                result = transpose_list(result, ast.literal_eval(pp))
+    return result
+
+
+def transpose_list(
+    li: list[tuple[slice, ...]] | tuple[tuple[slice, ...]],
+    permutation: list[int],
+) -> list[tuple[slice, ...]] | tuple[tuple[slice, ...]]:
+    trans_list = []
+    for idx in permutation:
+        trans_list.append(li[idx])
+    if isinstance(li, tuple):
+        return tuple(trans_list)
+    else:
+        return trans_list
+
+
+def invert_permutation(p: list[int]) -> list[int]:
+    q = [0] * len(p)
+    for i, pi in enumerate(p):
+        q[pi] = i
+    return q

python/paddle/distributed/flex_checkpoint/aoa/lexer.py

@@ -0,0 +1,150 @@
+# 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 re
+from enum import Enum, auto
+
+
+class Token:
+    def __init__(self, type, value):
+        self.type = type
+        self.value = value
+
+    def __repr__(self):
+        return f"Token({self.type}, {self.value!r})"
+
+
+class TokenType(Enum):
+    IDENTIFIER = auto()
+    NUMBER = auto()
+    COLON = auto()
+    LBRACKET = auto()
+    RBRACKET = auto()
+    COMMA = auto()
+    RARROW = auto()
+    STRING = auto()
+    EQUAL = auto()
+    NEWLINE = auto()
+    EOF = auto()
+
+
+class Lexer:
+    token_specification = [
+        ('RARROW', r'->'),
+        ('EQUAL', r'='),
+        ('COLON', r':'),
+        ('LBRACKET', r'\['),
+        ('RBRACKET', r'\]'),
+        ('COMMA', r','),
+        ('NUMBER', r'\d+'),
+        ('STRING', r'"[^"]*"|\'[^\']*\''),
+        ('IDENTIFIER', r'[A-Za-z_][A-Za-z\.\$\_\*\d\^T]*'),
+        ('SKIP', r'[ \t]+'),
+        ('NEWLINE', r'[\r\n]+'),
+        ('MISMATCH', r'.'),
+    ]
+
+    def __init__(self, context, traceback=None):
+        from .macros import macro_registry
+
+        self.macros = [list(d.values())[1] for d in macro_registry.macros]
+        self.get_token = re.compile(
+            '|'.join(
+                f'(?P<{name}>{regex})'
+                for name, regex in self.token_specification
+            )
+        ).match
+        self.context = context
+        self.traceback = traceback
+
+    def tokenize(self, text):
+        pos = 0
+        mo = self.get_token(text, pos)
+        tokens = []
+        if not text.endswith('\n'):
+            text += '\n'
+        while mo is not None:
+            kind = mo.lastgroup
+            value = mo.group()
+            if kind == 'SKIP':
+                pass
+            elif kind == 'MISMATCH':
+                raise RuntimeError(
+                    f'Unexpected character {value!r} at position {pos}'
+                )
+            else:
+                tokens.append(Token(TokenType[kind], value))
+            pos = mo.end()
+            mo = self.get_token(text, pos)
+        return tokens
+
+    def apply_macro(self, expression, macro):
+        if isinstance(expression, str):
+            expression = [expression]
+        new_expression = []
+        for expr in expression:
+            results = macro(self.tokenize(expr), expr, self.context)
+            if isinstance(results, str):
+                new_expression.append(results)
+            else:
+                new_expression.extend(results)
+        return new_expression
+
+    def apply_single_macro_to_all(self, expressions, macro):
+        new_expressions = []
+        macro_name = getattr(macro, "__name__", "macro")
+        for expr in expressions:
+            try:
+                results = macro(self.tokenize(expr), expr, self.context)
+            except (AssertionError, ValueError, KeyError, RuntimeError) as e:
+                if self.traceback:
+                    chain = self.traceback.build_chain(expr)
+                    self.traceback.add_error(
+                        error_message=str(e),
+                        stage=f"{macro_name}",
+                        chain=chain,
+                        error_type=type(e).__name__,
+                    )
+                    self.traceback.print()
+                raise
+
+            if isinstance(results, str):
+                results_list = [results]
+            else:
+                results_list = list(results)
+
+            if self.traceback:
+                if results_list != [expr]:
+                    self.traceback.record_children(
+                        expr, results_list, macro_name
+                    )
+
+            new_expressions.extend(results_list)
+        return new_expressions
+
+    def all_tokens(self, expressions):
+        if self.traceback:
+            self.traceback.register_roots(list(expressions))
+
+        current_expressions = expressions
+        for macro in self.macros:
+            current_expressions = self.apply_single_macro_to_all(
+                current_expressions, macro
+            )
+
+        self.final_expressions = list(current_expressions)
+        tokens = []
+        for expr in current_expressions:
+            tokens.extend(self.tokenize(expr))
+        return tokens

python/paddle/distributed/flex_checkpoint/aoa/macros.py

@@ -0,0 +1,864 @@
+# 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 math
+import re
+from itertools import product
+
+from .lexer import Token, TokenType
+
+
+def macro(name, priority):
+    def decorator(func):
+        macro_registry.register_macro(name, func, priority)
+        return func
+
+    return decorator
+
+
+class MacroRegistry:
+    _instance = None
+
+    def __new__(cls, *args, **kwargs):
+        if cls._instance is None:
+            cls._instance = super().__new__(cls)
+        return cls._instance
+
+    def __init__(self):
+        if not hasattr(self, 'macros'):
+            self.macros = []
+
+    def register_macro(self, name, func, priority):
+        if any(macro['name'] == name for macro in self.macros):
+            raise ValueError(f"Macro '{name}' is already registered.")
+        self.macros.append({'name': name, 'func': func, 'priority': priority})
+        self.macros.sort(key=lambda x: x['priority'], reverse=False)
+
+
+macro_registry = MacroRegistry()
+
+GLOBAL_ATTRIBUTE_KEYWORDS = [
+    "axis",
+    'fused_ffn',
+    'fused_qkv_old',
+    'num_heads',
+    'num_key_value_groups',
+    'permute',
+    'dtype',
+    'fused_qkv',
+    'src_dtype',
+    'dst_dtype',
+]
+
+EXTRA_SUFFIX = [
+    "^T",
+]
+
+
+def extract_axis_and_clean_tokens(tokens):
+    axis = 1
+    for idx, tkn in enumerate(tokens):
+        if tkn.value == "axis" and idx + 2 < len(tokens):
+            axis = int(tokens[idx + 2].value)
+            end_idx = idx + 3
+            if end_idx < len(tokens) - 1:
+                assert tokens[end_idx].value == ",", (
+                    f"The different attributes must split by a comma, but now the token is {tokens[end_idx].value}."
+                )
+                end_idx += 1
+            tokens = tokens[:idx] + tokens[end_idx:]
+            break
+    return axis, tokens
+
+
+# star_macro must be called after layer_id_macro
+@macro(name='star_macro', priority=3)
+def star_macro(tokens, expression, context):
+    STAR_TAG = "*"
+    if STAR_TAG not in expression:
+        return expression
+
+    def _sort_keys_by_numeric_part(prefix, suffix, allkeys):
+        pattern = re.compile(rf"{re.escape(prefix)}(\d+){re.escape(suffix)}")
+        filtered_keys = []
+        for key in allkeys:
+            match = pattern.fullmatch(key)
+            if match:
+                num = int(match.group(1))
+                filtered_keys.append((key, num))
+        sorted_keys = sorted(filtered_keys, key=lambda x: x[1])
+        return [key for key, _ in sorted_keys]
+
+    pre_rarrow = True
+    new_tokens = []
+    for token in tokens:
+        if token.type == TokenType.RARROW:
+            pre_rarrow = False
+        if token.type == TokenType.IDENTIFIER and STAR_TAG in token.value:
+            prefix, suffix = token.value.split(STAR_TAG)
+            allkeys = (
+                context.get_all_dst_state_keys()
+                if not pre_rarrow
+                else context.get_all_src_state_keys()
+            )
+            assert len(allkeys) != 0, (
+                f"No keys found with prefix '{prefix}' and suffix '{suffix}' in "
+                f"{'destination_state_shard_info' if not pre_rarrow else 'source_state_shard_info'}, please check!"
+            )
+            keys = list(_sort_keys_by_numeric_part(prefix, suffix, allkeys))
+            for key in keys:
+                new_tokens.append(Token(TokenType.IDENTIFIER, key))
+                if key != keys[-1]:
+                    new_tokens.append(Token(TokenType.COMMA, ","))
+        else:
+            new_tokens.append(token)
+    new_expression = "".join([token.value for token in new_tokens])
+    return new_expression
+
+
+@macro(name='layer_id_offset_macro', priority=1)
+def layer_id_offset_macro(tokens, expression, context):
+    LAYER_ID_OFFSET_MACRO_TAG = "$LAYER_ID_OFFSET"
+    if LAYER_ID_OFFSET_MACRO_TAG not in expression:
+        return expression
+
+    name_with_layer_id_offset = next(
+        (
+            token.value
+            for token in tokens
+            if token.type == TokenType.IDENTIFIER
+            and LAYER_ID_OFFSET_MACRO_TAG in token.value
+        ),
+        None,
+    )
+    assert name_with_layer_id_offset, (
+        "No $LAYER_ID_OFFSET found in NAME tokens.Please check the aoa_config."
+    )
+    assert all(
+        (t.type != TokenType.IDENTIFIER)
+        or (LAYER_ID_OFFSET_MACRO_TAG in t.value)
+        or (t.value in GLOBAL_ATTRIBUTE_KEYWORDS)
+        for t in tokens
+    ), (
+        f"All IDENTIFIER tokens must contain {LAYER_ID_OFFSET_MACRO_TAG} when a NAME with it is present, except for GLOBAL_ATTRIBUTE_KEYWORDS."
+    )
+
+    match_layer_id_offset = context.get_num_hidden_layers(
+        name_with_layer_id_offset, LAYER_ID_OFFSET_MACRO_TAG
+    )
+    expanded_expressions = []
+
+    match_layer_id_offset = sorted(match_layer_id_offset)
+
+    for layer_id in match_layer_id_offset:
+        expr = ""
+        before_rarrow = True
+        for token in tokens:
+            if token.type == TokenType.RARROW:
+                before_rarrow = False
+            if before_rarrow:
+                cur_layer_id = layer_id
+            else:
+                cur_layer_id = layer_id - 1
+            if token.type == TokenType.IDENTIFIER:
+                if LAYER_ID_OFFSET_MACRO_TAG in token.value:
+                    expr += token.value.replace(
+                        LAYER_ID_OFFSET_MACRO_TAG, str(cur_layer_id)
+                    )
+                elif token.value not in GLOBAL_ATTRIBUTE_KEYWORDS:
+                    expr += f"{token.value}.layer.{cur_layer_id}"
+                else:
+                    expr += token.value
+            else:
+                expr += token.value
+        expanded_expressions.append(expr)
+    return expanded_expressions
+
+
+@macro(name='array_macro', priority=2)
+def array_macro(tokens, expression, context):
+    if "[" not in expression:
+        return expression
+    new_tokens = []
+    idx = 0
+    while idx < len(tokens):
+        if tokens[idx].type == TokenType.LBRACKET:
+            name = tokens[idx - 1].value
+            assert (
+                tokens[idx + 1].type == TokenType.NUMBER
+                and tokens[idx + 2].type == TokenType.COLON
+                and tokens[idx + 3].type == TokenType.NUMBER
+                and tokens[idx + 4].type == TokenType.RBRACKET
+            ), (
+                f"The array macro format is incorrect which is must be like: NAME[START:END], but now the format is {tokens[idx].value}{tokens[idx + 1].value}:{tokens[idx + 3].value}{tokens[idx + 4].value}."
+            )
+            new_tokens.pop()
+            start = int(tokens[idx + 1].value)
+            end = int(tokens[idx + 3].value)
+            for i in range(start, end):
+                new_tokens.append(
+                    Token(TokenType.IDENTIFIER, name + "_" + str(i))
+                )
+                if i != end - 1:
+                    new_tokens.append(Token(TokenType.COMMA, ","))
+            idx += 5
+        else:
+            new_tokens.append(tokens[idx])
+            idx += 1
+    new_expression = "".join([token.value for token in new_tokens])
+    return new_expression
+
+
+@macro(name='fused_qkv_old_macro', priority=6)
+def fused_qkv_old_macro(tokens, expression, context):
+    FUSED_QKV_OLD_TAG = "fused_qkv_old"
+    if not any(tkn.value == FUSED_QKV_OLD_TAG for tkn in tokens):
+        return expression
+
+    axis, tokens = extract_axis_and_clean_tokens(tokens)
+
+    attn_head_num = None
+    num_key_value_groups = None
+    fused_qkv_old_pos = None
+    rarrow_pos = None
+    right_var_end_pos = None
+
+    for idx, token in enumerate(tokens):
+        if token.type == TokenType.IDENTIFIER:
+            if token.value == "num_heads" and idx + 2 < len(tokens):
+                attn_head_num = int(tokens[idx + 2].value)
+            elif token.value == "num_key_value_groups" and idx + 2 < len(
+                tokens
+            ):
+                num_key_value_groups = int(tokens[idx + 2].value)
+            elif token.value == FUSED_QKV_OLD_TAG:
+                fused_qkv_old_pos = idx
+        elif token.type == TokenType.RARROW and rarrow_pos is None:
+            rarrow_pos = idx
+        if (
+            right_var_end_pos is None
+            and token.type == TokenType.IDENTIFIER
+            and token.value
+            in {FUSED_QKV_OLD_TAG, "num_heads", "num_key_value_groups"}
+        ):
+            right_var_end_pos = idx + 1
+
+    assert attn_head_num and attn_head_num > 0, (
+        f"num_heads must be positive.(got: {attn_head_num})."
+    )
+    assert num_key_value_groups and num_key_value_groups > 0, (
+        f"num_key_value_groups must be positive.(got: {num_key_value_groups})."
+    )
+    assert fused_qkv_old_pos is not None, (
+        f"No fused_qkv_old tag found in expression. The tag must be {FUSED_QKV_OLD_TAG}."
+    )
+    assert rarrow_pos is not None, "No -> found in expression."
+    assert attn_head_num % num_key_value_groups == 0, (
+        f"num_heads ({attn_head_num}) must be divisible by num_key_value_groups ({num_key_value_groups})."
+    )
+
+    results = []
+    num_key_value_heads = num_key_value_groups
+    if rarrow_pos == 1:
+        src_qkv_weight_name = tokens[0].value
+        if fused_qkv_old_pos > 4:
+            dst_qkv_weight_name = None
+        else:
+            dst_qkv_weight_name = tokens[2].value
+
+        if context.aoa_config_reverse:
+            dst_state_shard_num = context.get_src_state_shard_num(
+                dst_qkv_weight_name
+            )
+            src_state_shard_num = (
+                context.get_dst_state_shard_num(src_qkv_weight_name)
+                if src_qkv_weight_name is not None
+                else 1
+            )
+        else:
+            src_state_shard_num = context.get_src_state_shard_num(
+                src_qkv_weight_name
+            )
+            dst_state_shard_num = (
+                context.get_dst_state_shard_num(dst_qkv_weight_name)
+                if dst_qkv_weight_name is not None
+                else 1
+            )
+
+        configs = [
+            (src_state_shard_num, src_qkv_weight_name),
+            (dst_state_shard_num, dst_qkv_weight_name),
+        ]
+
+        head_config = [
+            ("Q", attn_head_num),
+            ("K", num_key_value_heads),
+            ("V", num_key_value_heads),
+        ]
+
+        def gen_expr(tp_degree, num_heads, tp_rank, comp):
+            start = tp_rank * num_heads // tp_degree
+            count = num_heads // tp_degree
+            return ",".join(
+                f"fused_qkv_old_tmp.{comp}_{i}"
+                for i in range(start, start + count)
+            )
+
+        for idx, (tp_degree, qkv_weight_name) in enumerate(configs):
+            qkv_parts = [
+                gen_expr(tp_degree, n, tp_rank, c)
+                for tp_rank in range(tp_degree)
+                for c, n in head_config
+            ]
+            if idx == 0:
+                mapping = (
+                    f"{qkv_weight_name} -> {','.join(qkv_parts)}, axis={axis}"
+                )
+                results.append(mapping)
+            elif qkv_weight_name is not None:
+                mapping = (
+                    f"{','.join(qkv_parts)} -> {qkv_weight_name}, axis={axis}"
+                )
+                results.append(mapping)
+
+        if fused_qkv_old_pos > 4:
+
+            def _generate_expr(prefix, count, target_name):
+                elements = ",".join(
+                    f"fused_qkv_old_tmp.{prefix}_{i}" for i in range(count)
+                )
+                return f"{elements} -> {target_name}, axis={axis}"
+
+            q_name = tokens[2].value
+            k_name = tokens[4].value
+            v_name = tokens[6].value
+
+            results.append(_generate_expr("Q", attn_head_num, q_name))
+            results.append(_generate_expr("K", num_key_value_heads, k_name))
+            results.append(_generate_expr("V", num_key_value_heads, v_name))
+    elif rarrow_pos == 5:
+        q_name = tokens[0].value
+        k_name = tokens[2].value
+        v_name = tokens[4].value
+        dst_qkv_weight_name = tokens[6].value
+
+        fused_qkv_tmp_name = f"{q_name}.{k_name}.{v_name}.tmp"
+        results.append(
+            f"{q_name},{k_name},{v_name}  ->  {fused_qkv_tmp_name}, axis={axis}"
+        )
+        dst_state_shard_num = context.get_dst_state_shard_num(
+            dst_qkv_weight_name
+        )
+
+        configs = [
+            (1, fused_qkv_tmp_name),
+            (dst_state_shard_num, dst_qkv_weight_name),
+        ]
+
+        head_config = [
+            ("Q", attn_head_num),
+            ("K", num_key_value_heads),
+            ("V", num_key_value_heads),
+        ]
+
+        def gen_expr(tp_degree, num_heads, tp_rank, comp):
+            start = tp_rank * num_heads // tp_degree
+            count = num_heads // tp_degree
+            return ",".join(
+                f"fused_qkv_old_tmp.{comp}_{i}"
+                for i in range(start, start + count)
+            )
+
+        for idx, (tp_degree, qkv_weight_name) in enumerate(configs):
+            qkv_parts = [
+                gen_expr(tp_degree, n, tp_rank, c)
+                for tp_rank in range(tp_degree)
+                for c, n in head_config
+            ]
+            if idx == 0:
+                mapping = (
+                    f"{qkv_weight_name} -> {','.join(qkv_parts)}, axis={axis}"
+                )
+            else:
+                mapping = (
+                    f"{','.join(qkv_parts)} -> {qkv_weight_name}, axis={axis}"
+                )
+            results.append(mapping)
+    else:
+        raise ValueError(
+            f"Unsupported fused_qkv_old macro format: {expression}."
+        )
+    return results
+
+
+@macro(name='fused_ffn_macro', priority=6)
+def fused_ffn_macro(tokens, expression, context):
+    FUSED_FFN_TAG = "fused_ffn"
+    if not any(tkn.value == FUSED_FFN_TAG for tkn in tokens):
+        return expression
+
+    axis, tokens = extract_axis_and_clean_tokens(tokens)
+
+    rarrow_pos = None
+    fused_ffn_pos = None
+    for idx, token in enumerate(tokens):
+        if token.type == TokenType.RARROW and rarrow_pos is None:
+            rarrow_pos = idx
+        elif (
+            token.type == TokenType.IDENTIFIER and token.value == FUSED_FFN_TAG
+        ):
+            fused_ffn_pos = idx
+    assert rarrow_pos is not None, "No -> found in expression."
+    assert fused_ffn_pos is not None, (
+        f"No fused_ffn tag found in expression. The tag must be {FUSED_FFN_TAG}."
+    )
+    results = []
+    if rarrow_pos == 1:
+        src_ffn_weight_name = tokens[0].value
+        if fused_ffn_pos == 4:
+            dst_ffn_weight_name = tokens[2].value
+        else:
+            dst_ffn_weight_name = None
+        if context.aoa_config_reverse:
+            dst_state_shard_num = context.get_src_state_shard_num(
+                dst_ffn_weight_name
+            )
+            src_state_shard_num = (
+                context.get_dst_state_shard_num(src_ffn_weight_name)
+                if src_ffn_weight_name is not None
+                else 1
+            )
+        else:
+            src_state_shard_num = context.get_src_state_shard_num(
+                src_ffn_weight_name
+            )
+            dst_state_shard_num = (
+                context.get_dst_state_shard_num(dst_ffn_weight_name)
+                if dst_ffn_weight_name is not None
+                else 1
+            )
+        splited_num = math.lcm(src_state_shard_num, dst_state_shard_num)
+
+        configs = [
+            (src_state_shard_num, src_ffn_weight_name),
+            (dst_state_shard_num, dst_ffn_weight_name),
+        ]
+        split_config = [("GATE", splited_num), ("UP", splited_num)]
+
+        def gen_expr(tp_degree, splited_num, tp_rank, comp):
+            return ",".join(
+                f"fused_ffn_tmp.{comp}_{tp_rank * splited_num // tp_degree + idx}"
+                for idx in range(splited_num // tp_degree)
+            )
+
+        for idx, (tp_degree, ffn_weight_name) in enumerate(configs):
+            ffn_parts = [
+                gen_expr(tp_degree, n, tp_rank, c)
+                for tp_rank in range(tp_degree)
+                for c, n in split_config
+            ]
+            if idx == 0:
+                results.append(
+                    f"{ffn_weight_name}  -> {','.join(ffn_parts)}, axis={axis}"
+                )
+            elif ffn_weight_name is not None:
+                results.append(
+                    f"{','.join(ffn_parts)} -> {ffn_weight_name}, axis={axis}"
+                )
+        if fused_ffn_pos > 4:
+
+            def _generate_expr(prefix, count, target_name):
+                elements = ",".join(
+                    f"fused_ffn_tmp.{prefix}_{i}" for i in range(count)
+                )
+                return f"{elements} -> {target_name}, axis={axis}"
+
+            gate_name = tokens[2].value
+            up_name = tokens[4].value
+
+            results.append(_generate_expr("GATE", splited_num, gate_name))
+            results.append(_generate_expr("UP", splited_num, up_name))
+
+    elif rarrow_pos == 3:
+        gate_name = tokens[0].value
+        up_name = tokens[2].value
+        dst_ffn_weight_name = tokens[4].value
+
+        fused_gate_up_tmp_name = f"{gate_name}.{up_name}.tmp"
+        results.append(
+            f"{gate_name},{up_name}  ->  {fused_gate_up_tmp_name}, axis={axis}"
+        )
+        dst_state_shard_num = context.get_dst_state_shard_num(
+            dst_ffn_weight_name
+        )
+
+        configs = [
+            (1, fused_gate_up_tmp_name),
+            (dst_state_shard_num, dst_ffn_weight_name),
+        ]
+
+        split_config = [
+            ("GATE", dst_state_shard_num),
+            ("UP", dst_state_shard_num),
+        ]
+
+        def gen_expr(tp_degree, splited_num, tp_rank, comp):
+            return ",".join(
+                f"fused_ffn_tmp.{comp}_{tp_rank * splited_num // tp_degree + idx}"
+                for idx in range(splited_num // tp_degree)
+            )
+
+        for idx, (tp_degree, ffn_weight_name) in enumerate(configs):
+            ffn_parts = [
+                gen_expr(tp_degree, n, tp_rank, c)
+                for tp_rank in range(tp_degree)
+                for c, n in split_config
+            ]
+            if idx == 0:
+                results.append(
+                    f"{ffn_weight_name}  -> {','.join(ffn_parts)}, axis={axis}"
+                )
+            else:
+                results.append(
+                    f"{','.join(ffn_parts)} -> {ffn_weight_name}, axis={axis}"
+                )
+    else:
+        raise ValueError(f"Unsupported fused_ffn macro format: {expression}.")
+    return results
+
+
+@macro(name='transpose_macro', priority=5)
+def transpose_macro(tokens, expression, context):
+    TRANSPOSE_TAG = "^T"
+
+    if TRANSPOSE_TAG not in expression:
+        return expression
+
+    transpose_vars = set()
+    new_expression = ""
+    rarrow_pos = None
+
+    for idx, token in enumerate(tokens):
+        if token.type == TokenType.RARROW:
+            rarrow_pos = idx
+            break
+
+    assert rarrow_pos is not None, "No -> found in expression."
+
+    for token in tokens[rarrow_pos + 1 :]:
+        if token.type == TokenType.IDENTIFIER and token.value.endswith(
+            TRANSPOSE_TAG
+        ):
+            raise ValueError(
+                "Cannot assign to transpose (e.g., 'A -> B^T').\n"
+                "B^T is not a real variable, just a view.\n"
+                "Assign first:  A -> B\n"
+                "Then transpose: B^T -> B"
+            )
+    for token in tokens:
+        if token.type == TokenType.IDENTIFIER and token.value.endswith(
+            TRANSPOSE_TAG
+        ):
+            var_name = token.value[: -len(TRANSPOSE_TAG)]
+            transpose_vars.add(var_name)
+            new_expression += var_name + "_transpose_tmp"
+        else:
+            new_expression += token.value
+
+    results = [
+        f'{var} -> {var}_transpose_tmp, permute = "[]"'
+        for var in transpose_vars
+    ]
+    results.append(new_expression)
+    return results
+
+
+@macro(name='fused_qkv_macro', priority=6)
+def fused_qkv_macro(tokens, expression, context):
+    FUSED_QKV_TAG = "fused_qkv"
+    if not any(tkn.value == FUSED_QKV_TAG for tkn in tokens):
+        return expression
+
+    axis, tokens = extract_axis_and_clean_tokens(tokens)
+
+    attn_head_num = num_heads = None
+    num_key_value_groups = None
+    fused_qkv_pos = None
+    rarrow_pos = None
+
+    for idx, token in enumerate(tokens):
+        if token.type == TokenType.IDENTIFIER:
+            if token.value == "num_heads" and idx + 2 < len(tokens):
+                attn_head_num = int(tokens[idx + 2].value)
+            elif token.value == "num_key_value_groups" and idx + 2 < len(
+                tokens
+            ):
+                num_key_value_groups = int(tokens[idx + 2].value)
+            elif token.value == FUSED_QKV_TAG:
+                fused_qkv_pos = idx
+        elif token.type == TokenType.RARROW and rarrow_pos is None:
+            rarrow_pos = idx
+
+    assert attn_head_num and attn_head_num > 0, (
+        f"num_heads must be positive (got: {attn_head_num})"
+    )
+    assert num_key_value_groups and num_key_value_groups > 0, (
+        f"num_key_value_groups must be positive (got: {num_key_value_groups})"
+    )
+    assert fused_qkv_pos is not None, (
+        f"No fused_qkv tag found in expression. The tag must be {FUSED_QKV_TAG}."
+    )
+    assert rarrow_pos is not None, "No -> found in expression."
+    assert rarrow_pos == 1 or rarrow_pos == 5, (
+        "Only support q,k,v -> fused_qkv or fused_qkv -> q,k,v patterns"
+    )
+    assert attn_head_num % num_key_value_groups == 0, (
+        f"num_heads ({attn_head_num}) must be divisible by num_key_value_groups ({num_key_value_groups})."
+    )
+
+    num_key_value_heads = attn_head_num // num_key_value_groups
+
+    def make_names(base, n):
+        return [f"{base}{i}" for i in range(n)]
+
+    results = []
+
+    if rarrow_pos == 1:
+        fused_qkv_var = tokens[0].value
+        q_var = tokens[rarrow_pos + 1].value
+        k_var = tokens[rarrow_pos + 3].value
+        v_var = tokens[rarrow_pos + 5].value
+
+        q_names = make_names(q_var, attn_head_num)
+        k_names = make_names(k_var, num_key_value_groups)
+        v_names = make_names(v_var, num_key_value_groups)
+
+        fused_qkv_order = []
+        for g in range(num_key_value_groups):
+            fused_qkv_order.extend(
+                q_names[g * num_key_value_heads : (g + 1) * num_key_value_heads]
+            )
+            fused_qkv_order.append(k_names[g])
+            fused_qkv_order.append(v_names[g])
+        results.append(
+            f"{fused_qkv_var} -> {','.join(fused_qkv_order)}, axis={axis}"
+        )
+
+        results.append(f"{','.join(q_names)} -> {q_var}, axis={axis}")
+        results.append(f"{','.join(k_names)} -> {k_var}, axis={axis}")
+        results.append(f"{','.join(v_names)} -> {v_var}, axis={axis}")
+
+        return results
+
+    elif rarrow_pos == 5:
+        q_var = tokens[0].value
+        k_var = tokens[2].value
+        v_var = tokens[4].value
+        fused_qkv_var = tokens[rarrow_pos + 1].value
+
+        q_names = make_names(q_var, attn_head_num)
+        k_names = make_names(k_var, num_key_value_groups)
+        v_names = make_names(v_var, num_key_value_groups)
+
+        results.append(f"{q_var} -> {','.join(q_names)}, axis={axis}")
+        results.append(f"{k_var} -> {','.join(k_names)}, axis={axis}")
+        results.append(f"{v_var} -> {','.join(v_names)}, axis={axis}")
+
+        fused_qkv_order = []
+        for g in range(num_key_value_groups):
+            fused_qkv_order.extend(
+                q_names[g * num_key_value_heads : (g + 1) * num_key_value_heads]
+            )
+            fused_qkv_order.append(k_names[g])
+            fused_qkv_order.append(v_names[g])
+        results.append(
+            f"{','.join(fused_qkv_order)} -> {fused_qkv_var}, axis={axis}"
+        )
+        return results
+
+    else:
+        return expression
+
+
+class IDMatcher:
+    def __init__(
+        self,
+        source_keys: list[str],
+        extra_suffixes: list[str],
+        allowed_placeholders: list[str],
+    ):
+        self.source_keys = set(source_keys)
+        self.allowed_placeholders = allowed_placeholders
+        # Dynamically build regex pattern from allowed placeholders
+        placeholder_pattern = '|'.join(
+            re.escape(ph) for ph in self.allowed_placeholders
+        )
+        self._placeholder_pattern = re.compile(f'({placeholder_pattern})')
+        self.extra_suffixes = sorted(extra_suffixes, key=lambda x: (-len(x), x))
+
+    def _remove_extra_suffixes(self, key: str) -> str:
+        for sfx in self.extra_suffixes:
+            if key.endswith(sfx):
+                key = key[: -len(sfx)]
+                break
+        return key
+
+    def _pattern_to_regex(self, pattern: str) -> tuple[re.Pattern, list[str]]:
+        placeholders = sorted(set(self._placeholder_pattern.findall(pattern)))
+        regex_str = re.escape(pattern)
+        for ph in placeholders:
+            group_name = ph[1:]
+            regex_str = regex_str.replace(
+                re.escape(ph), f'(?P<{group_name}>\\d+)'
+            )
+        return re.compile(f'^{regex_str}$'), [ph[1:] for ph in placeholders]
+
+    def _substitute_ids(self, pattern: str, id_dict: dict[str, int]) -> str:
+        key = pattern
+        for ph, value in id_dict.items():
+            key = key.replace(f'${ph}', str(value))
+        return key
+
+    def find_matches(self, pattern: str) -> dict[str, list[int]]:
+        pattern = self._remove_extra_suffixes(pattern)
+        regex, ph_names = self._pattern_to_regex(pattern)
+        id_values = {ph: set() for ph in ph_names}
+        for key in self.source_keys:
+            match = regex.match(key)
+            if match:
+                for k, v in match.groupdict().items():
+                    id_values[k].add(int(v))
+        return {k: sorted(vs) for k, vs in id_values.items()}
+
+
+# Global registry for allowed_placeholders
+_REGISTERED_PLACEHOLDERS = ['$EXPERT_ID', '$LAYER_ID']
+
+
+# TODO: need to adapt the scene of temp_layers.\$LAYER_ID.weight -> dst_layers.\$LAYER_ID.weight
+@macro(name='id_macro', priority=1)
+def id(tokens, expression, context):
+    allowed_placeholders = _REGISTERED_PLACEHOLDERS
+    has_allowed_placeholder = any(
+        ph in expression for ph in allowed_placeholders
+    )
+    if not has_allowed_placeholder:
+        return expression
+
+    if not context.aoa_config_reverse:
+        name_with_id = next(
+            (
+                token.value
+                for token in tokens
+                if token.type == TokenType.IDENTIFIER
+                and any(ph in token.value for ph in allowed_placeholders)
+            ),
+            None,
+        )
+    else:
+        flag_right_var = False
+        for token in tokens:
+            if token.type == TokenType.RARROW:
+                flag_right_var = True
+            if token.type == TokenType.IDENTIFIER and any(
+                ph in token.value for ph in allowed_placeholders
+            ):
+                if flag_right_var:
+                    name_with_id = token.value
+                    break
+
+    assert name_with_id is not None, "No $ID found in NAME tokens"
+    all_src_state_keys = context.get_all_src_state_keys()
+    id_matcher = IDMatcher(
+        all_src_state_keys, EXTRA_SUFFIX, allowed_placeholders
+    )
+    valid_id_combos = id_matcher.find_matches(name_with_id)
+    valid_keys = list(valid_id_combos.keys())
+    IDENTIFIER_tokens = []
+    for token in tokens:
+        if token.value in GLOBAL_ATTRIBUTE_KEYWORDS:
+            break
+        if token.type == TokenType.IDENTIFIER:
+            IDENTIFIER_tokens.append(token)
+
+    for token in IDENTIFIER_tokens:
+        assert all(k in token.value for k in valid_keys), (
+            f"The token: {token.value} must contain all of the following keys: {valid_keys}.When use the id macro all IDENTIFIER tokens must contain the same ID placeholders."
+        )
+
+    def dict_cartesian_tuples(d: dict[str, list[int]]):
+        keys = list(d.keys())
+        value_lists = [d[k] for k in keys]
+        for prod in product(*value_lists):
+            yield tuple(zip(keys, prod))
+
+    results = []
+    id_combs = dict_cartesian_tuples(valid_id_combos)
+    id_combs = sorted(id_combs)
+    for id_comb in id_combs:
+        cur_statement = ""
+        for tkn in tokens:
+            tkn_val = tkn.value
+            if tkn.type == TokenType.IDENTIFIER and any(
+                ph in tkn.value for ph in allowed_placeholders
+            ):
+                for id_tag, id_val in id_comb:
+                    tkn_val = tkn_val.replace("$" + id_tag, str(id_val))
+                cur_statement += tkn_val
+            else:
+                cur_statement += tkn_val
+        results.append(cur_statement)
+
+    return results
+
+
+# This macro processes variable mappings between source and destination states,
+# but it requires that all expansion macros (layer_id_macro, expert_id_macro,
+# star_macro, array_macro, etc.) have already been executed to expand template
+# variables into concrete variable names.
+@macro(name='get_var_mapping_chain_macro', priority=4)
+def get_var_mapping_chain_macro(tokens, expression, context):
+    flag_left_var = True
+    left_var_list = []
+    right_var_list = []
+    for tkn in tokens:
+        if tkn.value in GLOBAL_ATTRIBUTE_KEYWORDS:
+            break
+        if tkn.type == TokenType.RARROW:
+            flag_left_var = False
+        if tkn.type == TokenType.IDENTIFIER:
+            extra_suffix_removed_value = tkn.value
+            for sfx in EXTRA_SUFFIX:
+                extra_suffix_removed_value = (
+                    extra_suffix_removed_value.removesuffix(sfx)
+                )
+            if flag_left_var:
+                left_var_list.append(extra_suffix_removed_value)
+            else:
+                right_var_list.append(extra_suffix_removed_value)
+    assert len(left_var_list) == 1 or len(right_var_list) == 1, (
+        "Left or right variable must have the only one element,the aoa_statements not support 'multiple var -> multiple var' pattern."
+    )
+    if len(left_var_list) == 1:
+        context.left_var_to_right_var_mapping[left_var_list[0]] = right_var_list
+        for right_var in right_var_list:
+            context.right_var_from_left_var_mapping[right_var] = left_var_list
+    else:
+        context.right_var_from_left_var_mapping[right_var_list[0]] = (
+            left_var_list
+        )
+        for left_var in left_var_list:
+            context.left_var_to_right_var_mapping[left_var] = right_var_list
+    return expression

python/paddle/distributed/flex_checkpoint/aoa/parser.py

@@ -0,0 +1,142 @@
+# 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 .lexer import Token, TokenType
+
+
+class Statement:
+    def __init__(self, left_vars, right_vars, attrs):
+        self.left_vars = left_vars  # List[Var]
+        self.right_vars = right_vars  # List[Var]
+        self.attrs = attrs  # List[Attribute]
+
+    def __repr__(self):
+        return f"Statement({self.left_vars} -> {self.right_vars}, attrs={self.attrs})"
+
+
+class Var:
+    def __init__(self, name):
+        self.name = name
+
+    def __repr__(self):
+        return self.name
+
+
+class Attribute:
+    def __init__(self, key, value):
+        self.key = key
+        self.value = value
+
+    def __repr__(self):
+        return f"{self.key}={self.value!r}"
+
+
+class Parser:
+    """
+    AOA Grammar
+    PROGRAM   ::= { STATEMENT }
+
+    STATEMENT ::= VAR_LIST '->' VAR ',' ATTR_LIST       // meige
+                | VAR '->' VAR_LIST ',' ATTR_LIST       // split
+                | VAR '->' VAR ',' ATTR_LIST            // single variable mapping + attributes
+                | VAR '->' VAR                          // single variable mapping, rename
+
+    VAR_LIST  ::= VAR { ',' VAR }
+    VAR       ::= IDENTIFIER
+    ATTR_LIST ::= ATTRIBUTE { ',' ATTRIBUTE }
+    ATTRIBUTE ::= IDENTIFIER '=' VALUE
+    VALUE     ::= NUMBER | STRING
+    """
+
+    def __init__(self, tokens):
+        self.tokens = tokens
+        self.pos = 0
+
+    def at_end(self):
+        return self.peek().type == TokenType.EOF
+
+    def peek(self, offset=0):
+        if self.pos + offset >= len(self.tokens):
+            return Token(TokenType.EOF, '')
+        return self.tokens[self.pos + offset]
+
+    def consume(self, expected_type=None):
+        tok = self.peek()
+        if expected_type and tok.type != expected_type:
+            raise SyntaxError(
+                f'Expected {expected_type}, got {tok.type} at pos {self.pos}'
+            )
+        self.pos += 1
+        return tok
+
+    def expect(self, expected_type):
+        return self.consume(expected_type)
+
+    def skip_newlines(self):
+        while self.peek().type == TokenType.NEWLINE:
+            self.consume()
+
+    def parse_program(self):
+        stmts = []
+        self.skip_newlines()
+        while not self.at_end():
+            stmt = self.parse_statement()
+            stmts.append(stmt)
+            self.skip_newlines()
+        return stmts
+
+    def parse_statement(self):
+        left_vars = [self.parse_var()]
+        while self.peek().type == TokenType.COMMA:
+            self.consume(TokenType.COMMA)
+            left_vars.append(self.parse_var())
+        self.expect(TokenType.RARROW)
+        right_vars = [self.parse_var()]
+        while self.peek().type == TokenType.COMMA:
+            # Lookahead for attribute: IDENT '=' after COMMA means attribute starts
+            if (
+                self.peek(1).type == TokenType.IDENTIFIER
+                and self.peek(2).type == TokenType.EQUAL
+            ):
+                break
+            self.consume(TokenType.COMMA)
+            right_vars.append(self.parse_var())
+        attrs = []
+        if self.peek().type == TokenType.COMMA:
+            self.consume(TokenType.COMMA)
+            attrs = self.parse_attr_list()
+        return Statement(left_vars, right_vars, attrs)
+
+    def parse_var(self):
+        name = self.expect(TokenType.IDENTIFIER).value
+        return Var(name)
+
+    def parse_attr_list(self):
+        attrs = [self.parse_attribute()]
+        while self.peek().type == TokenType.COMMA:
+            self.consume(TokenType.COMMA)
+            attrs.append(self.parse_attribute())
+        return attrs
+
+    def parse_attribute(self):
+        key = self.expect(TokenType.IDENTIFIER).value
+        self.expect(TokenType.EQUAL)
+        val_tok = self.consume()
+        if val_tok.type == TokenType.NUMBER:
+            val = int(val_tok.value)
+        elif val_tok.type == TokenType.STRING:
+            val = val_tok.value.strip('"').strip("'")
+        else:
+            raise SyntaxError(f'Unexpected value: {val_tok}')
+        return Attribute(key, val)

python/paddle/distributed/flex_checkpoint/aoa/traceback.py

@@ -0,0 +1,133 @@
+# 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
+
+
+class AOATraceback:
+    """
+    When error occurs, print the chain of "original aoa_statement -> ... -> current expression".
+    """
+
+    def __init__(self) -> None:
+        self.records: list[dict] = []
+        self.last_error_chain: list[str] = []
+        self.last_error_message: str = ""
+        self.last_error_stage: str = ""
+        self.last_error_type: str = ""
+        self.parent_map: dict[str, str | None] = {}
+        self.child_macro_map: dict[str, str] = {}
+
+    def register_roots(self, expressions: list[str]) -> None:
+        """Register the original aoa_statements as the root nodes of the chain."""
+        for expr in expressions:
+            self.parent_map.setdefault(expr, None)
+
+    def record_children(
+        self, parent: str, children: list[str], macro_name: str | None = None
+    ) -> None:
+        """Record the children expressions obtained by the parent expression, and mark the macro name used."""
+        macro = macro_name or "Expanded"
+        for child in children:
+            if child == parent:
+                continue
+            self.parent_map[child] = parent
+            self.child_macro_map[child] = macro
+
+    def build_chain(self, expr: str) -> list[str]:
+        """Build the chain from the root to expr by tracing back from the current expression."""
+        chain: list[str] = []
+        visited = set()
+        cur = expr
+        while cur is not None and cur not in visited:
+            chain.append(cur)
+            visited.add(cur)
+            cur = self.parent_map.get(cur)
+        chain.reverse()
+        return chain
+
+    def add_error(
+        self,
+        error_message: str,
+        stage: str,
+        chain: list[str],
+        error_type: str = "",
+    ) -> None:
+        """Record the error chain and information."""
+        self.last_error_chain = chain
+        self.last_error_message = error_message
+        self.last_error_stage = stage
+        self.last_error_type = error_type or ""
+        self.records.append(
+            {
+                "type": "error",
+                "stage": stage,
+                "message": error_message,
+                "error_type": self.last_error_type,
+                "chain": chain,
+            }
+        )
+
+    def format_traceback(self) -> str:
+        lines: list[str] = []
+        header_text = " AOA Traceback (related chain) "
+        header = f"===={header_text}===="
+        footer = "=" * len(header)
+
+        if self.last_error_chain:
+            lines.append(header)
+            indent_unit = "    "
+
+            lines.append("| Origin AOA Statement")
+            origin_expr = self.last_error_chain[0].replace("\n", " ")
+            lines.append(f"|-> {origin_expr}")
+
+            for level, expr in enumerate(self.last_error_chain[1:], start=1):
+                indent = indent_unit * level
+                single_line_expr = expr.replace("\n", " ")
+                macro = self.child_macro_map.get(
+                    expr, self.last_error_stage or "Expanded"
+                )
+                lines.append(f"{indent}| {macro}")
+                lines.append(f"{indent}|-> {single_line_expr}")
+
+            if self.last_error_message:
+                err_title = self.last_error_type or "Error"
+                stage_str = (
+                    f" [{self.last_error_stage}]"
+                    if self.last_error_stage
+                    else ""
+                )
+                err_level = len(self.last_error_chain)
+                indent = indent_unit * err_level
+                single_line_msg = self.last_error_message.replace("\n", " ")
+                lines.append(f"{indent}| Error")
+                lines.append(
+                    f"{indent}|-> ({err_title}{stage_str}) {single_line_msg}"
+                )
+
+            lines.append(footer)
+        else:
+            lines.append(header)
+            lines.append("(No trace records)")
+            lines.append(footer)
+
+        return "\n".join(lines)
+
+    def print(self, logger=None) -> None:
+        text = self.format_traceback()
+        if logger:
+            logger.error(text)
+        else:
+            print(text)

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