chore: import upstream snapshot with attribution

test/flex_checkpoint/CMakeLists.txt

@@ -0,0 +1,55 @@
+file(
+  GLOB TEST_OPS
+  RELATIVE "${CMAKE_CURRENT_SOURCE_DIR}"
+  "test_*.py")
+string(REPLACE ".py" "" TEST_OPS "${TEST_OPS}")
+
+foreach(TEST_OP ${TEST_OPS})
+  if(${TEST_OP} STREQUAL "test_strategy_conversion")
+    set(WORKFLOW_SCRIPT ${CMAKE_CURRENT_SOURCE_DIR}/${TEST_OP}.py)
+
+    execute_process(
+      COMMAND ${PYTHON_EXECUTABLE} ${WORKFLOW_SCRIPT} --list_tests
+      OUTPUT_VARIABLE TEST_CASE_LIST
+      OUTPUT_STRIP_TRAILING_WHITESPACE)
+    string(REPLACE "\n" ";" TEST_CASE_LIST "${TEST_CASE_LIST}")
+
+    foreach(TEST_CASE ${TEST_CASE_LIST})
+      string(REPLACE "__main__.TestStrategyConversion.test_" "" TEST_CASE_ALIAS
+                     ${TEST_CASE})
+
+      add_test(NAME ${TEST_OP}.${TEST_CASE_ALIAS}
+               COMMAND ${PYTHON_EXECUTABLE} -m unittest ${TEST_CASE})
+    endforeach()
+  else()
+    py_test_modules(${TEST_OP} MODULES ${TEST_OP})
+  endif()
+endforeach()
+
+set(GPU_ONLY_DISTRIBUTED_TESTS
+    test_sharded_state_dict test_strategy_conversion
+    test_load_state_dict_transpose test_model_full_param)
+
+if(TEST test_sharded_state_dict)
+  set_tests_properties(test_sharded_state_dict PROPERTIES TIMEOUT 480)
+endif()
+
+if(TEST test_model_full_param)
+  set_tests_properties(test_model_full_param PROPERTIES TIMEOUT 480)
+endif()
+
+if(NOT (WITH_DISTRIBUTE AND WITH_GPU))
+  get_property(
+    ALL_TESTS
+    DIRECTORY
+    PROPERTY TESTS)
+  foreach(CURRENT_TEST_NAME ${ALL_TESTS})
+    foreach(SUITE_NAME ${GPU_ONLY_DISTRIBUTED_TESTS})
+      if("${CURRENT_TEST_NAME}" STREQUAL "${SUITE_NAME}"
+         OR "${CURRENT_TEST_NAME}" MATCHES "^${SUITE_NAME}\\.")
+        message(STATUS "Disabling GPU/Dist test: ${CURRENT_TEST_NAME}")
+        set_tests_properties("${CURRENT_TEST_NAME}" PROPERTIES DISABLED TRUE)
+      endif()
+    endforeach()
+  endforeach()
+endif()

test/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.

test/flex_checkpoint/load_state_dict_cast_logic.py

@@ -0,0 +1,83 @@
+# 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 os
+import tempfile
+
+import numpy as np
+
+import paddle
+import paddle.distributed as dist
+from paddle.distributed import fleet
+from paddle.distributed.fleet.layers.mpu import (
+    ColumnParallelLinear,
+)
+from paddle.nn import Layer
+
+
+class SimpleMLP(Layer):
+    def __init__(self, in_features=1024, out_features=1024):
+        super().__init__()
+        self.linear = ColumnParallelLinear(
+            in_features, out_features, has_bias=False
+        )
+
+    def forward(self, x):
+        x = self.linear(x)
+        return x
+
+
+class TestLoadStateDictCastLogic:
+    def __init__(self):
+        self.aoa_config = {"aoa_statements": [os.getenv("aoa_statements")]}
+        self.ckpt_path = tempfile.TemporaryDirectory().name
+        self.in_features = 1024
+        self.out_features = 2048
+
+    def run_test(self):
+        self.run_save_state_dict()
+        model = SimpleMLP()
+        model_cast = SimpleMLP()
+        model_cast = paddle.amp.decorate(
+            models=model_cast,
+            optimizers=None,
+            level="O2",
+            dtype="float16",
+        )
+        sharded_state_dict = model.sharded_state_dict()
+        sharded_state_dict_trans = model_cast.sharded_state_dict()
+        dist.load_state_dict(sharded_state_dict, self.ckpt_path)
+        dist.load_state_dict(
+            sharded_state_dict_trans, self.ckpt_path, aoa_config=self.aoa_config
+        )
+        state_dict_1_after_load = model.state_dict()
+        state_dict_2_after_load = model_cast.state_dict()
+
+        np.testing.assert_array_equal(
+            state_dict_1_after_load['linear.weight'].astype("float16"),
+            state_dict_2_after_load['linear.weight'],
+        )
+
+    def setup_dist_env(self):
+        fleet.init(is_collective=True)
+
+    def run_save_state_dict(self):
+        self.setup_dist_env()
+        model = SimpleMLP()
+        sharded_state_dict = model.sharded_state_dict()
+        dist.save_state_dict(sharded_state_dict, self.ckpt_path)
+
+
+if __name__ == '__main__':
+    TestLoadStateDictCastLogic().run_test()

test/flex_checkpoint/load_state_dict_transpose_cast_logic.py

@@ -0,0 +1,112 @@
+# 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 os
+import tempfile
+
+import numpy as np
+
+import paddle
+import paddle.distributed as dist
+from paddle.distributed import fleet
+from paddle.distributed.fleet.layers.mpu import (
+    ColumnParallelLinear,
+)
+from paddle.distributed.flex_checkpoint.dcp.sharded_weight import (
+    build_sharded_state_dict,
+)
+from paddle.nn import Layer
+
+
+class ColumnParallelLinearTransWeight(ColumnParallelLinear):
+    def sharded_state_dict(
+        self,
+        structured_name_prefix: str = "",
+    ):
+        state_dict = self.state_dict(structured_name_prefix="")
+        for k, v in state_dict.items():
+            if "weight" in k:
+                state_dict[k] = v.T
+        return build_sharded_state_dict(
+            state_dict, {"weight": 0}, structured_name_prefix
+        )
+
+
+class SimpleMLP(Layer):
+    def __init__(self, in_features=1024, out_features=1024):
+        super().__init__()
+        self.linear = ColumnParallelLinear(
+            in_features, out_features, has_bias=False
+        )
+
+    def forward(self, x):
+        x = self.linear(x)
+        return x
+
+
+class SimpleMLPTransCastWeight(Layer):
+    def __init__(self, in_features=1024, out_features=1024):
+        super().__init__()
+        self.linear = ColumnParallelLinearTransWeight(
+            in_features, out_features, has_bias=False
+        )
+
+    def forward(self, x):
+        x = self.linear(x)
+        return x
+
+
+class TestLoadStateDictTransposeCastLogic:
+    def __init__(self):
+        self.aoa_config = {"aoa_statements": [os.getenv("aoa_statements")]}
+        self.ckpt_path = tempfile.TemporaryDirectory().name
+        self.in_features = 1024
+        self.out_features = 2048
+
+    def run_test(self):
+        self.run_save_state_dict()
+        model = SimpleMLP()
+        model_trans_cast = SimpleMLPTransCastWeight()
+        model_trans_cast = paddle.amp.decorate(
+            models=model_trans_cast,
+            optimizers=None,
+            level="O2",
+            dtype="float16",
+        )
+        sharded_state_dict = model.sharded_state_dict()
+        sharded_state_dict_trans = model_trans_cast.sharded_state_dict()
+        dist.load_state_dict(sharded_state_dict, self.ckpt_path)
+        dist.load_state_dict(
+            sharded_state_dict_trans, self.ckpt_path, aoa_config=self.aoa_config
+        )
+        state_dict_1_after_load = model.state_dict()
+        state_dict_2_after_load = model_trans_cast.state_dict()
+
+        np.testing.assert_array_equal(
+            state_dict_1_after_load['linear.weight'].astype("float16"),
+            state_dict_2_after_load['linear.weight'],
+        )
+
+    def setup_dist_env(self):
+        fleet.init(is_collective=True)
+
+    def run_save_state_dict(self):
+        self.setup_dist_env()
+        model = SimpleMLP()
+        sharded_state_dict = model.sharded_state_dict()
+        dist.save_state_dict(sharded_state_dict, self.ckpt_path)
+
+
+if __name__ == '__main__':
+    TestLoadStateDictTransposeCastLogic().run_test()

test/flex_checkpoint/load_state_dict_transpose_logic.py

@@ -0,0 +1,113 @@
+# 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 os
+import tempfile
+
+import numpy as np
+
+import paddle.distributed as dist
+from paddle.distributed import fleet
+from paddle.distributed.fleet.layers.mpu import (
+    ColumnParallelLinear,
+)
+from paddle.distributed.flex_checkpoint.dcp.sharded_weight import (
+    build_sharded_state_dict,
+)
+from paddle.nn import Layer
+
+
+class ColumnParallelLinearTransWeight(ColumnParallelLinear):
+    def sharded_state_dict(
+        self,
+        structured_name_prefix: str = "",
+    ):
+        state_dict = self.state_dict(structured_name_prefix="")
+        for k, v in state_dict.items():
+            if "weight" in k:
+                state_dict[k] = v.T
+        return build_sharded_state_dict(
+            state_dict, {"weight": 0, "bias": 0}, structured_name_prefix
+        )
+
+
+class SimpleMLP(Layer):
+    def __init__(self, in_features=1024, out_features=1024):
+        super().__init__()
+        self.linear = ColumnParallelLinear(
+            in_features, out_features, has_bias=True
+        )
+
+    def forward(self, x):
+        x = self.linear(x)
+        return x
+
+
+class SimpleMLPTransWeight(Layer):
+    def __init__(self, in_features=1024, out_features=1024):
+        super().__init__()
+        self.linear = ColumnParallelLinearTransWeight(
+            in_features, out_features, has_bias=True
+        )
+
+    def forward(self, x):
+        x = self.linear(x)
+        return x
+
+
+class TestLoadStateDictTransposeLogic:
+    def __init__(self):
+        self.aoa_config = {"aoa_statements": [os.getenv("aoa_statements")]}
+        self.ckpt_path = tempfile.TemporaryDirectory().name
+        self.in_features = 1024
+        self.out_features = 2048
+
+    def run_test(self):
+        self.run_save_state_dict()
+        model = SimpleMLP()
+        model_trans = SimpleMLPTransWeight()
+        sharded_state_dict = model.sharded_state_dict()
+        sharded_state_dict_trans = model_trans.sharded_state_dict()
+        dist.load_state_dict(sharded_state_dict, self.ckpt_path)
+        dist.load_state_dict(
+            sharded_state_dict_trans, self.ckpt_path, aoa_config=self.aoa_config
+        )
+        state_dict_1_after_load = model.state_dict()
+        state_dict_2_after_load = model_trans.state_dict()
+
+        np.testing.assert_array_equal(
+            state_dict_1_after_load['linear.weight'],
+            state_dict_2_after_load['linear.weight'],
+        )
+
+    def setup_dist_env(self):
+        fleet.init(is_collective=True)
+
+    def run_save_state_dict(self):
+        self.setup_dist_env()
+        model = SimpleMLP()
+        sharded_state_dict = model.sharded_state_dict()
+        dist.save_state_dict(sharded_state_dict, self.ckpt_path)
+
+
+class TestLoadStateDictTransposeLogic2(TestLoadStateDictTransposeLogic):
+    def __init__(self):
+        super().__init__()
+        self.in_features = 1024
+        self.out_features = 1024
+
+
+if __name__ == '__main__':
+    TestLoadStateDictTransposeLogic().run_test()
+    TestLoadStateDictTransposeLogic2().run_test()

test/flex_checkpoint/merge_sharded_state_dict.py

@@ -0,0 +1,101 @@
+# 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 os
+
+import numpy as np
+
+import paddle.distributed as dist
+from paddle.distributed import fleet
+from paddle.distributed.fleet.layers.mpu import (
+    ColumnParallelLinear,
+)
+from paddle.nn import Layer
+
+
+class SimpleMLP(Layer):
+    def __init__(self, hidden_size=1024):
+        super().__init__()
+        self.linear = ColumnParallelLinear(
+            hidden_size, hidden_size * 2, has_bias=True
+        )
+        self.linear1 = ColumnParallelLinear(
+            hidden_size, hidden_size * 2, has_bias=True
+        )
+
+    def forward(self, x):
+        x = self.linear(x)
+        x = self.linear1(x)
+        return x
+
+
+class TestDistCheckpoint:
+    def __init__(self):
+        np.random.seed(42)
+        self.temp_dir = "./state_dict_merge"
+        self.test_type = os.getenv("test_type")
+        self.layer_type = os.getenv("layer_type")
+        self.tp_degree = int(os.getenv("tp"))
+        self.dp_degree = int(os.getenv("dp"))
+        self.world_size = int(os.getenv("world_size"))
+        self.has_bias = os.getenv("has_bias", "True").lower() == "true"
+
+        self.hidden_size = 32
+        self.vocab_size = 1024
+
+    def run_layer_test(self):
+        strategy = fleet.DistributedStrategy()
+        strategy.hybrid_configs = {
+            "dp_degree": self.dp_degree,
+            "mp_degree": self.tp_degree,
+            "pp_degree": 1,
+        }
+        fleet.init(is_collective=True, strategy=strategy)
+        hcg = fleet.get_hybrid_communicate_group()
+        tp_group = hcg.get_model_parallel_group()
+
+        model_path = os.path.join(self.temp_dir, 'model')
+        single_path = os.path.join(self.temp_dir, 'single_model')
+        model = SimpleMLP()
+        sharded_state_dict = model.sharded_state_dict()
+        state_dict = model.state_dict()
+
+        dist.save_state_dict(sharded_state_dict, model_path, safetensors=False)
+
+        dist.flex_checkpoint.dcp.load_state_dict.merge_sharded_state_dict(
+            model_path,
+            single_path,
+            offload=True,
+            safetensors=False,
+        )
+
+        import paddle.distributed
+
+        paddle.distributed.barrier()
+
+        if paddle.distributed.get_rank() == 0:
+            import safetensors
+
+            load_result = {}
+            for i in range(1, 3):
+                load_result.update(
+                    safetensors.paddle.load_file(
+                        f"{single_path}/model-0000{i}-of-00002.safetensors"
+                    )
+                )
+            assert len(load_result) == 4
+
+
+if __name__ == '__main__':
+    TestDistCheckpoint().run_layer_test()

test/flex_checkpoint/model_full_param_logic.py

@@ -0,0 +1,414 @@
+# 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 os
+
+import paddle
+from paddle import nn
+from paddle.distributed import fleet
+from paddle.distributed.fleet.meta_parallel import (
+    ColumnParallelLinear,
+    LayerDesc,
+    PipelineLayer,
+    RowParallelLinear,
+    SharedLayerDesc,
+    VocabParallelEmbedding,
+)
+from paddle.distributed.fleet.meta_parallel.sharding.group_sharded_stage3 import (
+    GroupShardedStage3,
+)
+
+
+class SimpleMLP(nn.Layer):
+    def __init__(self, hidden_size=100, has_bias=False):
+        super().__init__()
+        self.embedding = VocabParallelEmbedding(24, hidden_size)
+        self.linear1 = ColumnParallelLinear(
+            hidden_size, hidden_size, gather_output=False, has_bias=has_bias
+        )
+        self.linear2 = RowParallelLinear(
+            hidden_size, hidden_size, input_is_parallel=True, has_bias=has_bias
+        )
+        self.llm_head = self.embedding
+
+    def forward(self, x):
+        x = self.embedding(x)
+        x = self.linear1(x)
+        x = self.linear2(x)
+        x = paddle.matmul(x, self.llm_head.weight, transpose_y=True)
+        return x
+
+
+class MLP(paddle.nn.Layer):
+    def __init__(self, linear_size=1000, param_attr=None, bias_attr=None):
+        super().__init__()
+        self._linear1 = nn.Linear(linear_size, linear_size)
+        self._linear2 = nn.Linear(linear_size, linear_size)
+        self._linear3 = nn.Linear(linear_size, 10)
+
+    def forward(self, inputs):
+        y = self._linear1(inputs)
+        y = self._linear2(y)
+        y = self._linear3(y)
+        return y
+
+
+class SimpleMLPPipeline(PipelineLayer):
+    def __init__(
+        self, hcg, hidden_size=100, has_bias=False, vocab_size=24, pp_degree=2
+    ):
+        shared_embedding = SharedLayerDesc(
+            key="shared_embedding",
+            layer_func=VocabParallelEmbedding,
+            num_embeddings=vocab_size,
+            embedding_dim=hidden_size,
+        )
+
+        shared_head = SharedLayerDesc(
+            key="shared_embedding",
+            layer_func=VocabParallelEmbedding,
+            num_embeddings=vocab_size,
+            embedding_dim=hidden_size,
+            forward_func=lambda layer, x: paddle.matmul(
+                x, layer.weight, transpose_y=True
+            ),
+        )
+
+        layers = [
+            shared_embedding,
+            LayerDesc(
+                ColumnParallelLinear,
+                hidden_size,
+                hidden_size,
+                gather_output=False,
+                has_bias=has_bias,
+            ),
+            LayerDesc(
+                RowParallelLinear,
+                hidden_size,
+                hidden_size,
+                input_is_parallel=True,
+                has_bias=has_bias,
+            ),
+            shared_head,
+        ]
+
+        super().__init__(
+            layers=layers,
+            loss_fn=None,
+            seg_method="uniform",
+            topology=hcg._topo,
+        )
+
+
+class TestFullParamLogic:
+    def __init__(self):
+        self.tp_degree = int(os.getenv("tp", "1"))
+        self.dp_degree = int(os.getenv("dp", "1"))
+        self.sharding_degree = int(os.getenv("sharding_degree", "1"))
+        self.pp_degree = int(os.getenv("pp", "1"))
+        self.world_size = int(os.getenv("world_size"))
+        self.has_bias = os.getenv("has_bias", "True").lower() == "true"
+        self.batch_size = 2
+        self.hidden_size = 32
+        self.vocab_size = 24
+        self.seq_len = 2
+        self.hcg = None
+
+    def run_test(self):
+        strategy = fleet.DistributedStrategy()
+        strategy.hybrid_configs = {
+            "dp_degree": self.dp_degree,
+            "mp_degree": self.tp_degree,
+            "sharding_degree": self.sharding_degree,
+            "pp_degree": 1,
+        }
+        fleet.init(is_collective=True, strategy=strategy)
+        self.run_full_param_test()
+        self.run_full_param_with_aoa_test()
+
+    def run_full_param_test(self):
+        model = SimpleMLP(hidden_size=self.hidden_size, has_bias=self.has_bias)
+        model = fleet.distributed_model(model)
+        model.train()
+        model_state_dict = model.state_dict()
+
+        for k, v in model_state_dict.items():
+            ones = paddle.ones_like(v)
+            paddle.assign(ones, v)
+
+        full_param_iter = model.full()
+        full_param = dict(full_param_iter)
+
+        param_shape = {
+            "_layers.embedding.weight": [24, 32],
+            "_layers.linear1.weight": [32, 32],
+            "_layers.linear1.bias": [32],
+            "_layers.linear2.weight": [32, 32],
+            "_layers.linear2.bias": [32],
+            "_layers.llm_head.weight": [24, 32],
+        }
+        for name, shape in param_shape.items():
+            if not self.has_bias:
+                if ".bias" in name:
+                    continue
+            assert name in full_param.keys()
+            tensor = full_param[name]
+            answer = paddle.ones_like(tensor)
+            assert tensor._md5sum() == answer._md5sum()
+
+    def run_full_param_with_aoa_test(self):
+        model = SimpleMLP(hidden_size=self.hidden_size, has_bias=self.has_bias)
+        model = paddle.amp.decorate(
+            models=model, optimizers=None, level="O2", dtype="float16"
+        )
+        model = fleet.distributed_model(model)
+        model.train()
+        model_state_dict = model.state_dict()
+
+        for k, v in model_state_dict.items():
+            ones = paddle.ones_like(v)
+            paddle.assign(ones, v)
+            if k == "_layers.linear1.weight":
+                zeros = paddle.zeros_like(v)
+                paddle.assign(zeros, v)
+
+        aoa_config = {
+            "aoa_statements": [
+                "_layers.linear1.weight, _layers.linear2.weight -> _layers.fused_weight, axis=1"
+                "_layers.embedding.weight -> _layers.embedding.weight, dtype = 'float32'"
+            ]
+        }
+
+        full_param_iter = model.full(aoa_config)
+        full_param = dict(full_param_iter)
+
+        param_shape = {
+            # "_layers.linear1.weight" : [32,32],
+            # "_layers.linear2.weight" : [32, 32],
+            "_layers.embedding.weight": [24, 32],
+            "_layers.linear1.bias": [32],
+            "_layers.linear2.bias": [32],
+            "_layers.llm_head.weight": [24, 32],
+            "_layers.fused_weight": [32, 64],
+        }
+
+        for name, shape in param_shape.items():
+            if name == "_layers.fused_weight":
+                continue
+            if not self.has_bias:
+                if ".bias" in name:
+                    continue
+            assert name in full_param.keys()
+            tensor = full_param[name]
+            answer = paddle.ones_like(tensor)
+            assert tensor._md5sum() == answer._md5sum()
+            if name == "_layers.embedding.weight":
+                assert tensor.dtype == paddle.float32
+        assert "_layers.fused_weight" in full_param.keys()
+        ones = paddle.ones([32, 32], 'float16')
+        zeros = paddle.zeros([32, 32], 'float16')
+        answer = paddle.concat([zeros, ones], axis=1)
+        assert full_param["_layers.fused_weight"]._md5sum() == answer._md5sum()
+
+
+class TestFullParamHVGroupLogic(TestFullParamLogic):
+    def __init__(self):
+        super().__init__()
+
+    def run_test(self):
+        strategy = fleet.DistributedStrategy()
+        strategy.hybrid_configs = {
+            "dp_degree": self.dp_degree,
+            "mp_degree": self.tp_degree,
+            "sharding_degree": self.sharding_degree,
+            "pp_degree": self.pp_degree,
+        }
+        fleet.init(is_collective=True, strategy=strategy)
+        self.run_full_param_test()
+        self.run_full_param_with_aoa_test()
+        self.run_full_param_memory_growth_threshold_test()
+
+    def run_full_param_test(self, memory_growth_threshold=8 * (2**30)):
+        hcg = fleet.get_hybrid_communicate_group()
+        model = SimpleMLPPipeline(
+            hcg=hcg,
+            hidden_size=self.hidden_size,
+            has_bias=self.has_bias,
+            vocab_size=self.vocab_size,
+            pp_degree=self.pp_degree,
+        )
+        model = fleet.distributed_model(model)
+        model.train()
+
+        model_state_dict = model.state_dict()
+        tp_group = hcg.get_model_parallel_group()
+        pp_group = hcg.get_pipe_parallel_group()
+        full_param_iter = model.full(
+            h_group=tp_group,
+            v_group=pp_group,
+            memory_growth_threshold=memory_growth_threshold,
+        )
+
+        full_param = dict(full_param_iter)
+        param_shape = {
+            "_layers.shared_layers.shared_embedding.weight": [24, 32],
+            "_layers.1.weight": [32, 32],
+            "_layers.1.bias": [32],
+            "_layers.2.weight": [32, 32],
+            "_layers.2.bias": [32],
+        }
+        param_split_axis = {
+            "_layers.shared_layers.shared_embedding.weight": 0,
+            "_layers.1.weight": 1,
+            "_layers.1.bias": 0,
+            "_layers.2.weight": 0,
+            "_layers.2.bias": -1,
+        }
+        model_parallel_rank = hcg.get_model_parallel_rank()
+
+        for name, shape in param_shape.items():
+            assert name in full_param.keys()
+            assert tuple(full_param[name].shape) == tuple(shape)
+
+        for name, param in full_param.items():
+            if name not in model_state_dict:
+                continue
+            splited_axis = param_split_axis[name]
+
+            if splited_axis == -1:
+                splited = [param, param]
+            else:
+                splited = paddle.split(
+                    param, num_or_sections=2, axis=splited_axis
+                )
+            t = splited[model_parallel_rank]
+            assert t._md5sum() == model_state_dict[name]._md5sum()
+
+    def run_full_param_with_aoa_test(self, memory_growth_threshold=8 * (2**30)):
+        hcg = fleet.get_hybrid_communicate_group()
+        model = SimpleMLPPipeline(
+            hcg=hcg,
+            hidden_size=self.hidden_size,
+            has_bias=self.has_bias,
+            vocab_size=self.vocab_size,
+            pp_degree=self.pp_degree,
+        )
+        model = fleet.distributed_model(model)
+        model.train()
+
+        model_state_dict = model.state_dict()
+
+        tp_group = hcg.get_model_parallel_group()
+        pp_group = hcg.get_pipe_parallel_group()
+
+        aoa_config = {
+            "aoa_statements": [
+                "_layers.1.weight, _layers.2.weight -> _layers.fused_weight, axis=1",
+                "_layers.2.bias -> _",
+            ]
+        }
+
+        full_param_iter = model.full(
+            aoa_config,
+            h_group=tp_group,
+            v_group=pp_group,
+            memory_growth_threshold=memory_growth_threshold,
+        )
+
+        full_param = dict(full_param_iter)
+
+        param_shape = {
+            "_layers.shared_layers.shared_embedding.weight": [24, 32],
+            "_layers.1.bias": [32],
+            "_layers.fused_weight": [32, 64],
+        }
+
+        param_split_axis = {
+            "_layers.shared_layers.shared_embedding.weight": 0,
+            "_layers.1.weight": 1,
+            "_layers.1.bias": 0,
+            "_layers.2.weight": 0,
+        }
+
+        for name, shape in param_shape.items():
+            assert name in full_param.keys()
+            assert tuple(full_param[name].shape) == tuple(shape)
+
+        assert "_layers.2.bias" not in full_param
+        assert "_layers.1.weight" not in full_param
+        assert "_layers.2.weight" not in full_param
+
+        fused_weight = full_param.pop("_layers.fused_weight")
+        splited = paddle.split(fused_weight, num_or_sections=2, axis=1)
+
+        full_param["_layers.1.weight"] = splited[0]
+        full_param["_layers.2.weight"] = splited[1]
+
+        model_parallel_rank = hcg.get_model_parallel_rank()
+
+        for name, param in full_param.items():
+            if name not in model_state_dict:
+                continue
+            splited_axis = param_split_axis[name]
+
+            if splited_axis == -1:
+                splited = [param, param]
+            else:
+                splited = paddle.split(
+                    param, num_or_sections=2, axis=splited_axis
+                )
+            t = splited[model_parallel_rank]
+            assert t._md5sum() == model_state_dict[name]._md5sum()
+
+    def run_full_param_memory_growth_threshold_test(self):
+        self.run_full_param_test(memory_growth_threshold=1)
+        self.run_full_param_with_aoa_test(memory_growth_threshold=1)
+
+
+class TestFullParamLogicWithSharding3:
+    def __init__(self):
+        paddle.distributed.init_parallel_env()
+
+    def run_test(self):
+        model = MLP()
+        opt = paddle.optimizer.AdamW(parameters=model.parameters())
+        model = GroupShardedStage3(model, opt, segment_size=256)
+        model.train()
+        x = paddle.randn([4, 1000])
+        loss = model(x).mean()
+        loss.backward()
+        opt.step()
+        opt.clear_grad()
+        model.get_all_parameters(convert2cpu=True)
+        param_md5_list = []
+        for param in model.parameters():
+            param_md5_list.append(param._md5sum())
+        full_param_iter = model.full()
+        full_param = dict(full_param_iter)
+        full_param_md5_list = []
+        for name, param in full_param.items():
+            full_param_md5_list.append(param._md5sum())
+        assert param_md5_list == full_param_md5_list
+
+
+if __name__ == '__main__':
+    test_using_hv_group = int(os.getenv("test_using_hv_group", "0")) == 1
+    test_with_sharding3 = int(os.getenv("test_with_sharding3", "0")) == 1
+    if test_using_hv_group:
+        TestFullParamHVGroupLogic().run_test()
+    elif test_with_sharding3:
+        TestFullParamLogicWithSharding3().run_test()
+    else:
+        TestFullParamLogic().run_test()

test/flex_checkpoint/sharded_state_dict_logic.py

@@ -0,0 +1,531 @@
+# 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 os
+
+import paddle
+from paddle import nn
+from paddle.distributed import ShardedWeight, fleet
+from paddle.distributed.fleet.layers.mpu import (
+    ColumnParallelLinear,
+    RowParallelLinear,
+    VocabParallelEmbedding,
+)
+from paddle.distributed.fleet.meta_optimizers.dygraph_optimizer.dygraph_sharding_optimizer import (
+    DygraphShardingOptimizer,
+    DygraphShardingOptimizerV2,
+)
+from paddle.distributed.fleet.meta_parallel.sharding.group_sharded_optimizer_stage2 import (
+    GroupShardedOptimizerStage2,
+)
+from paddle.distributed.fleet.meta_parallel.sharding.group_sharded_stage3 import (
+    GroupShardedStage3,
+)
+from paddle.distributed.fleet.utils.sequence_parallel_utils import (
+    ColumnSequenceParallelLinear,
+    RowSequenceParallelLinear,
+)
+
+
+class SimpleMLP(
+    nn.Layer
+):  # embedding_weight_size=24*100=2400,it can't be divided by 256,which is using to check the padding logic
+    def __init__(self, hidden_size=100, has_bias=False):
+        super().__init__()
+        self.embedding = VocabParallelEmbedding(24, hidden_size)
+        self.linear1 = ColumnParallelLinear(
+            hidden_size, hidden_size, gather_output=False, has_bias=has_bias
+        )
+        self.linear2 = RowParallelLinear(
+            hidden_size, hidden_size, input_is_parallel=True, has_bias=has_bias
+        )
+        self.llm_head = self.embedding  # test the shared weight
+
+    def forward(self, x):
+        x = self.embedding(x)
+        x = self.linear1(x)
+        x = self.linear2(x)
+        x = paddle.matmul(x, self.llm_head.weight, transpose_y=True)
+        return x
+
+
+class TestParallelLayersLogic:
+    def __init__(self):
+        self.optimizer_var_suffix = [".moment1_0", ".moment2_0", ".w_0"]
+        self.test_type = os.getenv("test_type")
+        self.layer_type = os.getenv("layer_type")
+        self.tp_degree = int(os.getenv("tp", "1"))
+        self.dp_degree = int(os.getenv("dp", "1"))
+        self.sharding_degree = int(os.getenv("sharding_degree", "1"))
+        self.world_size = int(os.getenv("world_size"))
+        self.has_bias = os.getenv("has_bias", "True").lower() == "true"
+        self.master_weight = (
+            os.getenv("master_weight", "False").lower() == "true"
+        )
+        self.batch_size = 2
+        self.hidden_size = 32
+        self.vocab_size = 24
+        self.seq_len = 2
+        self.hcg = None
+
+    def run_test(self):
+        strategy = fleet.DistributedStrategy()
+        strategy.hybrid_configs = {
+            "dp_degree": self.dp_degree,
+            "mp_degree": self.tp_degree,
+            "sharding_degree": self.sharding_degree,
+            "pp_degree": 1,
+        }
+        fleet.init(is_collective=True, strategy=strategy)
+        self.hcg = fleet.get_hybrid_communicate_group()
+        if self.test_type == "layer":
+            self.run_layer_test()
+        elif self.test_type == "optimizer":
+            self.run_optimizer_test()
+        else:
+            raise ValueError(f"Unknown test_type: {self.test_type}")
+
+    def run_layer_test(self):
+        tp_group = self.hcg.get_model_parallel_group()
+        layer = self._get_layer()
+        sharded_dict = layer.sharded_state_dict()
+        self._verify_parallel_layer(
+            sharded_dict, tp_group.rank, tp_group.nranks
+        )
+
+    def _get_layer(self):
+        if self.layer_type == "ColumnParallelLinear":
+            return ColumnParallelLinear(
+                self.hidden_size, self.hidden_size * 2, has_bias=self.has_bias
+            )
+        elif self.layer_type == "RowParallelLinear":
+            return RowParallelLinear(
+                self.hidden_size * 2, self.hidden_size, has_bias=self.has_bias
+            )
+        elif self.layer_type == "VocabParallelEmbedding":
+            return VocabParallelEmbedding(self.vocab_size, self.hidden_size)
+        elif self.layer_type == "ColumnSequenceParallelLinear":
+            return ColumnSequenceParallelLinear(
+                self.hidden_size,
+                self.hidden_size * 2,
+                has_bias=self.has_bias,
+                gather_output=False,
+            )
+        elif self.layer_type == "RowSequenceParallelLinear":
+            return RowSequenceParallelLinear(
+                self.hidden_size * 2,
+                self.hidden_size,
+                has_bias=self.has_bias,
+                input_is_parallel=True,
+            )
+        raise ValueError(f"Unknown layer_type: {self.layer_type}")
+
+    def _verify_parallel_layer(self, sharded_dict, tp_rank, tp_world_size):
+        if self.has_bias:
+            assert 'bias' in sharded_dict
+            bias_shard = sharded_dict['bias']
+            assert isinstance(bias_shard, ShardedWeight)
+        else:
+            assert 'bias' not in sharded_dict
+
+        assert 'weight' in sharded_dict
+        weight_shard = sharded_dict['weight']
+        assert isinstance(weight_shard, ShardedWeight)
+
+        if self.layer_type == "ColumnParallelLinear":
+            in_f, out_f = self.hidden_size, self.hidden_size * 2
+            assert weight_shard.global_shape == (in_f, out_f)
+            assert weight_shard.local_shape == (in_f, out_f // tp_world_size)
+            assert weight_shard.global_offset == (
+                0,
+                tp_rank * (out_f // tp_world_size),
+            )
+            if self.has_bias:
+                assert bias_shard.global_shape == (out_f,)
+                assert bias_shard.local_shape == (out_f // tp_world_size,)
+                assert bias_shard.global_offset == (
+                    tp_rank * (out_f // tp_world_size),
+                )
+
+        elif self.layer_type == "RowParallelLinear":
+            in_f, out_f = self.hidden_size * 2, self.hidden_size
+            # Weight is sharded on axis 1
+            assert weight_shard.global_shape == (in_f, out_f)
+            assert weight_shard.local_shape == (in_f // tp_world_size, out_f)
+            assert weight_shard.global_offset == (
+                tp_rank * (in_f // tp_world_size),
+                0,
+            )
+
+            if self.has_bias:
+                # Bias is replicated, not sharded
+                assert bias_shard.global_shape == (out_f,)
+                assert bias_shard.local_shape == bias_shard.global_shape
+                assert bias_shard.global_offset == (0,)
+
+        elif self.layer_type == "VocabParallelEmbedding":
+            assert weight_shard.global_shape == (
+                self.vocab_size,
+                self.hidden_size,
+            )
+            assert weight_shard.local_shape == (
+                self.vocab_size // tp_world_size,
+                self.hidden_size,
+            )
+            assert weight_shard.global_offset == (
+                tp_rank * (self.vocab_size // tp_world_size),
+                0,
+            )
+
+        elif self.layer_type == "ColumnSequenceParallelLinear":
+            in_f, out_f = self.hidden_size, self.hidden_size * 2
+            assert weight_shard.global_shape == (in_f, out_f)
+            assert weight_shard.local_shape == (in_f, out_f // tp_world_size)
+            assert weight_shard.global_offset == (
+                0,
+                tp_rank * (out_f // tp_world_size),
+            )
+            if self.has_bias:
+                assert bias_shard.global_shape == (out_f,)
+                assert bias_shard.local_shape == (out_f // tp_world_size,)
+                assert bias_shard.global_offset == (
+                    tp_rank * (out_f // tp_world_size),
+                )
+
+        elif self.layer_type == "RowSequenceParallelLinear":
+            in_f, out_f = self.hidden_size * 2, self.hidden_size
+            assert weight_shard.global_shape == (in_f, out_f)
+            assert weight_shard.local_shape == (in_f // tp_world_size, out_f)
+            assert weight_shard.global_offset == (
+                tp_rank * (in_f // tp_world_size),
+                0,
+            )
+            if self.has_bias:
+                assert bias_shard.global_shape == (out_f,)
+                assert bias_shard.local_shape == bias_shard.global_shape
+                assert bias_shard.global_offset == (0,)
+
+    def run_optimizer_test(self):
+        model = SimpleMLP(has_bias=self.has_bias)
+        model = paddle.amp.decorate(
+            models=model, optimizers=None, level="O2", dtype="float16"
+        )
+        if self.master_weight:  # test the master_weight
+            opt = paddle.optimizer.AdamW(
+                learning_rate=0.01,
+                parameters=model.parameters(),
+                multi_precision=True,
+            )
+        else:
+            opt = paddle.optimizer.AdamW(
+                learning_rate=0.01,
+                parameters=model.parameters(),
+                multi_precision=False,
+            )
+        if self.layer_type == "AdamW":
+            model = fleet.distributed_model(model)
+            model.train()
+            x = paddle.randint(
+                low=0,
+                high=self.vocab_size,
+                shape=[self.batch_size, self.seq_len, self.hidden_size],
+                dtype='int64',
+            )
+            y = model(x).mean()
+            y.backward()
+            opt.step()
+            opt.clear_grad()
+
+            model_sharded_state_dict = model.sharded_state_dict()
+            opt_sharded_state_dict = opt.sharded_state_dict(
+                model_sharded_state_dict
+            )
+            for key, value in model_sharded_state_dict.items():
+                for state_name in self.optimizer_var_suffix:
+                    opt__var_name = key + state_name
+                    if opt__var_name in opt_sharded_state_dict:
+                        assert tuple(
+                            opt_sharded_state_dict[opt__var_name].local_shape
+                        ) == tuple(value.local_shape)
+                        assert tuple(
+                            opt_sharded_state_dict[opt__var_name].global_shape
+                        ) == tuple(value.global_shape)
+                        assert tuple(
+                            opt_sharded_state_dict[opt__var_name].global_offset
+                        ) == tuple(value.global_offset)
+        elif self.layer_type == "DygraphShardingOptimizer":
+            opt = DygraphShardingOptimizer(opt, self.hcg)
+            model.train()
+            x = paddle.randint(
+                low=0,
+                high=self.vocab_size,
+                shape=[self.batch_size, self.seq_len, self.hidden_size],
+                dtype='int64',
+            )
+            rank = paddle.distributed.get_rank()
+            sharidng_x = (
+                x[0 : self.batch_size // 2]
+                if rank == 0
+                else x[self.batch_size // 2 :]
+            )
+            y = model(sharidng_x).mean()
+            y.backward()
+            opt.step()
+            opt.clear_grad()
+
+            model_sharded_state_dict = model.sharded_state_dict()
+            opt_sharded_state_dict = opt.sharded_state_dict(
+                model_sharded_state_dict
+            )
+
+            for key, value in model_sharded_state_dict.items():
+                for state_name in self.optimizer_var_suffix:
+                    opt__var_name = key + state_name
+                    if opt__var_name in opt_sharded_state_dict:
+                        assert tuple(
+                            opt_sharded_state_dict[opt__var_name].local_shape
+                        ) == tuple(value.local_shape)
+                        assert tuple(
+                            opt_sharded_state_dict[opt__var_name].global_shape
+                        ) == tuple(value.global_shape)
+                        assert tuple(
+                            opt_sharded_state_dict[opt__var_name].global_offset
+                        ) == tuple(value.global_offset)
+        elif self.layer_type == "DygraphShardingOptimizerV2":
+            opt = DygraphShardingOptimizerV2(opt, self.hcg)
+            model.train()
+            x = paddle.randint(
+                low=0,
+                high=self.vocab_size,
+                shape=[self.batch_size, self.seq_len, self.hidden_size],
+                dtype='int64',
+            )
+            rank = paddle.distributed.get_rank()
+            sharidng_x = (
+                x[0 : self.batch_size // 2]
+                if rank == 0
+                else x[self.batch_size // 2 :]
+            )
+            y = model(sharidng_x).mean()
+            y.backward()
+            opt.step()
+            opt.clear_grad()
+
+            model_sharded_state_dict = model.sharded_state_dict()
+            opt_sharded_state_dict = opt.sharded_state_dict(
+                model_sharded_state_dict
+            )
+            for key, value in model_sharded_state_dict.items():
+                for state_name in self.optimizer_var_suffix:
+                    opt__var_name = key + state_name
+                    if opt__var_name in opt_sharded_state_dict:
+                        if opt_sharded_state_dict[
+                            opt__var_name
+                        ].flattened_range.stop - opt_sharded_state_dict[
+                            opt__var_name
+                        ].flattened_range.start != math.prod(
+                            value.local_shape
+                        ):  # check the optimizer_var which isFragment
+                            opt_var_globle_flattened_range = []
+                            paddle.distributed.all_gather_object(
+                                opt_var_globle_flattened_range,
+                                opt_sharded_state_dict[
+                                    opt__var_name
+                                ].flattened_range,
+                            )
+
+                            first_fragment = opt_var_globle_flattened_range[0]
+                            second_fragment = opt_var_globle_flattened_range[1]
+                            assert (
+                                first_fragment.stop == second_fragment.start
+                            )  # the first_flattened_range_stop == the second_flattened_range_start
+                            opt_var_globle_size_flattened = (
+                                second_fragment.stop - first_fragment.start
+                            )
+                            model_var_globle_size_flattened = math.prod(
+                                value.local_shape
+                            )
+                            assert (
+                                opt_var_globle_size_flattened
+                                == model_var_globle_size_flattened
+                            )
+
+                        assert tuple(
+                            opt_sharded_state_dict[opt__var_name].local_shape
+                        ) == tuple(value.local_shape)
+                        assert tuple(
+                            opt_sharded_state_dict[opt__var_name].global_shape
+                        ) == tuple(value.global_shape)
+                        assert tuple(
+                            opt_sharded_state_dict[opt__var_name].global_offset
+                        ) == tuple(value.global_offset)
+
+        elif self.layer_type == "GroupShardedOptimizerStage2":
+            opt = GroupShardedOptimizerStage2(
+                opt._parameter_list, opt, self.hcg.get_sharding_parallel_group()
+            )
+
+            model.train()
+            x = paddle.randint(
+                low=0,
+                high=self.vocab_size,
+                shape=[self.batch_size, self.seq_len, self.hidden_size],
+                dtype='int64',
+            )
+            rank = paddle.distributed.get_rank()
+            sharidng_x = (
+                x[0 : self.batch_size // 2]
+                if rank == 0
+                else x[self.batch_size // 2 :]
+            )
+            y = model(sharidng_x).mean()
+            y.backward()
+            opt.step()
+            opt.clear_grad()
+
+            model_sharded_state_dict = model.sharded_state_dict()
+            opt_sharded_state_dict = opt.sharded_state_dict(
+                model_sharded_state_dict
+            )
+
+            for key, value in model_sharded_state_dict.items():
+                for state_name in self.optimizer_var_suffix:
+                    opt__var_name = key + state_name
+                    if opt__var_name in opt_sharded_state_dict:
+                        assert tuple(
+                            opt_sharded_state_dict[opt__var_name].local_shape
+                        ) == tuple(value.local_shape)
+                        assert tuple(
+                            opt_sharded_state_dict[opt__var_name].global_shape
+                        ) == tuple(value.global_shape)
+                        assert tuple(
+                            opt_sharded_state_dict[opt__var_name].global_offset
+                        ) == tuple(value.global_offset)
+
+        elif self.layer_type == "GroupShardedStage3":
+            model = fleet.distributed_model(model)
+            wrapped_model = GroupShardedStage3(
+                model, opt, segment_size=2**12
+            )  # slice the linear1、linear2 weight
+            for param in opt._parameter_list:
+                if hasattr(param, "fw_storage"):
+                    assert len(param.shape) != 1
+            wrapped_model.init_optimizer_for_slice_param()
+            for param in opt._parameter_list:
+                if hasattr(param, "fw_storage"):
+                    assert len(param.shape) == 1
+            model_sharded_state_dict = model.sharded_state_dict()
+            for k, v in model_sharded_state_dict.items():
+                if (
+                    k == "_layers.linear1.weight"
+                    or k == "_layers.linear2.weight"
+                ):
+                    assert not v.local_tensor._is_initialized()
+            wrapped_model.init_slice_param()
+            for k, v in model_sharded_state_dict.items():
+                if (
+                    k == "_layers.linear1.weight"
+                    or k == "_layers.linear2.weight"
+                ):
+                    assert v.local_tensor._is_initialized()
+            wrapped_model.align_param_to_buffer_and_clear_slice_param()
+            for k, v in model_sharded_state_dict.items():
+                if (
+                    k == "_layers.linear1.weight"
+                    or k == "_layers.linear2.weight"
+                ):
+                    assert not v.local_tensor._is_initialized()
+            model.train()
+            x = paddle.randint(
+                low=0,
+                high=self.vocab_size,
+                shape=[self.batch_size, self.seq_len, self.hidden_size],
+                dtype='int64',
+            )
+            rank = paddle.distributed.get_rank()
+            sharidng_x = (
+                x[0 : self.batch_size // 2]
+                if rank == 0
+                else x[self.batch_size // 2 :]
+            )
+            y = model(sharidng_x).mean()
+            y.backward()
+            opt.step()
+            opt.clear_grad()
+            model_sharded_state_dict = model.sharded_state_dict()
+            for k, v in model_sharded_state_dict.items():
+                if (
+                    k == "_layers.linear1.weight"
+                    or k == "_layers.linear2.weight"
+                ):
+                    assert not v.local_tensor._is_initialized()
+            wrapped_model.get_all_parameters()
+            opt_sharded_state_dict = opt.sharded_state_dict(
+                model_sharded_state_dict
+            )
+
+            for k, v in model_sharded_state_dict.items():
+                if (
+                    k == "_layers.linear1.weight"
+                    or k == "_layers.linear2.weight"
+                ):
+                    assert v.local_tensor._is_initialized()
+
+            for key, value in model_sharded_state_dict.items():
+                for state_name in self.optimizer_var_suffix:
+                    opt__var_name = key + state_name
+                    if opt__var_name in opt_sharded_state_dict:
+                        if hasattr(
+                            value.local_tensor, "fw_storage"
+                        ):  # check the optimizer_var which isFragment
+                            opt_var_globle_flattened_range = []
+                            paddle.distributed.all_gather_object(
+                                opt_var_globle_flattened_range,
+                                opt_sharded_state_dict[
+                                    opt__var_name
+                                ].flattened_range,
+                            )
+
+                            first_fragment = opt_var_globle_flattened_range[0]
+                            second_fragment = opt_var_globle_flattened_range[1]
+                            assert (
+                                first_fragment.stop == second_fragment.start
+                            )  # the first_flattened_range_stop == the second_flattened_range_start
+                            opt_var_globle_size_flattened = (
+                                second_fragment.stop - first_fragment.start
+                            )
+                            model_var_globle_size_flattened = math.prod(
+                                value.local_shape
+                            )
+                            assert (
+                                opt_var_globle_size_flattened
+                                == model_var_globle_size_flattened
+                            )
+
+                        assert tuple(
+                            opt_sharded_state_dict[opt__var_name].local_shape
+                        ) == tuple(value.local_shape)
+                        assert tuple(
+                            opt_sharded_state_dict[opt__var_name].global_shape
+                        ) == tuple(value.global_shape)
+                        assert tuple(
+                            opt_sharded_state_dict[opt__var_name].global_offset
+                        ) == tuple(value.global_offset)
+        else:
+            raise ValueError(f"Unknown layer_type: {self.layer_type}")
+
+
+if __name__ == '__main__':
+    TestParallelLayersLogic().run_test()

test/flex_checkpoint/strategy_conversion_engine.py

@@ -0,0 +1,294 @@
+# 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.
+
+# strategy_conversion_engine.py
+import argparse
+import hashlib
+
+import paddle
+import paddle.distributed as dist
+from paddle import nn
+from paddle.distributed import fleet
+from paddle.distributed.fleet.layers.mpu import (
+    ColumnParallelLinear,
+    RowParallelLinear,
+)
+
+# ==============================================================================
+# 1. Model Definitions
+# A model zoo with simple models supporting different parallelism strategies.
+# ==============================================================================
+
+
+class MLPBlock(nn.Layer):
+    """
+    A basic building block compatible with Tensor Parallelism,
+    mimicking a transformer's FFN layer.
+    """
+
+    def __init__(self, hidden_size=32):
+        super().__init__()
+        self.linear1 = ColumnParallelLinear(
+            hidden_size, hidden_size * 4, has_bias=True, gather_output=False
+        )
+        self.relu = nn.ReLU()
+        self.linear2 = RowParallelLinear(
+            hidden_size * 4, hidden_size, has_bias=True, input_is_parallel=True
+        )
+
+    def forward(self, x):
+        return self.linear2(self.relu(self.linear1(x)))
+
+
+class UnifiedMLP(nn.Sequential):
+    """
+    A unified model composed of multiple MLPBlocks.
+    This sequential structure is suitable for all parallelism types:
+    - TP is handled inside each MLPBlock.
+    - PP wraps this entire Sequential model.
+    - DP/EP treats this entire Sequential model as a single unit.
+    """
+
+    def __init__(self, hidden_size=32, num_blocks=4):
+        super().__init__(*[MLPBlock(hidden_size) for _ in range(num_blocks)])
+
+
+class Top1Router(nn.Layer):
+    """A simple Top-1 Gating network for MoE."""
+
+    def __init__(self, d_model, num_experts):
+        super().__init__()
+        self.gate = nn.Linear(d_model, num_experts)
+
+    def forward(self, x):
+        gate_logits = self.gate(x)
+        expert_weights, expert_indices = paddle.topk(gate_logits, k=1, axis=-1)
+        return nn.functional.softmax(expert_weights, axis=-1), expert_indices
+
+
+class MoELayer(nn.Layer):
+    """
+    A more robust MoE layer that handles both EP > 1 (distributed)
+    and EP = 1 (local) scenarios.
+    """
+
+    def __init__(self, d_model, num_experts, num_blocks=2, moe_group=None):
+        super().__init__()
+        self.d_model = d_model
+        self.num_experts = num_experts
+        self.moe_group = moe_group
+        self.ep_world_size = moe_group.nranks if moe_group else 1
+
+        self.router = Top1Router(d_model, num_experts)
+        self.experts = nn.LayerList(
+            [UnifiedMLP(d_model, num_blocks) for _ in range(self.num_experts)]
+        )
+
+    def forward(self, x):
+        original_shape = x.shape
+        x = x.reshape([-1, self.d_model])
+        expert_weights, expert_indices = self.router(x)
+        final_output = paddle.zeros_like(x)
+
+        if self.ep_world_size > 1:
+            # Simplified distributed routing for testing purposes.
+            ep_rank = dist.get_rank(self.moe_group)
+            for i in range(self.num_experts):
+                if i % self.ep_world_size == ep_rank:
+                    mask = (expert_indices == i).astype('float32')
+                    expert_output = self.experts[i](x)
+                    final_output += expert_output * mask
+        else:
+            # Local routing for EP = 1
+            for i in range(self.num_experts):
+                token_mask = (expert_indices == i).squeeze(-1)
+                if not token_mask.any():
+                    continue
+                selected_tokens = x[token_mask]
+                selected_weights = expert_weights[token_mask]
+                expert_output = self.experts[i](selected_tokens)
+                indices_to_scatter = paddle.where(token_mask)[0]
+                final_output = paddle.scatter(
+                    final_output,
+                    indices_to_scatter,
+                    expert_output * selected_weights,
+                    overwrite=False,
+                )
+
+        return final_output.reshape(original_shape)
+
+
+# ==============================================================================
+# 2. Core Logic (Environment Setup, Execution, and Verification)
+# ==============================================================================
+
+
+def get_model_and_strategy(args, hcg):
+    """Builds model and DistributedStrategy based on parsed arguments."""
+    strategy = fleet.DistributedStrategy()
+    strategy.hybrid_configs = {
+        "dp_degree": args.dp,
+        "mp_degree": args.tp,
+        "pp_degree": args.pp,
+    }
+
+    if args.model_type == "moe":
+        model = MoELayer(d_model=32, num_experts=4)
+    else:
+        model = UnifiedMLP()
+
+    if args.ep > 1:
+        model = MoELayer(
+            d_model=32, num_experts=4, moe_group=hcg.get_data_parallel_group()
+        )
+        strategy.hybrid_configs["ep_degree"] = args.ep
+    elif args.pp > 1:
+        # For PP, the model must be wrapped by PipelineLayer
+        model = fleet.meta_parallel.PipelineLayer(
+            layers=model, num_stages=args.pp, topology=hcg.topology()
+        )
+
+    return model, strategy
+
+
+def setup_execution_environment(config_args):
+    """A unified function to initialize Fleet and the model."""
+    strategy = fleet.DistributedStrategy()
+    strategy.hybrid_configs = {
+        "dp_degree": config_args.dp,
+        "mp_degree": config_args.tp,
+        "pp_degree": config_args.pp,
+    }
+
+    fleet.init(is_collective=True, strategy=strategy)
+    hcg = fleet.get_hybrid_communicate_group()
+
+    model, strategy = get_model_and_strategy(config_args, hcg)
+
+    # Re-initialize with the final strategy (in case ep_degree was added)
+    fleet.init(is_collective=True, strategy=strategy)
+
+    return model
+
+
+def verify_by_md5(sd1, sd2):
+    """Compares two state_dicts by the MD5 hash of each parameter."""
+
+    def get_tensor_md5(tensor):
+        return hashlib.md5(tensor.numpy().tobytes()).hexdigest()
+
+    assert sd1.keys() == sd2.keys(), (
+        f"State dicts have different keys! Got {sd1.keys()} vs {sd2.keys()}"
+    )
+    for key in sd1.keys():
+        md5_1 = get_tensor_md5(sd1[key])
+        md5_2 = get_tensor_md5(sd2[key])
+        assert md5_1 == md5_2, (
+            f"MD5 mismatch for param '{key}': baseline={md5_1} vs roundtrip={md5_2}"
+        )
+
+
+def run_step1_save_source(args):
+    """Step 1: In the source configuration, save a distributed checkpoint."""
+    model = setup_execution_environment(args.src)
+    dist.save_state_dict(model.sharded_state_dict(), args.src_ckpt_path)
+
+
+def run_step2_convert(args):
+    """Step 2: In the target configuration, load the source checkpoint and resave."""
+    model = setup_execution_environment(args.tgt)
+    dist.load_state_dict(model.sharded_state_dict(), args.src_ckpt_path)
+    dist.save_state_dict(model.sharded_state_dict(), args.tgt_ckpt_path)
+
+
+def run_step3_verify(args):
+    """Step 3: In the source configuration, load both checkpoints and compare them."""
+    # 1. Create the "round-trip" model by loading the target checkpoint
+    model_roundtrip = setup_execution_environment(args.src)
+    dist.load_state_dict(
+        model_roundtrip.sharded_state_dict(), args.tgt_ckpt_path
+    )
+
+    # 2. Create the "baseline" model by loading the original source checkpoint
+    model_baseline = setup_execution_environment(args.src)
+    dist.load_state_dict(
+        model_baseline.sharded_state_dict(), args.src_ckpt_path
+    )
+
+    dist.barrier()
+
+    # 3. Each rank verifies its own part of the state_dict.
+    # This works for all strategies, including Pipeline Parallelism.
+    final_sd = model_roundtrip.state_dict()
+    initial_sd = model_baseline.state_dict()
+
+    if final_sd and initial_sd:
+        verify_by_md5(initial_sd, final_sd)
+
+
+# ==============================================================================
+# 3. Main Entry Point
+# ==============================================================================
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--step",
+        type=str,
+        required=True,
+        choices=["save_source", "convert", "verify"],
+    )
+    parser.add_argument("--src_ckpt_path", type=str)
+    parser.add_argument("--tgt_ckpt_path", type=str)
+    parser.add_argument(
+        "--model_type",
+        default="mlp",
+        choices=["mlp", "moe"],
+        help="Model architecture.",
+    )
+
+    # Add all strategy parameters dynamically for source and target
+    for prefix in ["src", "tgt"]:
+        for p in ["world_size", "tp", "dp", "pp", "ep"]:
+            parser.add_argument(f"--{prefix}_{p}", type=int, default=0)
+
+    args = parser.parse_args()
+
+    # Reorganize parsed args into src/tgt namespaces
+    def organize_args(prefix):
+        config = {
+            p: getattr(args, f"{prefix}_{p}")
+            for p in ["world_size", "tp", "dp", "pp", "ep"]
+        }
+        config["model_type"] = args.model_type
+        # Default parallelism degree to 1 if not specified
+        if config["tp"] == 0:
+            config["tp"] = 1
+        if config["dp"] == 0:
+            config["dp"] = 1
+        if config["pp"] == 0:
+            config["pp"] = 1
+        if config["ep"] == 0:
+            config["ep"] = 1
+        return argparse.Namespace(**config)
+
+    args.src = organize_args("src")
+    args.tgt = organize_args("tgt")
+
+    # Execute the requested step
+    engine = {
+        "save_source": run_step1_save_source,
+        "convert": run_step2_convert,
+        "verify": run_step3_verify,
+    }
+    engine[args.step](args)

test/flex_checkpoint/test_AOATraceback.py

@@ -0,0 +1,103 @@
+# 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 unittest
+
+from paddle.distributed.flex_checkpoint.aoa.aoa_engine import AOAEngine
+from paddle.distributed.flex_checkpoint.dcp.sharded_weight import (
+    ShardedWeightDesc,
+)
+
+
+def create_shard_info(keys, shape=(4, 4), dtype="float32"):
+    info = {}
+    for k in keys:
+        desc = ShardedWeightDesc(
+            key=k,
+            local_shape=shape,
+            global_shape=shape,
+            global_offset=(0, 0),
+            dtype=dtype,
+        )
+        info[k] = [desc]
+    return info
+
+
+class TestMacroLayerOffsetError(unittest.TestCase):
+    def setUp(self):
+        self.source_keys = [f"model.layers.{i}.weight" for i in range(10)]
+        self.dest_keys = [f"model.layers.{i}.weight_out" for i in range(10)] + [
+            f"model.layers.{i}.weight_out2" for i in range(10)
+        ]
+
+        self.src_info = create_shard_info(self.source_keys)
+        self.dst_info = create_shard_info(self.dest_keys)
+
+    def test_macro_error_chain(self):
+        """
+        The statement contains fused_qkv_old and is missing a comma, expecting to trigger the assertion and print the chain.
+        """
+        aoa_config = {
+            "aoa_statements": [
+                "model.layers.$LAYER_ID.weight^T -> model.layers.$LAYER_ID.weight_out, axis=0 fused_qkv_old, num_heads=20,num_key_value_groups=4",
+            ],
+            "enable_traceback": True,
+        }
+
+        with self.assertRaises(AssertionError):
+            AOAEngine(
+                aoa_config=aoa_config,
+                source_state_shard_info=self.src_info,
+                destination_state_shard_info=self.dst_info,
+            )
+
+    def test_no_error_should_be_raised(self):
+        # No error should be raised
+        source_keys = ["model.layers.0.weight"]
+        dest_keys = ["model.layers.0.weight_out"]
+        src_info = create_shard_info(source_keys)
+        dst_info = create_shard_info(dest_keys)
+        aoa_config = {
+            "aoa_statements": [
+                "model.layers.0.weight^T -> model.layers.0.weight_out",
+            ],
+            "enable_traceback": True,
+        }
+        AOAEngine(
+            aoa_config=aoa_config,
+            source_state_shard_info=src_info,
+            destination_state_shard_info=dst_info,
+        )
+
+    def test_shape_propagation_error_chain(self):
+        """
+        when split/concat, only support one attr named `axis`, but got multiple attrs.
+        """
+        aoa_config = {
+            "aoa_statements": [
+                "model.layers.0.weight -> model.layers.0.weight_out,model.layers.0.weight_out2,axis=0,axis=1",
+            ],
+            "enable_traceback": True,
+        }
+
+        with self.assertRaises(ValueError):
+            AOAEngine(
+                aoa_config=aoa_config,
+                source_state_shard_info=self.src_info,
+                destination_state_shard_info=self.dst_info,
+            )
+
+
+if __name__ == "__main__":
+    unittest.main()

test/flex_checkpoint/test_aoa_engine.py

@@ -0,0 +1,741 @@
+# 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 unittest
+
+from paddle.distributed.flex_checkpoint.aoa.aoa_engine import (
+    AOAEngine,
+    ShardedWeightDesc,
+    ShardMappingEntry,
+)
+
+
+class TestAOAEngine(unittest.TestCase):
+    def test_aoa_spilt_merge(self):
+        # ------------------------------------------------------
+        # 1. Define source tensor shards (s0 and s1).
+        # Each is a (2,2) tensor, fully covering its global shape.
+        #
+        #  s0 (2,2):         s1 (2,2):
+        #  +----+----+       +----+----+
+        #  |    |    |       |    |    |
+        #  +----+----+       +----+----+
+        #  |    |    |       |    |    |
+        #  +----+----+       +----+----+
+        s0 = ShardedWeightDesc(
+            key="s0",
+            local_shape=(2, 2),
+            global_shape=(2, 2),
+            global_offset=(0, 0),
+        )
+        s1 = ShardedWeightDesc(
+            key="s1",
+            local_shape=(2, 2),
+            global_shape=(2, 2),
+            global_offset=(0, 0),
+        )
+
+        # ------------------------------------------------------
+        # 2. Define destination tensor shards (d0 and d1).
+        # Both are (1,4) tensors, i.e., a single row with 4 columns.
+        #
+        #  d0 (1,4):       d1 (1,4):
+        #  +--+--+--+--+   +--+--+--+--+
+        #  |  |  |  |  |   |  |  |  |  |
+        #  +--+--+--+--+   +--+--+--+--+
+        d0 = ShardedWeightDesc(
+            key="d0",
+            local_shape=(1, 4),
+            global_shape=(1, 4),
+            global_offset=(0, 0),
+        )
+        d1 = ShardedWeightDesc(
+            key="d1",
+            local_shape=(1, 4),
+            global_shape=(1, 4),
+            global_offset=(0, 0),
+        )
+
+        # ------------------------------------------------------
+        # 3. Record the shard info for sources and destinations
+        source_state_shard_info = {
+            "s0": [s0],
+            "s1": [s1],
+        }
+        destination_state_shard_info = {
+            "d0": [d0],
+            "d1": [d1],
+        }
+
+        # ------------------------------------------------------
+        # 4. AOA statements define axis mapping for concatenation and splitting:
+        #    - "s" is formed by concatenating s0 and s1 along axis 1 (columns).
+        #    - d0 and d1 are obtained by splitting "s" along axis 0 (rows).
+        aoa_statements = [
+            "s0, s1 -> s, axis = 1 \n",
+            "s -> d0, d1, axis = 0 \n",
+        ]
+
+        # ------------------------------------------------------
+        # 5. Create the AOAEngine with this configuration
+        aoa_engine = AOAEngine(
+            aoa_config={"aoa_statements": aoa_statements},
+            source_state_shard_info=source_state_shard_info,
+            destination_state_shard_info=destination_state_shard_info,
+        )
+
+        queries = []
+        answers = []
+
+        # ======================================================
+        # Query 1: Find source for the first half of d0 (columns 0-1)
+        # d0 shard: key="d0", local_shape=(1,2), global_shape=(1,4), global_offset=(0,0)
+        # Covers d0[:, 0:2]
+        #
+        #  d0 (1,4):
+        #  +------+------+------+------+
+        #  |(0,0) |(0,1) |      |      |
+        #  +------+------+------+------+
+        #
+        # This region is mapped from s0, row 0, columns 0-1
+        query = ShardedWeightDesc(
+            key="d0",
+            local_shape=(1, 2),
+            global_shape=(1, 4),
+            global_offset=(0, 0),
+        )
+        src_sharded_weight_desc = ShardedWeightDesc(
+            key="s0",
+            local_shape=(1, 2),
+            global_shape=(2, 2),
+            global_offset=(0, 0),
+        )
+        shard_mapping_entry = ShardMappingEntry(
+            target_slice=query,
+            source_slice=src_sharded_weight_desc,
+            postprocess_list=None,
+        )
+        answer = [shard_mapping_entry]
+        queries.append(query)
+        answers.append(answer)
+
+        # ======================================================
+        # Query 2: Find source for the second half of d1 (columns 2-3)
+        # d1 shard: key="d1", local_shape=(1,2), global_shape=(1,4), global_offset=(0,2)
+        # Covers d1[:, 2:4]
+        #
+        #  d1 (1,4):
+        #  +------+------+------+------+
+        #  |      |      |(0,2)|(0,3)|
+        #  +------+------+------+------+
+        #
+        # This region is mapped from s1, row 1, columns 0-1
+        query = ShardedWeightDesc(
+            key="d1",
+            local_shape=(1, 2),
+            global_shape=(1, 4),
+            global_offset=(0, 2),
+        )
+        src_sharded_weight_desc = ShardedWeightDesc(
+            key="s1",
+            local_shape=(1, 2),
+            global_shape=(2, 2),
+            global_offset=(1, 0),
+        )
+        shard_mapping_entry = ShardMappingEntry(
+            target_slice=query,
+            source_slice=src_sharded_weight_desc,
+            postprocess_list=None,
+        )
+        answer = [shard_mapping_entry]
+        queries.append(query)
+        answers.append(answer)
+
+        # ======================================================
+        # Query 3: Find sources for the entire d1 (full row)
+        # d1 shard: key="d1", local_shape=(1,4), global_shape=(1,4), global_offset=(0,0)
+        # Layout: covers all columns
+        #
+        #  d1 (1,4):
+        #  +------+------+------+------+
+        #  | s0   | s0   | s1   | s1   |
+        #  |(0,0) |(0,1) |(0,2) |(0,3) |
+        #  +------+------+------+------+
+        # The first two columns come from s0, the last two from s1.
+        #
+        # Source slices:
+        #  s0, local_shape=(1,2), global_shape=(2,2), global_offset=(1,0)
+        #      +----+----+
+        #      |(1,0)|(1,1)|   <- used for d1 (0,0)-(0,1)
+        #      +----+----+
+        #
+        #  s1, local_shape=(1,2), global_shape=(2,2), global_offset=(1,0)
+        #      +----+----+
+        #      |(1,0)|(1,1)|   <- used for d1 (0,2)-(0,3)
+        #      +----+----+
+        #
+        # The answer consists of two mapping entries:
+        # 1. d1[:, 0:2] <-- s0[1, :]
+        # 2. d1[:, 2:4] <-- s1[1, :]
+        query = ShardedWeightDesc(
+            key="d1",
+            local_shape=(1, 4),
+            global_shape=(1, 4),
+            global_offset=(0, 0),
+        )
+
+        # d1[:, 0:2] <--- s0[1, :]
+        src_sharded_weight_desc0 = ShardedWeightDesc(
+            key="s0",
+            local_shape=(1, 2),
+            global_shape=(2, 2),
+            global_offset=(1, 0),  # row 1, columns 0:2
+        )
+        dst_sharded_weight_desc0 = ShardedWeightDesc(
+            key="d1",
+            local_shape=(1, 2),
+            global_shape=(1, 4),
+            global_offset=(0, 0),
+        )
+        # Visual mapping:
+        # d1 (0,0)-(0,1) <--- s0 (1,0)-(1,1)
+        #  +------+------+------+------+
+        #  |==s0==|==s0==|      |      |
+        #  +------+------+------+------+
+        src_sharded_weight_desc1 = ShardedWeightDesc(
+            key="s1",
+            local_shape=(1, 2),
+            global_shape=(2, 2),
+            global_offset=(1, 0),
+        )
+        dst_sharded_weight_desc1 = ShardedWeightDesc(
+            key="d1",
+            local_shape=(1, 2),
+            global_shape=(1, 4),
+            global_offset=(0, 2),
+        )
+        # Visual mapping:
+        # d1 (0,2)-(0,3) <--- s1 (1,0)-(1,1)
+        #  +------+------+------+------+
+        #  |      |      |==s1==|==s1==|
+        #  +------+------+------+------+
+
+        shard_mapping_entry0 = ShardMappingEntry(
+            target_slice=dst_sharded_weight_desc0,
+            source_slice=src_sharded_weight_desc0,
+            postprocess_list=None,
+        )
+        shard_mapping_entry1 = ShardMappingEntry(
+            target_slice=dst_sharded_weight_desc1,
+            source_slice=src_sharded_weight_desc1,
+            postprocess_list=None,
+        )
+        answer = [shard_mapping_entry0, shard_mapping_entry1]
+        queries.append(query)
+        answers.append(answer)
+        # Visual answer summary:
+        # d1 (row 0):
+        #  +------+------+------+------+
+        #  |==s0==|==s0==|==s1==|==s1==|
+        #  +------+------+------+------+
+        #   ^      ^      ^      ^
+        #   |      |      |      |
+        #   |______|      |______|
+        #    from s0       from s1
+
+        # ------------------------------------------------------
+
+        # ======================================================
+        # Query 4: for optimizer state
+        query = ShardedWeightDesc(
+            key="d1.moment1_0",
+            local_shape=(1, 4),
+            global_shape=(1, 4),
+            global_offset=(0, 0),
+        )
+
+        # d1[:, 0:2] <--- s0[1, :]
+        src_sharded_weight_desc0 = ShardedWeightDesc(
+            key="s0.moment1_0",
+            local_shape=(1, 2),
+            global_shape=(2, 2),
+            global_offset=(1, 0),  # row 1, columns 0:2
+        )
+        dst_sharded_weight_desc0 = ShardedWeightDesc(
+            key="d1.moment1_0",
+            local_shape=(1, 2),
+            global_shape=(1, 4),
+            global_offset=(0, 0),
+        )
+
+        src_sharded_weight_desc1 = ShardedWeightDesc(
+            key="s1.moment1_0",
+            local_shape=(1, 2),
+            global_shape=(2, 2),
+            global_offset=(1, 0),
+        )
+        dst_sharded_weight_desc1 = ShardedWeightDesc(
+            key="d1.moment1_0",
+            local_shape=(1, 2),
+            global_shape=(1, 4),
+            global_offset=(0, 2),
+        )
+
+        shard_mapping_entry0 = ShardMappingEntry(
+            target_slice=dst_sharded_weight_desc0,
+            source_slice=src_sharded_weight_desc0,
+            postprocess_list=None,
+        )
+        shard_mapping_entry1 = ShardMappingEntry(
+            target_slice=dst_sharded_weight_desc1,
+            source_slice=src_sharded_weight_desc1,
+            postprocess_list=None,
+        )
+
+        answer = [shard_mapping_entry0, shard_mapping_entry1]
+        queries.append(query)
+        answers.append(answer)
+
+        # ======================================================
+        # Query 5: for optimizer state
+        query = ShardedWeightDesc(
+            key="d1.w_0",
+            local_shape=(1, 4),
+            global_shape=(1, 4),
+            global_offset=(0, 0),
+        )
+
+        # d1[:, 0:2] <--- s0[1, :]
+        src_sharded_weight_desc0 = ShardedWeightDesc(
+            key="s0.w_0",
+            local_shape=(1, 2),
+            global_shape=(2, 2),
+            global_offset=(1, 0),  # row 1, columns 0:2
+        )
+        dst_sharded_weight_desc0 = ShardedWeightDesc(
+            key="d1.w_0",
+            local_shape=(1, 2),
+            global_shape=(1, 4),
+            global_offset=(0, 0),
+        )
+
+        src_sharded_weight_desc1 = ShardedWeightDesc(
+            key="s1.w_0",
+            local_shape=(1, 2),
+            global_shape=(2, 2),
+            global_offset=(1, 0),
+        )
+        dst_sharded_weight_desc1 = ShardedWeightDesc(
+            key="d1.w_0",
+            local_shape=(1, 2),
+            global_shape=(1, 4),
+            global_offset=(0, 2),
+        )
+
+        shard_mapping_entry0 = ShardMappingEntry(
+            target_slice=dst_sharded_weight_desc0,
+            source_slice=src_sharded_weight_desc0,
+            postprocess_list=None,
+        )
+        shard_mapping_entry1 = ShardMappingEntry(
+            target_slice=dst_sharded_weight_desc1,
+            source_slice=src_sharded_weight_desc1,
+            postprocess_list=None,
+        )
+
+        answer = [shard_mapping_entry0, shard_mapping_entry1]
+        queries.append(query)
+        answers.append(answer)
+
+        # 6. Run the queries and check results
+        for idx in range(len(queries)):
+            query = queries[idx]
+            answer = answers[idx]
+            result = aoa_engine.find_shard_sources(query)
+            self.assertEqual(result, answer)
+
+    def test_aoa_cast(self):
+        """Test AOA cast primitive for dtype conversion."""
+
+        s0 = ShardedWeightDesc(
+            key="s0",
+            local_shape=(2, 2),
+            global_shape=(2, 2),
+            global_offset=(0, 0),
+            dtype="int32",
+        )
+
+        d0 = ShardedWeightDesc(
+            key="d0",
+            local_shape=(2, 2),
+            global_shape=(2, 2),
+            global_offset=(0, 0),
+            dtype="float32",
+        )
+
+        source_state_shard_info = {
+            "s0": [s0],
+        }
+        destination_state_shard_info = {
+            "d0": [d0],
+        }
+
+        aoa_statements = [
+            's0 -> d0, dtype="float32" \n',
+        ]
+
+        aoa_engine = AOAEngine(
+            aoa_config={"aoa_statements": aoa_statements},
+            source_state_shard_info=source_state_shard_info,
+            destination_state_shard_info=destination_state_shard_info,
+        )
+
+        query = ShardedWeightDesc(
+            key="d0",
+            local_shape=(2, 2),
+            global_shape=(2, 2),
+            global_offset=(0, 0),
+            dtype="float32",
+        )
+        src_sharded_weight_desc = ShardedWeightDesc(
+            key="s0",
+            local_shape=(2, 2),
+            global_shape=(2, 2),
+            global_offset=(0, 0),
+            dtype="int32",
+        )
+        shard_mapping_entry = ShardMappingEntry(
+            target_slice=query,
+            source_slice=src_sharded_weight_desc,
+            postprocess_list=['float32'],
+        )
+        answer = [shard_mapping_entry]
+
+        result = aoa_engine.find_shard_sources(query)
+        self.assertEqual(result, answer)
+
+    def test_aoa_add(self):
+        """Test AOA add primitive for adding new keys that don't exist in source."""
+
+        d0 = ShardedWeightDesc(
+            key="d0",
+            local_shape=(2, 2),
+            global_shape=(2, 2),
+            global_offset=(0, 0),
+            dtype="float32",
+        )
+
+        source_state_shard_info = {}
+
+        destination_state_shard_info = {
+            "d0": [d0],
+        }
+
+        aoa_statements = [
+            "_ -> d0 \n",
+        ]
+
+        aoa_engine = AOAEngine(
+            aoa_config={"aoa_statements": aoa_statements},
+            source_state_shard_info=source_state_shard_info,
+            destination_state_shard_info=destination_state_shard_info,
+        )
+
+        query = ShardedWeightDesc(
+            key="d0",
+            local_shape=(2, 2),
+            global_shape=(2, 2),
+            global_offset=(0, 0),
+            dtype="float32",
+        )
+
+        answer = []
+
+        result = aoa_engine.find_shard_sources(query)
+        self.assertEqual(result, answer)
+
+    def test_mixed_aoa_statements(self):
+        # test fused_ffn and transposed,rename,test_get_var_mapping_chain_macro
+        s0 = ShardedWeightDesc(
+            key="layers.0.gate_up_fused_proj.weight",
+            local_shape=(2, 4),
+            global_shape=(2, 8),
+            global_offset=(0, 0),
+            dtype="float32",
+        )
+
+        s1 = ShardedWeightDesc(
+            key="layers.0.gate_up_fused_proj.weight",
+            local_shape=(2, 4),
+            global_shape=(2, 8),
+            global_offset=(0, 4),
+            dtype="float32",
+        )
+
+        d0 = ShardedWeightDesc(
+            key="new_name_layers.0.gate_up_fused_proj.weight",
+            local_shape=(2, 2),
+            global_shape=(2, 8),
+            global_offset=(0, 0),
+            dtype="float32",
+        )
+
+        d1 = ShardedWeightDesc(
+            key="new_name_layers.0.gate_up_fused_proj.weight",
+            local_shape=(2, 2),
+            global_shape=(2, 8),
+            global_offset=(0, 2),
+            dtype="float32",
+        )
+
+        d2 = ShardedWeightDesc(
+            key="new_name_layers.0.gate_up_fused_proj.weight",
+            local_shape=(2, 2),
+            global_shape=(2, 8),
+            global_offset=(0, 4),
+            dtype="float32",
+        )
+
+        d3 = ShardedWeightDesc(
+            key="new_name_layers.0.gate_up_fused_proj.weight",
+            local_shape=(2, 2),
+            global_shape=(2, 8),
+            global_offset=(0, 6),
+            dtype="float32",
+        )
+
+        source_state_shard_info = {
+            "layers.0.gate_up_fused_proj.weight": [s0, s1],
+        }
+        destination_state_shard_info = {
+            "new_name_layers.0.gate_up_fused_proj.weight": [d0, d1, d2, d3],
+        }
+
+        # find temp_var -> dst
+        aoa_statements = [
+            "layers.0.gate_up_fused_proj.weight -> temp_var, fused_ffn \n",
+            "temp_var^T -> new_name_layers.0.gate_up_fused_proj.weight \n",
+        ]
+
+        aoa_engine = AOAEngine(
+            aoa_config={"aoa_statements": aoa_statements},
+            source_state_shard_info=source_state_shard_info,
+            destination_state_shard_info=destination_state_shard_info,
+        )
+
+        # new_name_up_proj_0
+        query = ShardedWeightDesc(
+            key="new_name_layers.0.gate_up_fused_proj.weight",
+            local_shape=(2, 1),
+            global_shape=(2, 8),
+            global_offset=(0, 0),
+            dtype="float32",
+        )
+
+        target_slice_1 = ShardedWeightDesc(
+            key='new_name_layers.0.gate_up_fused_proj.weight',
+            local_shape=(1, 1),
+            global_shape=(2, 8),
+            global_offset=(0, 0),
+            dtype='float32',
+        )
+
+        target_slice_2 = ShardedWeightDesc(
+            key='new_name_layers.0.gate_up_fused_proj.weight',
+            local_shape=(1, 1),
+            global_shape=(2, 8),
+            global_offset=(1, 0),
+            dtype='float32',
+        )
+
+        src_slice_1 = ShardedWeightDesc(
+            key='layers.0.gate_up_fused_proj.weight',
+            local_shape=(1, 1),
+            global_shape=(2, 8),
+            global_offset=(0, 0),
+            dtype='float32',
+        )
+
+        src_slice_2 = ShardedWeightDesc(
+            key='layers.0.gate_up_fused_proj.weight',
+            local_shape=(1, 1),
+            global_shape=(2, 8),
+            global_offset=(0, 2),
+            dtype='float32',
+        )
+
+        shard_mapping_entry_1 = ShardMappingEntry(
+            target_slice=target_slice_1,
+            source_slice=src_slice_1,
+            postprocess_list=['[1, 0]'],
+        )
+        shard_mapping_entry_2 = ShardMappingEntry(
+            target_slice=target_slice_2,
+            source_slice=src_slice_2,
+            postprocess_list=['[1, 0]'],
+        )
+        answer = [shard_mapping_entry_1, shard_mapping_entry_2]
+
+        result = aoa_engine.find_shard_sources(query)
+        self.assertEqual(result, answer)
+
+        s0 = ShardedWeightDesc(
+            key="layers.0.gate_up_fused_proj.weight",
+            local_shape=(4, 2),
+            global_shape=(8, 2),
+            global_offset=(0, 0),
+            dtype="float32",
+        )
+
+        s1 = ShardedWeightDesc(
+            key="layers.0.gate_up_fused_proj.weight",
+            local_shape=(4, 2),
+            global_shape=(8, 2),
+            global_offset=(4, 0),
+            dtype="float32",
+        )
+
+        d0 = ShardedWeightDesc(
+            key="new_name_layers.0.gate_up_fused_proj.weight",
+            local_shape=(2, 2),
+            global_shape=(2, 8),
+            global_offset=(0, 0),
+            dtype="float32",
+        )
+
+        d1 = ShardedWeightDesc(
+            key="new_name_layers.0.gate_up_fused_proj.weight",
+            local_shape=(2, 2),
+            global_shape=(2, 8),
+            global_offset=(0, 2),
+            dtype="float32",
+        )
+
+        d2 = ShardedWeightDesc(
+            key="new_name_layers.0.gate_up_fused_proj.weight",
+            local_shape=(2, 2),
+            global_shape=(2, 8),
+            global_offset=(0, 4),
+            dtype="float32",
+        )
+
+        d3 = ShardedWeightDesc(
+            key="new_name_layers.0.gate_up_fused_proj.weight",
+            local_shape=(2, 2),
+            global_shape=(2, 8),
+            global_offset=(0, 6),
+            dtype="float32",
+        )
+
+        source_state_shard_info = {
+            "layers.0.gate_up_fused_proj.weight": [s0, s1],
+        }
+        destination_state_shard_info = {
+            "new_name_layers.0.gate_up_fused_proj.weight": [d0, d1, d2, d3],
+        }
+
+        # find temp_var -> src
+        aoa_statements = [
+            "layers.0.gate_up_fused_proj.weight^T -> temp_var \n",
+            "temp_var -> new_name_layers.0.gate_up_fused_proj.weight,fused_ffn\n",
+        ]
+
+        aoa_engine = AOAEngine(
+            aoa_config={"aoa_statements": aoa_statements},
+            source_state_shard_info=source_state_shard_info,
+            destination_state_shard_info=destination_state_shard_info,
+        )
+
+        # new_name_up_proj_0
+        query = ShardedWeightDesc(
+            key="new_name_layers.0.gate_up_fused_proj.weight",
+            local_shape=(2, 1),
+            global_shape=(2, 8),
+            global_offset=(0, 0),
+            dtype="float32",
+        )
+
+        src_slice_1 = ShardedWeightDesc(
+            key='layers.0.gate_up_fused_proj.weight',
+            local_shape=(1, 2),
+            global_shape=(8, 2),
+            global_offset=(0, 0),
+            dtype='float32',
+        )
+
+        shard_mapping_entry_1 = ShardMappingEntry(
+            target_slice=query,
+            source_slice=src_slice_1,
+            postprocess_list=['[1, 0]'],
+        )
+
+        answer = [shard_mapping_entry_1]
+
+        result = aoa_engine.find_shard_sources(query)
+        self.assertEqual(result, answer)
+
+    def test_aoa_transpose_reverse(self):
+        s0 = ShardedWeightDesc(
+            key="s0",
+            local_shape=(1, 2, 4),
+            global_shape=(1, 2, 4),
+            global_offset=(0, 0, 0),
+            dtype="float32",
+        )
+        d0 = ShardedWeightDesc(
+            key="d0",
+            local_shape=(4, 1, 2),
+            global_shape=(4, 1, 2),
+            global_offset=(0, 0, 0),
+            dtype="float32",
+        )
+        aoa_statements = [
+            "s0 -> d0, permute= '[2, 0, 1]' \n",
+        ]
+        aoa_config = {
+            "aoa_statements": aoa_statements,
+            "aoa_config_reverse": True,
+        }
+        source_state_shard_info = {
+            "d0": [d0],
+        }
+        destination_state_shard_info = {
+            "s0": [s0],
+        }
+        aoa_engine = AOAEngine(
+            aoa_config=aoa_config,
+            source_state_shard_info=source_state_shard_info,
+            destination_state_shard_info=destination_state_shard_info,
+        )
+
+        query = ShardedWeightDesc(
+            key="s0",
+            local_shape=(1, 2, 4),
+            global_shape=(1, 2, 4),
+            global_offset=(0, 0, 0),
+            dtype="float32",
+        )
+        answer = [
+            ShardMappingEntry(
+                target_slice=s0,
+                source_slice=d0,
+                postprocess_list=['[1, 2, 0]'],
+            )
+        ]
+        result = aoa_engine.find_shard_sources(query)
+        self.assertEqual(result, answer)
+
+
+if __name__ == '__main__':
+    unittest.main()

test/flex_checkpoint/test_aoa_engine_transpose_cast.py

@@ -0,0 +1,674 @@
+# 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 unittest
+
+from paddle.distributed.flex_checkpoint.aoa.aoa_engine import (
+    AOAEngine,
+    ShardedWeightDesc,
+    ShardMappingEntry,
+)
+
+
+class TestAOAEngineTransposeCast(unittest.TestCase):
+    def setUp(self):
+        self.setup_statements()
+        self.aoa_engine = AOAEngine(
+            aoa_config={"aoa_statements": self.aoa_statements},
+            source_state_shard_info=self.source_state_shard_info,
+            destination_state_shard_info=self.destination_state_shard_info,
+        )
+        self.generate_query_answer()
+
+    def setup_statements(self):
+        s0 = ShardedWeightDesc(
+            key="s0",
+            local_shape=(2, 2),
+            global_shape=(2, 2),
+            global_offset=(0, 0),
+        )
+        s1 = ShardedWeightDesc(
+            key="s1",
+            local_shape=(2, 2),
+            global_shape=(2, 2),
+            global_offset=(0, 0),
+        )
+
+        d0 = ShardedWeightDesc(
+            key="d0",
+            local_shape=(4, 1),
+            global_shape=(4, 1),
+            global_offset=(0, 0),
+        )
+        d1 = ShardedWeightDesc(
+            key="d1",
+            local_shape=(4, 1),
+            global_shape=(4, 1),
+            global_offset=(0, 0),
+        )
+
+        self.source_state_shard_info = {
+            "s0": [s0],
+            "s1": [s1],
+        }
+        self.destination_state_shard_info = {
+            "d0": [d0],
+            "d1": [d1],
+        }
+
+        self.aoa_statements = [
+            "s0, s1 -> s, axis = 1 \n",
+            "s -> s, dtype = 'float32'\n",
+            "s^T -> d\n",
+            "d -> d0, d1, axis = 1",
+        ]
+
+    def generate_query_answer(self):
+        self.queries = []
+        self.answers = []
+
+        # ======================================================
+        # Query 1:
+        query = ShardedWeightDesc(
+            key="d0",
+            local_shape=(2, 1),
+            global_shape=(4, 1),
+            global_offset=(2, 0),
+        )
+        src_sharded_weight_desc = ShardedWeightDesc(
+            key="s1",
+            local_shape=(1, 2),
+            global_shape=(2, 2),
+            global_offset=(0, 0),
+        )
+        shard_mapping_entry = ShardMappingEntry(
+            target_slice=query,
+            source_slice=src_sharded_weight_desc,
+            postprocess_list=["float32", "[1, 0]"],
+        )
+        answer = [shard_mapping_entry]
+        self.queries.append(query)
+        self.answers.append(answer)
+
+        # ======================================================
+        # Query 2:
+        query = ShardedWeightDesc(
+            key="d1",
+            local_shape=(2, 1),
+            global_shape=(4, 1),
+            global_offset=(0, 0),
+        )
+        src_sharded_weight_desc = ShardedWeightDesc(
+            key="s0",
+            local_shape=(1, 2),
+            global_shape=(2, 2),
+            global_offset=(1, 0),
+        )
+        shard_mapping_entry = ShardMappingEntry(
+            target_slice=query,
+            source_slice=src_sharded_weight_desc,
+            postprocess_list=["float32", "[1, 0]"],
+        )
+        answer = [shard_mapping_entry]
+        self.queries.append(query)
+        self.answers.append(answer)
+
+        # ======================================================
+        # Query 3:
+        query = ShardedWeightDesc(
+            key="d1",
+            local_shape=(4, 1),
+            global_shape=(4, 1),
+            global_offset=(0, 0),
+        )
+
+        # d1[0:2, :] <--- s0[1, :]^T
+        src_sharded_weight_desc0 = ShardedWeightDesc(
+            key="s0",
+            local_shape=(1, 2),
+            global_shape=(2, 2),
+            global_offset=(1, 0),
+        )
+        dst_sharded_weight_desc0 = ShardedWeightDesc(
+            key="d1",
+            local_shape=(2, 1),
+            global_shape=(4, 1),
+            global_offset=(0, 0),
+        )
+
+        # d1[2:4, :] <--- s1[1, :]^T
+        src_sharded_weight_desc1 = ShardedWeightDesc(
+            key="s1",
+            local_shape=(1, 2),
+            global_shape=(2, 2),
+            global_offset=(1, 0),
+        )
+        dst_sharded_weight_desc1 = ShardedWeightDesc(
+            key="d1",
+            local_shape=(2, 1),
+            global_shape=(4, 1),
+            global_offset=(2, 0),
+        )
+
+        shard_mapping_entry0 = ShardMappingEntry(
+            target_slice=dst_sharded_weight_desc0,
+            source_slice=src_sharded_weight_desc0,
+            postprocess_list=["float32", "[1, 0]"],
+        )
+        shard_mapping_entry1 = ShardMappingEntry(
+            target_slice=dst_sharded_weight_desc1,
+            source_slice=src_sharded_weight_desc1,
+            postprocess_list=["float32", "[1, 0]"],
+        )
+        answer = [shard_mapping_entry0, shard_mapping_entry1]
+        self.queries.append(query)
+        self.answers.append(answer)
+
+    def test_transpose(self):
+        for idx in range(len(self.queries)):
+            query = self.queries[idx]
+            answer = self.answers[idx]
+            result = self.aoa_engine.find_shard_sources(query)
+            self.assertEqual(result, answer)
+
+
+class TestAOAEngineTransposeCast2(TestAOAEngineTransposeCast):
+    def setup_statements(self):
+        s0 = ShardedWeightDesc(
+            key="s0",
+            local_shape=(4, 1),
+            global_shape=(4, 1),
+            global_offset=(0, 0),
+        )
+        s1 = ShardedWeightDesc(
+            key="s1",
+            local_shape=(4, 1),
+            global_shape=(4, 1),
+            global_offset=(0, 0),
+        )
+
+        d0 = ShardedWeightDesc(
+            key="d0",
+            local_shape=(2, 2),
+            global_shape=(2, 2),
+            global_offset=(0, 0),
+        )
+        d1 = ShardedWeightDesc(
+            key="d1",
+            local_shape=(2, 2),
+            global_shape=(2, 2),
+            global_offset=(0, 0),
+        )
+
+        self.source_state_shard_info = {
+            "s0": [s0],
+            "s1": [s1],
+        }
+        self.destination_state_shard_info = {
+            "d0": [d0],
+            "d1": [d1],
+        }
+
+        self.aoa_statements = [
+            "s0^T -> s0\n",
+            "s1^T -> s1\n",
+            "s0, s1 -> s, axis = 0\n",
+            "s -> s, dtype = 'float16'\n",
+            "s -> d0, d1, axis = 1",
+        ]
+
+    def generate_query_answer(self):
+        self.queries = []
+        self.answers = []
+
+        # ======================================================
+        # Query 1:
+        query = ShardedWeightDesc(
+            key="d0",
+            local_shape=(1, 2),
+            global_shape=(2, 2),
+            global_offset=(1, 0),
+        )
+        src_sharded_weight_desc = ShardedWeightDesc(
+            key="s1",
+            local_shape=(2, 1),
+            global_shape=(4, 1),
+            global_offset=(0, 0),
+        )
+        shard_mapping_entry = ShardMappingEntry(
+            target_slice=query,
+            source_slice=src_sharded_weight_desc,
+            postprocess_list=["[1, 0]", "float16"],
+        )
+        answer = [shard_mapping_entry]
+        self.queries.append(query)
+        self.answers.append(answer)
+
+        # ======================================================
+        # Query 2:
+        query = ShardedWeightDesc(
+            key="d1",
+            local_shape=(1, 2),
+            global_shape=(2, 2),
+            global_offset=(0, 0),
+        )
+        src_sharded_weight_desc = ShardedWeightDesc(
+            key="s0",
+            local_shape=(2, 1),
+            global_shape=(4, 1),
+            global_offset=(2, 0),
+        )
+        shard_mapping_entry = ShardMappingEntry(
+            target_slice=query,
+            source_slice=src_sharded_weight_desc,
+            postprocess_list=["[1, 0]", "float16"],
+        )
+        answer = [shard_mapping_entry]
+        self.queries.append(query)
+        self.answers.append(answer)
+
+        # ======================================================
+        # Query 3:
+        query = ShardedWeightDesc(
+            key="d1",
+            local_shape=(2, 2),
+            global_shape=(2, 2),
+            global_offset=(0, 0),
+        )
+
+        # d1[0:1, :] <--- s0[2:4, :]^T
+        src_sharded_weight_desc0 = ShardedWeightDesc(
+            key="s0",
+            local_shape=(2, 1),
+            global_shape=(4, 1),
+            global_offset=(2, 0),
+        )
+        dst_sharded_weight_desc0 = ShardedWeightDesc(
+            key="d1",
+            local_shape=(1, 2),
+            global_shape=(2, 2),
+            global_offset=(0, 0),
+        )
+
+        # d1[1:2, :] <--- s1[2:4, :]^T
+        src_sharded_weight_desc1 = ShardedWeightDesc(
+            key="s1",
+            local_shape=(2, 1),
+            global_shape=(4, 1),
+            global_offset=(2, 0),
+        )
+        dst_sharded_weight_desc1 = ShardedWeightDesc(
+            key="d1",
+            local_shape=(1, 2),
+            global_shape=(2, 2),
+            global_offset=(1, 0),
+        )
+
+        shard_mapping_entry0 = ShardMappingEntry(
+            target_slice=dst_sharded_weight_desc0,
+            source_slice=src_sharded_weight_desc0,
+            postprocess_list=["[1, 0]", "float16"],
+        )
+        shard_mapping_entry1 = ShardMappingEntry(
+            target_slice=dst_sharded_weight_desc1,
+            source_slice=src_sharded_weight_desc1,
+            postprocess_list=["[1, 0]", "float16"],
+        )
+        answer = [shard_mapping_entry0, shard_mapping_entry1]
+        self.queries.append(query)
+        self.answers.append(answer)
+
+
+class TestAOAEngineTransposeCast3(TestAOAEngineTransposeCast):
+    def setup_statements(self):
+        s0 = ShardedWeightDesc(
+            key="s0",
+            local_shape=(3, 4),
+            global_shape=(3, 4),
+            global_offset=(0, 0),
+        )
+
+        d0 = ShardedWeightDesc(
+            key="d0",
+            local_shape=(1, 6),
+            global_shape=(1, 6),
+            global_offset=(0, 0),
+        )
+        d1 = ShardedWeightDesc(
+            key="d1",
+            local_shape=(6, 1),
+            global_shape=(6, 1),
+            global_offset=(0, 0),
+        )
+
+        self.source_state_shard_info = {
+            "s0": [s0],
+        }
+        self.destination_state_shard_info = {
+            "d0": [d0],
+            "d1": [d1],
+        }
+
+        self.aoa_statements = [
+            "s0 -> a1, a2, a3, a4, axis = 1\n",
+            "a2^T -> b2\n",
+            "a3^T -> b3\n",
+            "b2, b3 -> d0, axis = 1\n",
+            "a3, a4 -> d1, axis = 0\n",
+        ]
+
+    def generate_query_answer(self):
+        self.queries = []
+        self.answers = []
+
+        # ======================================================
+        # Query 1:
+        query = ShardedWeightDesc(
+            key="d0",
+            local_shape=(1, 6),
+            global_shape=(1, 6),
+            global_offset=(0, 0),
+        )
+        # d0[:, 0:3] <--- s0[:, 1:2]^T
+        src_sharded_weight_desc0 = ShardedWeightDesc(
+            key="s0",
+            local_shape=(3, 1),
+            global_shape=(3, 4),
+            global_offset=(0, 1),
+        )
+        dst_sharded_weight_desc0 = ShardedWeightDesc(
+            key="d0",
+            local_shape=(1, 3),
+            global_shape=(1, 6),
+            global_offset=(0, 0),
+        )
+
+        # d0[:, 3:6] <--- s0[:, 2:3]^T
+        src_sharded_weight_desc1 = ShardedWeightDesc(
+            key="s0",
+            local_shape=(3, 1),
+            global_shape=(3, 4),
+            global_offset=(0, 2),
+        )
+        dst_sharded_weight_desc1 = ShardedWeightDesc(
+            key="d0",
+            local_shape=(1, 3),
+            global_shape=(1, 6),
+            global_offset=(0, 3),
+        )
+
+        shard_mapping_entry0 = ShardMappingEntry(
+            target_slice=dst_sharded_weight_desc0,
+            source_slice=src_sharded_weight_desc0,
+            postprocess_list=["[1, 0]"],
+        )
+        shard_mapping_entry1 = ShardMappingEntry(
+            target_slice=dst_sharded_weight_desc1,
+            source_slice=src_sharded_weight_desc1,
+            postprocess_list=["[1, 0]"],
+        )
+        answer = [shard_mapping_entry0, shard_mapping_entry1]
+        self.queries.append(query)
+        self.answers.append(answer)
+
+        # ======================================================
+        # Query 2:
+        query = ShardedWeightDesc(
+            key="d1",
+            local_shape=(6, 1),
+            global_shape=(6, 1),
+            global_offset=(0, 0),
+        )
+        # d1[0:3, :] <--- s0[:, 2:3]
+        src_sharded_weight_desc0 = ShardedWeightDesc(
+            key="s0",
+            local_shape=(3, 1),
+            global_shape=(3, 4),
+            global_offset=(0, 2),
+        )
+        dst_sharded_weight_desc0 = ShardedWeightDesc(
+            key="d1",
+            local_shape=(3, 1),
+            global_shape=(6, 1),
+            global_offset=(0, 0),
+        )
+
+        # d1[3:6, :] <--- s0[:, 3:4]
+        src_sharded_weight_desc1 = ShardedWeightDesc(
+            key="s0",
+            local_shape=(3, 1),
+            global_shape=(3, 4),
+            global_offset=(0, 3),
+        )
+        dst_sharded_weight_desc1 = ShardedWeightDesc(
+            key="d1",
+            local_shape=(3, 1),
+            global_shape=(6, 1),
+            global_offset=(3, 0),
+        )
+
+        shard_mapping_entry0 = ShardMappingEntry(
+            target_slice=dst_sharded_weight_desc0,
+            source_slice=src_sharded_weight_desc0,
+            postprocess_list=None,
+        )
+        shard_mapping_entry1 = ShardMappingEntry(
+            target_slice=dst_sharded_weight_desc1,
+            source_slice=src_sharded_weight_desc1,
+            postprocess_list=None,
+        )
+        answer = [shard_mapping_entry0, shard_mapping_entry1]
+        self.queries.append(query)
+        self.answers.append(answer)
+
+
+class TestAOAEngineTransposeCast4(TestAOAEngineTransposeCast):
+    def setup_statements(self):
+        s0 = ShardedWeightDesc(
+            key="s0",
+            local_shape=(4, 1, 3),
+            global_shape=(4, 1, 3),
+            global_offset=(0, 0, 0),
+        )
+        s1 = ShardedWeightDesc(
+            key="s1",
+            local_shape=(4, 1, 3),
+            global_shape=(4, 1, 3),
+            global_offset=(0, 0, 0),
+        )
+
+        d0 = ShardedWeightDesc(
+            key="d0",
+            local_shape=(1, 4, 4),
+            global_shape=(1, 4, 4),
+            global_offset=(0, 0, 0),
+        )
+        d1 = ShardedWeightDesc(
+            key="d1",
+            local_shape=(1, 4, 2),
+            global_shape=(1, 4, 2),
+            global_offset=(0, 0, 0),
+        )
+
+        self.source_state_shard_info = {
+            "s0": [s0],
+            "s1": [s1],
+        }
+        self.destination_state_shard_info = {
+            "d0": [d0],
+            "d1": [d1],
+        }
+
+        self.aoa_statements = [
+            "s0, s1 -> s, axis = 1\n",
+            "s -> s, dtype = 'bfloat16'\n",
+            "s -> a, permute = '[2, 0, 1]'\n",
+            "a -> b1, b2, b3, axis = 0\n",
+            "b1 -> b1, permute = '[0, 2, 1]'\n",
+            "b2 -> b2, permute = '[0, 2, 1]'\n",
+            "b1, b2 -> d0, axis = 1\n",
+            "b3 -> d1\n",
+            "d1 -> d1, dtype = 'float32'",
+        ]
+
+    def generate_query_answer(self):
+        self.queries = []
+        self.answers = []
+
+        # ======================================================
+        # Query 1:
+        query = ShardedWeightDesc(
+            key="d0",
+            local_shape=(1, 4, 4),
+            global_shape=(1, 4, 4),
+            global_offset=(0, 0, 0),
+        )
+        # d0[:, 0:1, :] <--- s0[:, :, 0:1].transpose([2, 0, 1]).transpose([0, 2, 1])
+        src_sharded_weight_desc0 = ShardedWeightDesc(
+            key="s0",
+            local_shape=(4, 1, 1),
+            global_shape=(4, 1, 3),
+            global_offset=(0, 0, 0),
+        )
+        dst_sharded_weight_desc0 = ShardedWeightDesc(
+            key="d0",
+            local_shape=(1, 1, 4),
+            global_shape=(1, 4, 4),
+            global_offset=(0, 0, 0),
+        )
+
+        # d0[:, 1:2, :] <--- s1[:, :, 0:1].transpose([2, 0, 1]).transpose([0, 2, 1])
+        src_sharded_weight_desc1 = ShardedWeightDesc(
+            key="s1",
+            local_shape=(4, 1, 1),
+            global_shape=(4, 1, 3),
+            global_offset=(0, 0, 0),
+        )
+        dst_sharded_weight_desc1 = ShardedWeightDesc(
+            key="d0",
+            local_shape=(1, 1, 4),
+            global_shape=(1, 4, 4),
+            global_offset=(0, 1, 0),
+        )
+
+        # d0[:, 2:3, :] <--- s0[:, :, 1:2].transpose([2, 0, 1]).transpose([0, 2, 1])
+        src_sharded_weight_desc2 = ShardedWeightDesc(
+            key="s0",
+            local_shape=(4, 1, 1),
+            global_shape=(4, 1, 3),
+            global_offset=(0, 0, 1),
+        )
+        dst_sharded_weight_desc2 = ShardedWeightDesc(
+            key="d0",
+            local_shape=(1, 1, 4),
+            global_shape=(1, 4, 4),
+            global_offset=(0, 2, 0),
+        )
+
+        # d0[:, 3:4, :] <--- s1[:, :, 1:2].transpose([2, 0, 1]).transpose([0, 2, 1])
+        src_sharded_weight_desc3 = ShardedWeightDesc(
+            key="s1",
+            local_shape=(4, 1, 1),
+            global_shape=(4, 1, 3),
+            global_offset=(0, 0, 1),
+        )
+        dst_sharded_weight_desc3 = ShardedWeightDesc(
+            key="d0",
+            local_shape=(1, 1, 4),
+            global_shape=(1, 4, 4),
+            global_offset=(0, 3, 0),
+        )
+
+        shard_mapping_entry0 = ShardMappingEntry(
+            target_slice=dst_sharded_weight_desc0,
+            source_slice=src_sharded_weight_desc0,
+            postprocess_list=["bfloat16", "[2, 0, 1]", "[0, 2, 1]"],
+        )
+        shard_mapping_entry1 = ShardMappingEntry(
+            target_slice=dst_sharded_weight_desc1,
+            source_slice=src_sharded_weight_desc1,
+            postprocess_list=["bfloat16", "[2, 0, 1]", "[0, 2, 1]"],
+        )
+        shard_mapping_entry2 = ShardMappingEntry(
+            target_slice=dst_sharded_weight_desc2,
+            source_slice=src_sharded_weight_desc2,
+            postprocess_list=["bfloat16", "[2, 0, 1]", "[0, 2, 1]"],
+        )
+        shard_mapping_entry3 = ShardMappingEntry(
+            target_slice=dst_sharded_weight_desc3,
+            source_slice=src_sharded_weight_desc3,
+            postprocess_list=["bfloat16", "[2, 0, 1]", "[0, 2, 1]"],
+        )
+        answer = [
+            shard_mapping_entry0,
+            shard_mapping_entry1,
+            shard_mapping_entry2,
+            shard_mapping_entry3,
+        ]
+        self.queries.append(query)
+        self.answers.append(answer)
+
+        # ======================================================
+        # Query 2:
+        query = ShardedWeightDesc(
+            key="d1",
+            local_shape=(1, 4, 2),
+            global_shape=(1, 4, 2),
+            global_offset=(0, 0, 0),
+        )
+        # d1[:, :, 0:1] <--- s0[:, :, 2:3].transpose([2, 0, 1])
+        src_sharded_weight_desc0 = ShardedWeightDesc(
+            key="s0",
+            local_shape=(4, 1, 1),
+            global_shape=(4, 1, 3),
+            global_offset=(0, 0, 2),
+        )
+        dst_sharded_weight_desc0 = ShardedWeightDesc(
+            key="d1",
+            local_shape=(1, 4, 1),
+            global_shape=(1, 4, 2),
+            global_offset=(0, 0, 0),
+        )
+
+        # d1[:, :, 1:2] <--- s1[:, :, 2:3].transpose([2, 0, 1])
+        src_sharded_weight_desc1 = ShardedWeightDesc(
+            key="s1",
+            local_shape=(4, 1, 1),
+            global_shape=(4, 1, 3),
+            global_offset=(0, 0, 2),
+        )
+        dst_sharded_weight_desc1 = ShardedWeightDesc(
+            key="d1",
+            local_shape=(1, 4, 1),
+            global_shape=(1, 4, 2),
+            global_offset=(0, 0, 1),
+        )
+
+        shard_mapping_entry0 = ShardMappingEntry(
+            target_slice=dst_sharded_weight_desc0,
+            source_slice=src_sharded_weight_desc0,
+            postprocess_list=["bfloat16", "[2, 0, 1]", "float32"],
+        )
+        shard_mapping_entry1 = ShardMappingEntry(
+            target_slice=dst_sharded_weight_desc1,
+            source_slice=src_sharded_weight_desc1,
+            postprocess_list=["bfloat16", "[2, 0, 1]", "float32"],
+        )
+        answer = [shard_mapping_entry0, shard_mapping_entry1]
+        self.queries.append(query)
+        self.answers.append(answer)
+
+
+if __name__ == '__main__':
+    unittest.main()

test/flex_checkpoint/test_balance_file_reading.py

@@ -0,0 +1,249 @@
+#!/usr/bin/env python3
+# 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.
+
+import unittest
+from unittest.mock import patch
+
+from paddle.distributed.flex_checkpoint.dcp.load_state_dict import (
+    get_rank_to_read_files,
+)
+
+
+class TestGetRankToReadFiles(unittest.TestCase):
+    """Unit tests for get_rank_to_read_files function."""
+
+    def setUp(self):
+        """Set up test fixtures."""
+        self.mock_rank = 0
+        self.patcher = patch(
+            'paddle.distributed.get_rank', return_value=self.mock_rank
+        )
+        self.mock_get_rank = self.patcher.start()
+        self.addCleanup(self.patcher.stop)
+
+    def test_local_files_assignment(self):
+        """Test assignment when rank has all required files locally."""
+        # Arrange
+        rank_to_required = {
+            0: ['file1.distcp', 'file2.distcp'],
+            1: ['file3.distcp', 'file4.distcp'],
+        }
+
+        rank_to_available_files = {
+            0: ['file1.distcp', 'file2.distcp'],
+            1: ['file3.distcp', 'file4.distcp'],
+        }
+
+        # Act
+        result = get_rank_to_read_files(
+            rank_to_required, rank_to_available_files
+        )
+
+        # Assert
+        expected_files = ['file1.distcp', 'file2.distcp']
+        self.assertEqual(sorted(result), sorted(expected_files))
+
+    def test_cross_node_files_assignment(self):
+        """Test assignment when rank needs files from other nodes."""
+        # Arrange
+        rank_to_required = {
+            0: ['file1.distcp', 'file2.distcp'],
+            1: ['file1.distcp', 'file3.distcp'],
+        }
+
+        rank_to_available_files = {
+            0: ['file1.distcp'],
+            1: ['file2.distcp', 'file3.distcp'],
+        }
+
+        # Act
+        result = get_rank_to_read_files(
+            rank_to_required, rank_to_available_files
+        )
+
+        # Assert
+        # Rank 0 should get file1 locally and file2 might be assigned from rank 1
+        self.assertIn('file1.distcp', result)
+        self.assertEqual(len(result), 1)  # Should balance workload
+
+    def test_empty_rank_assignment(self):
+        """Test when current rank has no files to read."""
+        # Arrange
+        rank_to_required = {
+            1: ['file1.distcp', 'file2.distcp'],
+            2: ['file3.distcp', 'file4.distcp'],
+        }
+
+        rank_to_available_files = {
+            1: ['file1.distcp', 'file2.distcp'],
+            2: ['file3.distcp', 'file4.distcp'],
+        }
+
+        self.mock_rank = 0  # Current rank has nothing to do
+
+        # Act
+        result = get_rank_to_read_files(
+            rank_to_required, rank_to_available_files
+        )
+
+        # Assert
+        self.assertEqual(result, [])
+
+    def test_load_balancing_multiple_candidates(self):
+        """Test load balancing when multiple ranks can read the same file."""
+        # Arrange
+        rank_to_required = {
+            0: ['shared_file.distcp', 'file2.distcp'],
+            1: ['shared_file.distcp', 'file3.distcp'],
+            2: ['file4.distcp', 'file5.distcp'],
+        }
+
+        rank_to_available_files = {
+            0: ['shared_file.distcp', 'file2.distcp'],
+            1: ['shared_file.distcp', 'file3.distcp'],
+            2: ['shared_file.distcp', 'file4.distcp', 'file5.distcp'],
+        }
+
+        # Act
+        result = get_rank_to_read_files(
+            rank_to_required, rank_to_available_files
+        )
+
+        # Assert
+        # Should include shared_file and file2, but workload should be balanced
+        self.assertIn('file2.distcp', result)
+        self.assertEqual(len(result), 2)
+
+    def test_missing_file_warning(self):
+        """Test behavior when required file is not available on any rank."""
+        # Arrange
+        rank_to_required = {
+            0: ['missing_file.distcp', 'existing_file.distcp'],
+            1: ['file2.distcp'],
+        }
+
+        rank_to_available_files = {
+            0: ['existing_file.distcp'],
+            1: ['file2.distcp'],
+        }
+        # missing_file.distcp is not in any rank_to_available_files
+
+        # Act & Assert - should not raise exception but handle gracefully
+        result = get_rank_to_read_files(
+            rank_to_required, rank_to_available_files
+        )
+
+        # Should still return files that are available
+        self.assertIn('existing_file.distcp', result)
+
+    def test_single_rank_scenario(self):
+        """Test single rank scenario (non-distributed mode)."""
+        # Arrange
+        rank_to_required = {0: ['file1.distcp', 'file2.distcp', 'file3.distcp']}
+        rank_to_available_files = {
+            0: ['file1.distcp', 'file2.distcp', 'file3.distcp']
+        }
+
+        # Act
+        result = get_rank_to_read_files(
+            rank_to_required, rank_to_available_files
+        )
+
+        # Assert
+        expected_files = ['file1.distcp', 'file2.distcp', 'file3.distcp']
+        self.assertEqual(sorted(result), sorted(expected_files))
+
+    def test_empty_inputs(self):
+        """Test with empty input dictionaries."""
+        # Arrange
+        rank_to_required = {}
+        rank_to_available_files = {}
+
+        # Act
+        result = get_rank_to_read_files(
+            rank_to_required, rank_to_available_files
+        )
+
+        # Assert
+        self.assertEqual(result, [])
+
+    def test_rank_with_no_local_files(self):
+        """Test when a rank has logical files but no local files available."""
+        # Arrange
+        rank_to_required = {
+            0: ['file1.distcp', 'file2.distcp'],
+            1: ['file3.distcp'],
+        }
+
+        rank_to_available_files = {
+            1: ['file1.distcp', 'file2.distcp', 'file3.distcp']
+        }
+        # Rank 0 has no local files but needs file1 and file2
+
+        # Act
+        result = get_rank_to_read_files(
+            rank_to_required, rank_to_available_files
+        )
+
+        # Assert
+        # Rank 0 should get files assigned from rank 1
+        self.assertEqual(len(result), 0)
+
+    def test_rank_not_in_mappings(self):
+        """Test when current rank is not present in input mappings."""
+        # Arrange
+        rank_to_required = {1: ['file1.distcp'], 2: ['file2.distcp']}
+
+        rank_to_available_files = {1: ['file1.distcp'], 2: ['file2.distcp']}
+
+        self.mock_rank = 0  # Current rank not in mappings
+
+        # Act
+        result = get_rank_to_read_files(
+            rank_to_required, rank_to_available_files
+        )
+
+        # Assert
+        self.assertEqual(result, [])
+
+    def test_duplicate_files_across_ranks(self):
+        """Test handling of duplicate files across different ranks."""
+        # Arrange
+        rank_to_required = {
+            0: ['file1.distcp', 'file2.distcp'],
+            1: ['file1.distcp', 'file3.distcp'],  # file1 duplicated
+            2: ['file4.distcp'],
+        }
+
+        rank_to_available_files = {
+            0: ['file1.distcp', 'file2.distcp'],
+            1: ['file1.distcp', 'file3.distcp'],
+            2: ['file4.distcp'],
+        }
+
+        # Act
+        result = get_rank_to_read_files(
+            rank_to_required, rank_to_available_files
+        )
+
+        # Assert
+        # file1 should be assigned to only one rank (load balanced)
+        self.assertIn('file1.distcp', result)
+        # Should have exactly 2 files (file1 plus one more for balance)
+        self.assertEqual(len(result), 2)
+
+
+if __name__ == '__main__':
+    unittest.main()

test/flex_checkpoint/test_key_validation.py

@@ -0,0 +1,856 @@
+# 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.
+
+import unittest
+from unittest.mock import MagicMock, patch
+
+import paddle
+from paddle.distributed.flex_checkpoint.dcp.key_validation import (
+    AOAMappingEntry,
+    AOASliceMapping,
+    KeyValidationResult,
+    ShapeMismatchInfo,
+    _append_src_lines,
+    _build_aoa_mappings,
+    _classify_mappings,
+    _describe_ops,
+    _emit,
+    _format_key_list,
+    _format_pattern_groups,
+    _format_slice_range,
+    _get_signature,
+    _group_by_signature,
+    _group_keys_adaptive,
+    _print_aoa_report,
+    _print_standard_report,
+    _slice_covers_full,
+    _try_fold_src_keys,
+    validate_and_report_keys_aoa,
+    validate_and_report_keys_standard,
+)
+from paddle.distributed.flex_checkpoint.dcp.metadata import (
+    LocalTensorIndex,
+    LocalTensorMetadata,
+    Metadata,
+)
+
+
+class TestSliceCoversFull(unittest.TestCase):
+    def test_covers_full(self):
+        sl = (slice(0, 4), slice(0, 8))
+        self.assertTrue(_slice_covers_full(sl, (4, 8)))
+
+    def test_not_covers_partial(self):
+        sl = (slice(0, 2), slice(0, 8))
+        self.assertFalse(_slice_covers_full(sl, (4, 8)))
+
+    def test_not_covers_mismatched_dims(self):
+        sl = (slice(0, 4),)
+        self.assertFalse(_slice_covers_full(sl, (4, 8)))
+
+    def test_non_zero_start(self):
+        sl = (slice(1, 4), slice(0, 8))
+        self.assertFalse(_slice_covers_full(sl, (4, 8)))
+
+
+class TestFormatSliceRange(unittest.TestCase):
+    def test_basic(self):
+        src_sl = (slice(0, 4), slice(0, 8))
+        dst_sl = (slice(0, 4), slice(0, 8))
+        result = _format_slice_range(src_sl, dst_sl)
+        self.assertIn("0:4", result)
+        self.assertIn("0:8", result)
+        self.assertIn("->", result)
+
+    def test_partial_slices(self):
+        src_sl = (slice(2, 6),)
+        dst_sl = (slice(0, 4),)
+        result = _format_slice_range(src_sl, dst_sl)
+        self.assertIn("2:6", result)
+        self.assertIn("0:4", result)
+
+
+class TestTryFoldSrcKeys(unittest.TestCase):
+    def test_fold_consecutive(self):
+        keys = [f"model.experts.{i}.weight" for i in range(8)]
+        result = _try_fold_src_keys(keys)
+        self.assertIsNotNone(result)
+        self.assertIn("{0..7}", result)
+
+    def test_no_fold_different_patterns(self):
+        keys = ["model.a.weight", "model.b.weight"]
+        result = _try_fold_src_keys(keys)
+        self.assertIsNone(result)
+
+    def test_no_fold_multiple_varying_positions(self):
+        keys = ["layer.0.expert.0.w", "layer.1.expert.1.w"]
+        result = _try_fold_src_keys(keys)
+        self.assertIsNone(result)
+
+    def test_single_key(self):
+        result = _try_fold_src_keys(["a.0.b"])
+        self.assertIsNone(result)
+
+    def test_empty(self):
+        result = _try_fold_src_keys([])
+        self.assertIsNone(result)
+
+
+class TestDescribeOps(unittest.TestCase):
+    def test_single_with_permute(self):
+        entry = AOAMappingEntry(
+            dst_key="a.weight",
+            dst_global_shape=(4, 8),
+            slice_mappings=[
+                AOASliceMapping(
+                    "b.weight",
+                    (slice(0, 4), slice(0, 8)),
+                    (slice(0, 4), slice(0, 8)),
+                    ["[1, 0]"],
+                )
+            ],
+        )
+        result = _describe_ops(entry)
+        self.assertIn("permute([1, 0])", result)
+
+    def test_concat_with_cast(self):
+        entry = AOAMappingEntry(
+            dst_key="a.weight",
+            dst_global_shape=(8, 4),
+            slice_mappings=[
+                AOASliceMapping(
+                    "b.weight",
+                    (slice(0, 4), slice(0, 4)),
+                    (slice(0, 4), slice(0, 4)),
+                    ["bfloat16"],
+                ),
+                AOASliceMapping(
+                    "c.weight",
+                    (slice(0, 4), slice(0, 4)),
+                    (slice(4, 8), slice(0, 4)),
+                    ["bfloat16"],
+                ),
+            ],
+        )
+        result = _describe_ops(entry)
+        self.assertIn("concat", result)
+        self.assertIn("cast(bfloat16)", result)
+
+    def test_no_ops(self):
+        entry = AOAMappingEntry(
+            dst_key="a.weight",
+            dst_global_shape=(4, 8),
+            slice_mappings=[
+                AOASliceMapping(
+                    "b.weight",
+                    (slice(0, 4), slice(0, 8)),
+                    (slice(0, 4), slice(0, 8)),
+                    None,
+                )
+            ],
+        )
+        result = _describe_ops(entry)
+        self.assertEqual(result, "")
+
+    def test_empty_slice_mappings(self):
+        entry = AOAMappingEntry(
+            dst_key="a.weight", dst_global_shape=(4,), slice_mappings=[]
+        )
+        result = _describe_ops(entry)
+        self.assertEqual(result, "")
+
+
+class TestClassifyMappings(unittest.TestCase):
+    def _make_entry(self, dst_key, src_key, pp=None, multi_src=False):
+        if multi_src:
+            sms = [
+                AOASliceMapping(src_key, (slice(0, 4),), (slice(0, 4),), pp),
+                AOASliceMapping(
+                    src_key + ".2", (slice(0, 4),), (slice(4, 8),), pp
+                ),
+            ]
+        else:
+            sms = [AOASliceMapping(src_key, (slice(0, 4),), (slice(0, 4),), pp)]
+        return AOAMappingEntry(
+            dst_key=dst_key, dst_global_shape=(8,), slice_mappings=sms
+        )
+
+    def test_rename_only(self):
+        entry = self._make_entry("model.layers.2.w", "model.layers.0.w")
+        rename, transform, struct = _classify_mappings([entry])
+        self.assertEqual(len(rename), 1)
+        self.assertEqual(len(transform), 0)
+        self.assertEqual(len(struct), 0)
+
+    def test_with_transform(self):
+        entry = self._make_entry(
+            "model.layers.2.w", "model.layers.0.w", ["[1, 0]"]
+        )
+        rename, transform, struct = _classify_mappings([entry])
+        self.assertEqual(len(rename), 0)
+        self.assertEqual(len(transform), 1)
+        self.assertEqual(len(struct), 0)
+
+    def test_structural_multi_src(self):
+        entry = self._make_entry(
+            "model.layers.2.qkv", "model.layers.0.q", multi_src=True
+        )
+        rename, transform, struct = _classify_mappings([entry])
+        self.assertEqual(len(rename), 0)
+        self.assertEqual(len(transform), 0)
+        self.assertEqual(len(struct), 1)
+
+    def test_structural_different_pattern(self):
+        entry = self._make_entry("model.decoder.0.w", "model.encoder.0.w")
+        rename, transform, struct = _classify_mappings([entry])
+        self.assertEqual(len(struct), 1)
+
+
+class TestGroupBySignature(unittest.TestCase):
+    def test_same_signature_grouped(self):
+        entries = []
+        for i in range(4):
+            entries.append(
+                AOAMappingEntry(
+                    dst_key=f"model.layers.{i}.w",
+                    dst_global_shape=(4,),
+                    slice_mappings=[
+                        AOASliceMapping(
+                            f"src.layers.{i}.w",
+                            (slice(0, 4),),
+                            (slice(0, 4),),
+                            ["[1, 0]"],
+                        )
+                    ],
+                )
+            )
+        groups = _group_by_signature(entries)
+        self.assertEqual(len(groups), 1)
+        self.assertEqual(len(next(iter(groups.values()))), 4)
+
+    def test_different_signatures(self):
+        e1 = AOAMappingEntry(
+            dst_key="model.layers.0.w",
+            dst_global_shape=(4,),
+            slice_mappings=[
+                AOASliceMapping(
+                    "src.layers.0.w", (slice(0, 4),), (slice(0, 4),), None
+                )
+            ],
+        )
+        e2 = AOAMappingEntry(
+            dst_key="model.layers.0.qkv",
+            dst_global_shape=(12,),
+            slice_mappings=[
+                AOASliceMapping(
+                    "src.layers.0.q", (slice(0, 4),), (slice(0, 4),), None
+                ),
+                AOASliceMapping(
+                    "src.layers.0.k", (slice(0, 4),), (slice(4, 8),), None
+                ),
+            ],
+        )
+        groups = _group_by_signature([e1, e2])
+        self.assertEqual(len(groups), 2)
+
+
+class TestGetSignature(unittest.TestCase):
+    def test_digits_normalized(self):
+        entry = AOAMappingEntry(
+            dst_key="model.layers.5.weight",
+            dst_global_shape=(4,),
+            slice_mappings=[
+                AOASliceMapping(
+                    "src.layers.5.weight", (slice(0, 4),), (slice(0, 4),), None
+                )
+            ],
+        )
+        sig = _get_signature(entry)
+        self.assertIn("{N}", sig)
+        self.assertNotIn("5", sig)
+
+
+class TestFormatPatternGroups(unittest.TestCase):
+    def test_basic_output(self):
+        entries = [
+            AOAMappingEntry(
+                dst_key="model.layers.0.w",
+                dst_global_shape=(4, 8),
+                slice_mappings=[
+                    AOASliceMapping(
+                        "src.layers.0.w",
+                        (slice(0, 4), slice(0, 8)),
+                        (slice(0, 4), slice(0, 8)),
+                        ["[1, 0]"],
+                    )
+                ],
+            )
+        ]
+        groups = {"sig1": entries}
+        lines, next_idx = _format_pattern_groups(groups, "test", 1)
+        self.assertTrue(any("Pattern #1" in l for l in lines))
+        self.assertEqual(next_idx, 2)
+
+    def test_numbering_continues(self):
+        e1 = [
+            AOAMappingEntry(
+                dst_key="a.0.w",
+                dst_global_shape=(4,),
+                slice_mappings=[
+                    AOASliceMapping(
+                        "b.0.w", (slice(0, 4),), (slice(0, 4),), None
+                    )
+                ],
+            )
+        ]
+        e2 = [
+            AOAMappingEntry(
+                dst_key="c.0.w",
+                dst_global_shape=(4,),
+                slice_mappings=[
+                    AOASliceMapping(
+                        "d.0.w", (slice(0, 4),), (slice(0, 4),), None
+                    )
+                ],
+            )
+        ]
+        groups = {"sig1": e1, "sig2": e2}
+        lines, next_idx = _format_pattern_groups(groups, "test", 5)
+        self.assertEqual(next_idx, 7)
+
+    def test_max_patterns_truncation(self):
+        import paddle.distributed.flex_checkpoint.dcp.key_validation as kv
+
+        old = kv._MAX_PATTERNS_SHOWN
+        kv._MAX_PATTERNS_SHOWN = 2
+        try:
+            groups = {}
+            for i in range(5):
+                groups[f"sig{i}"] = [
+                    AOAMappingEntry(
+                        dst_key=f"x.{i}.w",
+                        dst_global_shape=(4,),
+                        slice_mappings=[
+                            AOASliceMapping(
+                                f"y.{i}.w", (slice(0, 4),), (slice(0, 4),), None
+                            )
+                        ],
+                    )
+                ]
+            lines, _ = _format_pattern_groups(groups, "test", 1)
+            self.assertTrue(any("more" in l for l in lines))
+        finally:
+            kv._MAX_PATTERNS_SHOWN = old
+
+
+class TestAppendSrcLines(unittest.TestCase):
+    def test_few_srcs(self):
+        sms = [
+            AOASliceMapping("a.w", (slice(0, 4),), (slice(0, 4),), None),
+            AOASliceMapping("b.w", (slice(0, 4),), (slice(4, 8),), None),
+        ]
+        lines = []
+        _append_src_lines(lines, sms)
+        self.assertEqual(len(lines), 2)
+        self.assertIn("SRC:", lines[0])
+        self.assertIn("+", lines[1])
+
+    def test_many_srcs_foldable(self):
+        sms = [
+            AOASliceMapping(
+                f"experts.{i}.w",
+                (slice(0, 4),),
+                (slice(i * 4, (i + 1) * 4),),
+                None,
+            )
+            for i in range(10)
+        ]
+        lines = []
+        _append_src_lines(lines, sms)
+        # Should fold into single line with ×N
+        self.assertTrue(any("\u00d7" in l for l in lines))
+
+    def test_many_srcs_not_foldable(self):
+        sms = [
+            AOASliceMapping(
+                "src_alpha.w", (slice(0, 4),), (slice(0, 4),), None
+            ),
+            AOASliceMapping("src_beta.w", (slice(0, 4),), (slice(4, 8),), None),
+            AOASliceMapping(
+                "src_gamma.w", (slice(0, 4),), (slice(8, 12),), None
+            ),
+            AOASliceMapping(
+                "src_delta.w", (slice(0, 4),), (slice(12, 16),), None
+            ),
+            AOASliceMapping(
+                "src_epsilon.w", (slice(0, 4),), (slice(16, 20),), None
+            ),
+            AOASliceMapping(
+                "src_zeta.w", (slice(0, 4),), (slice(20, 24),), None
+            ),
+        ]
+        lines = []
+        _append_src_lines(lines, sms)
+        # Should show first 2, ..., last 1
+        self.assertTrue(any("more" in l for l in lines))
+
+
+class TestGroupKeysAdaptive(unittest.TestCase):
+    def test_basic_grouping(self):
+        keys = [
+            "model.layers.0.weight",
+            "model.layers.1.weight",
+            "model.layers.2.weight",
+            "model.embed.weight",
+        ]
+        groups = _group_keys_adaptive(keys)
+        self.assertEqual(len(groups), 2)
+        # layers.* grouped together
+        layer_group = [g for g in groups.values() if len(g) == 3]
+        self.assertEqual(len(layer_group), 1)
+
+    def test_no_digits(self):
+        keys = ["model.weight", "model.bias"]
+        groups = _group_keys_adaptive(keys)
+        self.assertEqual(len(groups), 2)
+
+
+class TestFormatKeyList(unittest.TestCase):
+    def test_few_keys(self):
+        keys = {"a.w", "b.w", "c.w"}
+        lines = _format_key_list(keys)
+        self.assertEqual(len(lines), 3)
+
+    def test_many_keys_grouped(self):
+        keys = {f"model.layers.{i}.weight" for i in range(100)}
+        lines = _format_key_list(keys)
+        # Should be grouped and folded
+        self.assertTrue(len(lines) < 100)
+        self.assertTrue(any("[" in l for l in lines))
+
+    def test_empty(self):
+        lines = _format_key_list(set())
+        self.assertEqual(lines, [])
+
+
+class TestEmit(unittest.TestCase):
+    @patch("paddle.distributed.flex_checkpoint.dcp.key_validation.logger")
+    def test_normal_output(self, mock_logger):
+        lines = ["line1", "line2", "line3"]
+        _emit(lines)
+        self.assertEqual(mock_logger.info.call_count, 3)
+
+    @patch("paddle.distributed.flex_checkpoint.dcp.key_validation.logger")
+    def test_truncation(self, mock_logger):
+        import paddle.distributed.flex_checkpoint.dcp.key_validation as kv
+
+        old_max = kv._MAX_TOTAL_LINES
+        kv._MAX_TOTAL_LINES = 5
+        try:
+            lines = ["x"] * 20
+            _emit(lines)
+            # 5 lines + 1 truncation msg = 6
+            self.assertEqual(mock_logger.info.call_count, 6)
+        finally:
+            kv._MAX_TOTAL_LINES = old_max
+
+
+class TestPrintStandardReport(unittest.TestCase):
+    @patch("paddle.distributed.flex_checkpoint.dcp.key_validation._emit")
+    def test_all_matched(self, mock_emit):
+        result = KeyValidationResult()
+        _print_standard_report(result, "/tmp/ckpt", 100)
+        lines = mock_emit.call_args[0][0]
+        self.assertTrue(any("[OK]" in l for l in lines))
+
+    @patch("paddle.distributed.flex_checkpoint.dcp.key_validation._emit")
+    def test_with_missing_and_unexpected(self, mock_emit):
+        result = KeyValidationResult(
+            missing_keys={"a.w", "b.w"},
+            unexpected_keys={"c.w"},
+            shape_mismatches=[ShapeMismatchInfo("d.w", (4, 8), (4, 16))],
+        )
+        _print_standard_report(result, "/tmp/ckpt", 100)
+        lines = mock_emit.call_args[0][0]
+        self.assertTrue(any("Missing" in l for l in lines))
+        self.assertTrue(any("Unexpected" in l for l in lines))
+        self.assertTrue(any("Shape" in l for l in lines))
+        self.assertTrue(any("Matched: 98/100" in l for l in lines))
+
+    @patch("paddle.distributed.flex_checkpoint.dcp.key_validation._emit")
+    def test_shape_mismatch_truncation(self, mock_emit):
+        import paddle.distributed.flex_checkpoint.dcp.key_validation as kv
+
+        old = kv._MAX_SHAPE_MISMATCHES
+        kv._MAX_SHAPE_MISMATCHES = 2
+        try:
+            mismatches = [
+                ShapeMismatchInfo(f"k{i}", (4,), (8,)) for i in range(5)
+            ]
+            result = KeyValidationResult(
+                missing_keys={"x"}, shape_mismatches=mismatches
+            )
+            _print_standard_report(result, "/tmp/ckpt", 10)
+            lines = mock_emit.call_args[0][0]
+            self.assertTrue(any("and 3 more" in l for l in lines))
+        finally:
+            kv._MAX_SHAPE_MISMATCHES = old
+
+
+class TestPrintAoaReport(unittest.TestCase):
+    @patch("paddle.distributed.flex_checkpoint.dcp.key_validation._emit")
+    def test_all_resolved(self, mock_emit):
+        mappings = [
+            AOAMappingEntry(
+                dst_key="a.w",
+                dst_global_shape=(4,),
+                slice_mappings=[
+                    AOASliceMapping("b.w", (slice(0, 4),), (slice(0, 4),), None)
+                ],
+                is_identity=False,
+            ),
+        ]
+        result = KeyValidationResult()
+        _print_aoa_report(result, mappings, set(), "/tmp/ckpt")
+        lines = mock_emit.call_args[0][0]
+        self.assertTrue(any("[OK]" in l for l in lines))
+
+    @patch("paddle.distributed.flex_checkpoint.dcp.key_validation._emit")
+    def test_with_missing(self, mock_emit):
+        mappings = [
+            AOAMappingEntry(
+                "a.w",
+                (4,),
+                [AOASliceMapping("b.w", (slice(0, 4),), (slice(0, 4),), None)],
+            )
+        ]
+        result = KeyValidationResult(
+            missing_keys={"c.w"}, unexpected_keys={"d.w"}
+        )
+        _print_aoa_report(result, mappings, {"removed.w"}, "/tmp/ckpt")
+        lines = mock_emit.call_args[0][0]
+        self.assertTrue(any("Missing" in l for l in lines))
+        self.assertTrue(any("Unexpected" in l for l in lines))
+        self.assertTrue(any("Removed" in l for l in lines))
+
+    @patch("paddle.distributed.flex_checkpoint.dcp.key_validation._emit")
+    def test_randomly_initialized_keys(self, mock_emit):
+        mappings = []
+        result = KeyValidationResult(
+            randomly_initialized_keys={"init.w", "init.b"}
+        )
+        _print_aoa_report(result, mappings, set(), "/tmp/ckpt")
+        lines = mock_emit.call_args[0][0]
+        self.assertTrue(any("Initialized (2)" in l for l in lines))
+
+    @patch("paddle.distributed.flex_checkpoint.dcp.key_validation._emit")
+    def test_removed_keys_truncation(self, mock_emit):
+        mappings = []
+        removed = {f"removed.key.{i}" for i in range(10)}
+        result = KeyValidationResult()
+        _print_aoa_report(result, mappings, removed, "/tmp/ckpt")
+        lines = mock_emit.call_args[0][0]
+        self.assertTrue(any("more" in l for l in lines))
+
+
+class TestBuildAoaMappings(unittest.TestCase):
+    def test_basic(self):
+        engine = MagicMock()
+        td1 = MagicMock()
+        td1.shape = [4, 8]
+        td1.slices = [
+            (
+                "src.w",
+                (slice(0, 4), slice(0, 8)),
+                (slice(0, 4), slice(0, 8)),
+                None,
+            )
+        ]
+        td2 = MagicMock()
+        td2.shape = [8, 8]
+        td2.slices = [
+            (
+                "src.q",
+                (slice(0, 4), slice(0, 8)),
+                (slice(0, 4), slice(0, 8)),
+                ["[1, 0]"],
+            ),
+            (
+                "src.k",
+                (slice(0, 4), slice(0, 8)),
+                (slice(4, 8), slice(0, 8)),
+                ["[1, 0]"],
+            ),
+        ]
+        ov = MagicMock()
+        ov.items.return_value = sorted({"dst.qkv": td2, "dst.w": td1}.items())
+        engine.output_vars = ov
+
+        results = _build_aoa_mappings(engine)
+        self.assertEqual(len(results), 2)
+        qkv = next(r for r in results if r.dst_key == "dst.qkv")
+        self.assertFalse(qkv.is_identity)
+        self.assertEqual(len(qkv.slice_mappings), 2)
+
+    def test_identity_detection(self):
+        engine = MagicMock()
+        td = MagicMock()
+        td.shape = [4, 8]
+        td.slices = [
+            (
+                "same.key",
+                (slice(0, 4), slice(0, 8)),
+                (slice(0, 4), slice(0, 8)),
+                None,
+            )
+        ]
+        ov = MagicMock()
+        ov.items.return_value = [("same.key", td)]
+        engine.output_vars = ov
+
+        results = _build_aoa_mappings(engine)
+        self.assertEqual(len(results), 1)
+        self.assertTrue(results[0].is_identity)
+
+    def test_none_tensor_desc_skipped(self):
+        engine = MagicMock()
+        ov = MagicMock()
+        ov.items.return_value = [("a", None), ("b", None)]
+        engine.output_vars = ov
+        results = _build_aoa_mappings(engine)
+        self.assertEqual(len(results), 0)
+
+
+class TestValidateAndReportKeysStandard(unittest.TestCase):
+    def _make_metadata(self, keys_shapes):
+        """keys_shapes: dict of {key: shape_tuple}"""
+        storage_metadata = {}
+        state_dict_metadata = {}
+        for key, shape in keys_shapes.items():
+            idx = LocalTensorIndex(
+                tensor_key=key,
+                global_offset=tuple([0] * len(shape)),
+                replica_id=0,
+            )
+            storage_metadata[idx] = f"{key}.distcp"
+            state_dict_metadata[key] = [
+                LocalTensorMetadata(
+                    global_offset=tuple([0] * len(shape)),
+                    local_shape=shape,
+                    dtype="float32",
+                    global_shape=shape,
+                )
+            ]
+        return Metadata(
+            state_dict_metadata=state_dict_metadata,
+            storage_metadata=storage_metadata,
+        )
+
+    @patch("paddle.distributed.get_rank", return_value=0)
+    @patch("paddle.distributed.flex_checkpoint.dcp.key_validation._emit")
+    def test_all_match(self, mock_emit, mock_rank):
+        metadata = self._make_metadata({"w1": (4, 8), "w2": (4, 8)})
+        state_dict = {
+            "w1": paddle.zeros([4, 8]),
+            "w2": paddle.zeros([4, 8]),
+        }
+        result = validate_and_report_keys_standard(
+            [metadata], {"w1", "w2"}, None, False, "/tmp/ckpt", state_dict
+        )
+        self.assertEqual(len(result.missing_keys), 0)
+        self.assertEqual(len(result.unexpected_keys), 0)
+
+    @patch("paddle.distributed.get_rank", return_value=0)
+    @patch("paddle.distributed.flex_checkpoint.dcp.key_validation._emit")
+    def test_missing_keys(self, mock_emit, mock_rank):
+        metadata = self._make_metadata({"w1": (4,)})
+        state_dict = {
+            "w1": paddle.zeros([4]),
+            "w2": paddle.zeros([4]),
+        }
+        result = validate_and_report_keys_standard(
+            [metadata], {"w1", "w2"}, None, False, "/tmp/ckpt", state_dict
+        )
+        self.assertIn("w2", result.missing_keys)
+
+    @patch("paddle.distributed.get_rank", return_value=0)
+    @patch("paddle.distributed.flex_checkpoint.dcp.key_validation._emit")
+    def test_unexpected_keys(self, mock_emit, mock_rank):
+        metadata = self._make_metadata({"w1": (4,), "w2": (4,), "w3": (4,)})
+        state_dict = {"w1": paddle.zeros([4])}
+        result = validate_and_report_keys_standard(
+            [metadata], {"w1"}, None, False, "/tmp/ckpt", state_dict
+        )
+        self.assertIn("w2", result.unexpected_keys)
+        self.assertIn("w3", result.unexpected_keys)
+
+    @patch("paddle.distributed.get_rank", return_value=0)
+    @patch("paddle.distributed.flex_checkpoint.dcp.key_validation._emit")
+    def test_shape_mismatch(self, mock_emit, mock_rank):
+        metadata = self._make_metadata({"w1": (4, 8)})
+        state_dict = {"w1": paddle.zeros([4, 16])}
+        result = validate_and_report_keys_standard(
+            [metadata], {"w1"}, None, False, "/tmp/ckpt", state_dict
+        )
+        self.assertEqual(len(result.shape_mismatches), 1)
+        self.assertEqual(result.shape_mismatches[0].src_global_shape, (4, 8))
+        self.assertEqual(result.shape_mismatches[0].dst_global_shape, (4, 16))
+
+    @patch("paddle.distributed.get_rank", return_value=0)
+    @patch("paddle.distributed.flex_checkpoint.dcp.key_validation._emit")
+    def test_replica_id_filtered(self, mock_emit, mock_rank):
+        """Keys with replica_id != 0 should be filtered out."""
+        storage_metadata = {
+            LocalTensorIndex(
+                tensor_key="w1", global_offset=(0,), replica_id=0
+            ): "f1",
+            LocalTensorIndex(
+                tensor_key="w2", global_offset=(0,), replica_id=1
+            ): "f2",
+        }
+        metadata = Metadata(
+            state_dict_metadata={
+                "w1": [LocalTensorMetadata((0,), (4,), "float32", (4,))]
+            },
+            storage_metadata=storage_metadata,
+        )
+        state_dict = {"w1": paddle.zeros([4])}
+        result = validate_and_report_keys_standard(
+            [metadata], {"w1"}, None, False, "/tmp/ckpt", state_dict
+        )
+        self.assertEqual(len(result.unexpected_keys), 0)
+
+    @patch(
+        "paddle.distributed.flex_checkpoint.dcp.key_validation._get_rank",
+        return_value=1,
+    )
+    @patch("paddle.distributed.flex_checkpoint.dcp.key_validation._emit")
+    @patch("paddle.distributed.all_gather_object")
+    def test_non_rank0_no_print(self, mock_gather, mock_emit, mock_rank):
+        metadata = self._make_metadata({"w1": (4,)})
+        state_dict = {"w1": paddle.zeros([4])}
+
+        def gather_side_effect(out_list, obj, group=None):
+            out_list.clear()
+            out_list.append(obj)
+
+        mock_gather.side_effect = gather_side_effect
+        validate_and_report_keys_standard(
+            [metadata], {"w1"}, None, True, "/tmp/ckpt", state_dict
+        )
+        mock_emit.assert_not_called()
+
+
+class TestValidateAndReportKeysAoa(unittest.TestCase):
+    def _make_mock_engine(self):
+        engine = MagicMock()
+        td1 = MagicMock()
+        td1.shape = [4, 8]
+        td1.slices = [
+            (
+                "src.w1",
+                (slice(0, 4), slice(0, 8)),
+                (slice(0, 4), slice(0, 8)),
+                None,
+            )
+        ]
+        td2 = MagicMock()
+        td2.shape = [8, 8]
+        td2.slices = [
+            (
+                "src.q",
+                (slice(0, 4), slice(0, 8)),
+                (slice(0, 4), slice(0, 8)),
+                ["[1, 0]"],
+            ),
+            (
+                "src.k",
+                (slice(0, 4), slice(0, 8)),
+                (slice(4, 8), slice(0, 8)),
+                ["[1, 0]"],
+            ),
+        ]
+        # output_vars: need .items() for _build_aoa_mappings and iteration for values()
+        ov = MagicMock()
+        ov.items.return_value = sorted({"dst.w1": td1, "dst.qkv": td2}.items())
+        ov.values.return_value = [td1, td2]
+        ov.__iter__ = lambda self: iter({"dst.w1": td1, "dst.qkv": td2})
+        ov.__getitem__ = lambda self, k: {"dst.w1": td1, "dst.qkv": td2}[k]
+        engine.output_vars = ov
+        engine.need_add_output_vars = ["dst.init"]
+        engine.need_remove_input_vars = ["src.removed"]
+        engine.input_vars = MagicMock()
+        engine.input_vars.keys.return_value = [
+            "src.w1",
+            "src.q",
+            "src.k",
+            "src.removed",
+            "src.leftover",
+        ]
+        engine.context = MagicMock()
+        engine.context.get_all_dst_state_keys.return_value = {
+            "dst.w1",
+            "dst.qkv",
+            "dst.init",
+        }
+        return engine
+
+    @patch("paddle.distributed.get_rank", return_value=0)
+    @patch("paddle.distributed.flex_checkpoint.dcp.key_validation._emit")
+    def test_all_resolved(self, mock_emit, mock_rank):
+        engine = self._make_mock_engine()
+        metadata = MagicMock()
+        result = validate_and_report_keys_aoa(engine, metadata, "/tmp/ckpt")
+        # dst.w1 and dst.qkv are covered; dst.init is randomly initialized
+        self.assertEqual(len(result.missing_keys), 0)
+        # src.leftover not consumed and not removed
+        self.assertIn("src.leftover", result.unexpected_keys)
+        self.assertIn("dst.init", result.randomly_initialized_keys)
+
+    @patch("paddle.distributed.get_rank", return_value=0)
+    @patch("paddle.distributed.flex_checkpoint.dcp.key_validation._emit")
+    def test_truly_missing(self, mock_emit, mock_rank):
+        engine = self._make_mock_engine()
+        # Add a dst key that is NOT covered
+        engine.context.get_all_dst_state_keys = lambda: {
+            "dst.w1",
+            "dst.qkv",
+            "dst.init",
+            "dst.missing",
+        }
+        metadata = MagicMock()
+        result = validate_and_report_keys_aoa(engine, metadata, "/tmp/ckpt")
+        self.assertIn("dst.missing", result.missing_keys)
+
+    @patch("paddle.distributed.get_rank", return_value=1)
+    @patch("paddle.distributed.flex_checkpoint.dcp.key_validation._emit")
+    def test_non_rank0_no_print(self, mock_emit, mock_rank):
+        engine = self._make_mock_engine()
+        metadata = MagicMock()
+        validate_and_report_keys_aoa(engine, metadata, "/tmp/ckpt")
+        mock_emit.assert_not_called()
+
+
+class TestColorHelpers(unittest.TestCase):
+    def test_no_color(self):
+        from paddle.distributed.flex_checkpoint.dcp.key_validation import _C
+
+        self.assertEqual(_C.green("test"), "test")
+        self.assertEqual(_C.yellow("test"), "test")
+        self.assertEqual(_C.red("test"), "test")
+        self.assertEqual(_C.cyan("test"), "test")
+
+
+if __name__ == "__main__":
+    unittest.main()

test/flex_checkpoint/test_load_state_dict_transpose.py

@@ -0,0 +1,63 @@
+# 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 unittest
+
+import collective.test_communication_api_base as test_base
+
+
+class TestLoadStateDictTranspose(test_base.CommunicationTestDistBase):
+    def setUp(self):
+        super().setUp(num_of_devices=2)
+
+    def test_metadata(self):
+        envs = {
+            "aoa_statements": "linear.weight^T -> linear.weight",
+        }
+        self.run_test_case(
+            "load_state_dict_transpose_logic.py",
+            user_defined_envs=envs,
+        )
+
+
+class TestLoadStateDictCast(test_base.CommunicationTestDistBase):
+    def setUp(self):
+        super().setUp(num_of_devices=2)
+
+    def test_cast(self):
+        envs = {
+            "aoa_statements": 'linear.weight -> linear.weight, dtype="float16"',
+        }
+        self.run_test_case(
+            "load_state_dict_cast_logic.py",
+            user_defined_envs=envs,
+        )
+
+
+class TestLoadStateDictTransposeCast(test_base.CommunicationTestDistBase):
+    def setUp(self):
+        super().setUp(num_of_devices=2)
+
+    def test_transpose_cast(self):
+        envs = {
+            "aoa_statements": 'linear.weight^T -> linear.weight, dtype="float16"',
+        }
+        self.run_test_case(
+            "load_state_dict_transpose_cast_logic.py",
+            user_defined_envs=envs,
+        )
+
+
+if __name__ == "__main__":
+    unittest.main()

test/flex_checkpoint/test_macros.py

@@ -0,0 +1,580 @@
+# 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 unittest
+
+from paddle.distributed.flex_checkpoint.aoa.aoa_engine import (
+    AOAShardInfoContext,
+)
+from paddle.distributed.flex_checkpoint.aoa.lexer import Lexer
+from paddle.distributed.flex_checkpoint.aoa.macros import macro_registry
+from paddle.distributed.flex_checkpoint.dcp.sharded_weight import (
+    ShardedWeightDesc,
+)
+
+
+class MacroContext:
+    def __init__(self):
+        self.source_keys = {
+            "embed_tokens.weight",
+            "layers.1.mlp.gate_up_fused_proj.weight",
+            "layers.1.post_attention_layernorm.weight",
+            "layers.2.self_attn.qkv_proj.weight",
+            "layers.2.self_attn.o_proj.weight",
+            "layers.2.mlp.gate_up_fused_proj.weight",
+            "layers.2.mlp.down_proj.weight",
+            "layers.2.input_layernorm.weight",
+            "layers.1.mlp.gate_up_fused_proj.weight_test1",
+            "layers.2.post_attention_layernorm.weight",
+            "layers.1.experts.0.weight",
+            "layers.0.qkv_proj.weight",
+            "fused_qkv_old_test_name",
+            "layers.shared.qkv_proj.weight",
+            "layers.5.experts.0.up_gate_proj.weight",
+            "layers.5.experts.1.up_gate_proj.weight",
+            "layers.2.experts.0.weight",
+            "layers.2.experts.1.weight",
+            "layers.2.self_attn.qkv_proj.bias",
+            "layers.2.mlp.gate_up_fused_proj.bias",
+            "layers.3.experts.0.up_gate_proj.weight",
+            "layers.3.experts.1.up_gate_proj.weight",
+        }
+
+        self.dst_keys = {
+            "embed_tokens.weight",
+            "layers.0.self_attn.qkv_proj.weight",
+            "layers.0.self_attn.o_proj.weight",
+            "layers.0.mlp.gate_up_fused_proj.weight",
+            "layers.0.mlp.down_proj.weight",
+            "layers.0.input_layernorm.weight",
+            "layers.0.post_attention_layernorm.weight",
+            "layers.1.mlp.gate_up_fused_proj.weight",
+            "layers.1.mlp.gate_up_fused_proj.weight_test2",
+            "layers.1.post_attention_layernorm.weight",
+            "layers.0.experts.0.weight",
+            "layers.0.experts.1.weight",
+            "layers.1.experts.0.weight",
+            "layers.0.q_proj.weight",
+            "layers.0.k_proj.weight",
+            "layers.0.v_proj.weight",
+            "q_test_name",
+            "k_test_name",
+            "v_test_name",
+            "layers.0.shared.q_proj.weight",
+            "layers.0.shared.k_proj.weight",
+            "layers.0.shared.v_proj.weight",
+            "layers.1.shared.q_proj.weight",
+            "layers.1.shared.k_proj.weight",
+            "layers.1.shared.v_proj.weight",
+            "layers.5.experts.0.gate_proj.weight",
+            "layers.5.experts.1.gate_proj.weight",
+            "layers.5.experts.0.up_proj.weight",
+            "layers.5.experts.1.up_proj.weight",
+            "layers.2.self_attn.qkv_proj.weight",
+            "layers.2.self_attn.qkv_proj.bias",
+            "layers.2.mlp.gate_up_fused_proj.bias",
+            "layers.2.mlp.gate_up_fused_proj.weight",
+            "layers.3.experts.0.up_gate_proj.weight",
+            "layers.3.experts.1.up_gate_proj.weight",
+        }
+
+        # Build _ShardInfo mapping for AOAShardInfoContext based on existing keys
+        def make_shard_info(keys: set[str], num_shards: int):
+            shard_info: dict[str, list[ShardedWeightDesc]] = {}
+            for k in keys:
+                descs: list[ShardedWeightDesc] = []
+                for i in range(num_shards):
+                    descs.append(
+                        ShardedWeightDesc(
+                            key=k,
+                            local_shape=(1,),
+                            global_shape=(num_shards,),
+                            global_offset=(i,),
+                        )
+                    )
+                shard_info[k] = descs
+            return shard_info
+
+        self.source_state_shard_info = make_shard_info(self.source_keys, 2)
+        self.destination_state_shard_info = make_shard_info(self.dst_keys, 4)
+
+        self._ctx = AOAShardInfoContext(
+            source_state_shard_info=self.source_state_shard_info,
+            destination_state_shard_info=self.destination_state_shard_info,
+        )
+
+    def set_aoa_config_reverse(
+        self,
+    ):  # when aoa_config_reverse is True, the src and dst of AOAShardInfoContext are reversed
+        self._ctx = AOAShardInfoContext(
+            source_state_shard_info=self.destination_state_shard_info,
+            destination_state_shard_info=self.source_state_shard_info,
+        )
+        self._ctx.aoa_config_reverse = True
+
+
+def get_macro(macro_name):
+    for macro in macro_registry.macros:
+        if macro["name"] == macro_name:
+            return macro["func"]
+    raise ValueError(f"Macro '{macro_name}' not found.")
+
+
+class TestMacro(unittest.TestCase):
+    def setUp(self):
+        self.macro_func = None
+        self.source = None
+        self.expected_expanded = None
+
+    def macro_name(self):
+        raise NotImplementedError
+
+    def source_code(self):
+        raise NotImplementedError
+
+    def expected(self):
+        raise NotImplementedError
+
+    def start_macro_test(self, aoa_config_reverse: bool = False):
+        self.macro_func = get_macro(self.macro_name())
+        self.source = self.source_code()
+        self.expected_expanded = self.expected()
+        self.ctx = MacroContext()
+        if aoa_config_reverse:
+            self.ctx.set_aoa_config_reverse()
+            self.lexer = Lexer(self.ctx._ctx)
+            self.lexer.apply_macro(
+                self.source, get_macro("get_var_mapping_chain_macro")
+            )
+        else:
+            self.lexer = Lexer(self.ctx._ctx)
+        actual_expanded = self.lexer.apply_macro(self.source, self.macro_func)
+        self.assertEqual(actual_expanded, self.expected_expanded)
+
+
+class TestStarMacro(TestMacro):
+    def macro_name(self):
+        return "star_macro"
+
+    def source_code(self):
+        return "layers.2.experts.*.weight -> fused_experts, axis = 1"
+
+    def expected(self):
+        return [
+            'layers.2.experts.0.weight,layers.2.experts.1.weight->fused_experts,axis=1\n'
+        ]
+
+    def test(self):
+        self.start_macro_test()
+
+
+class TestLayerIdMacro(TestMacro):
+    def macro_name(self):
+        return "id_macro"
+
+    def source_code(self):
+        return "layers.$LAYER_ID.qkv_proj.weight->layers.$LAYER_ID.q_proj.weight,layer.$LAYER_ID.k_proj.weight,layer.$LAYER_ID.v_proj.weight\n"
+
+    def expected(self):
+        return [
+            'layers.0.qkv_proj.weight->layers.0.q_proj.weight,layer.0.k_proj.weight,layer.0.v_proj.weight\n',
+        ]
+
+    def test(self):
+        self.start_macro_test()
+
+
+class Test_expert_id_Macro(TestMacro):
+    def macro_name(self):
+        return "id_macro"
+
+    def source_code(self):
+        return "layers.5.experts.$EXPERT_ID.up_gate_proj.weight -> layers.5.experts.$EXPERT_ID.gate_proj.weight, layers.5.experts.$EXPERT_ID.up_proj.weight"
+
+    def expected(self):
+        return [
+            'layers.5.experts.0.up_gate_proj.weight->layers.5.experts.0.gate_proj.weight,layers.5.experts.0.up_proj.weight\n',
+            'layers.5.experts.1.up_gate_proj.weight->layers.5.experts.1.gate_proj.weight,layers.5.experts.1.up_proj.weight\n',
+        ]
+
+    def test(self):
+        self.start_macro_test()
+
+
+class Test_ID_macro_reverse(TestMacro):
+    def macro_name(self):
+        return "id_macro"
+
+    def source_code(self):
+        return "layers.5.experts.$EXPERT_ID.up_gate_proj.weight -> layers.5.experts.$EXPERT_ID.gate_proj.weight, layers.5.experts.$EXPERT_ID.up_proj.weight"
+
+    def expected(self):
+        return [
+            'layers.5.experts.0.up_gate_proj.weight->layers.5.experts.0.gate_proj.weight,layers.5.experts.0.up_proj.weight\n',
+            'layers.5.experts.1.up_gate_proj.weight->layers.5.experts.1.gate_proj.weight,layers.5.experts.1.up_proj.weight\n',
+        ]
+
+    def test(self):
+        self.start_macro_test(aoa_config_reverse=True)
+
+
+class TestFusedQkvOldMacro(TestMacro):
+    def macro_name(self):
+        return "fused_qkv_old_macro"
+
+    def source_code(self):
+        return "layers.2.self_attn.qkv_proj.weight -> layers.2.self_attn.qkv_proj.weight, fused_qkv_old, num_heads = 8, num_key_value_groups = 4"
+
+    def expected(self):
+        return [
+            'layers.2.self_attn.qkv_proj.weight -> fused_qkv_old_tmp.Q_0,fused_qkv_old_tmp.Q_1,fused_qkv_old_tmp.Q_2,fused_qkv_old_tmp.Q_3,fused_qkv_old_tmp.K_0,fused_qkv_old_tmp.K_1,fused_qkv_old_tmp.V_0,fused_qkv_old_tmp.V_1,fused_qkv_old_tmp.Q_4,fused_qkv_old_tmp.Q_5,fused_qkv_old_tmp.Q_6,fused_qkv_old_tmp.Q_7,fused_qkv_old_tmp.K_2,fused_qkv_old_tmp.K_3,fused_qkv_old_tmp.V_2,fused_qkv_old_tmp.V_3, axis=1',
+            'fused_qkv_old_tmp.Q_0,fused_qkv_old_tmp.Q_1,fused_qkv_old_tmp.K_0,fused_qkv_old_tmp.V_0,fused_qkv_old_tmp.Q_2,fused_qkv_old_tmp.Q_3,fused_qkv_old_tmp.K_1,fused_qkv_old_tmp.V_1,fused_qkv_old_tmp.Q_4,fused_qkv_old_tmp.Q_5,fused_qkv_old_tmp.K_2,fused_qkv_old_tmp.V_2,fused_qkv_old_tmp.Q_6,fused_qkv_old_tmp.Q_7,fused_qkv_old_tmp.K_3,fused_qkv_old_tmp.V_3 -> layers.2.self_attn.qkv_proj.weight, axis=1',
+        ]
+
+    def test(self):
+        self.start_macro_test()
+
+
+class TestTransposeMacro(TestMacro):
+    def macro_name(self):
+        return "transpose_macro"
+
+    def source_code(self):
+        return (
+            "layers.2.mlp.down_proj.weight^T -> layers.2.mlp.down_proj.weight_T"
+        )
+
+    def expected(self):
+        return [
+            'layers.2.mlp.down_proj.weight -> layers.2.mlp.down_proj.weight_transpose_tmp, permute = "[]"',
+            'layers.2.mlp.down_proj.weight_transpose_tmp->layers.2.mlp.down_proj.weight_T\n',
+        ]
+
+    def test(self):
+        self.start_macro_test()
+
+
+class TestFusedQKVMacro(TestMacro):
+    def macro_name(self):
+        return "fused_qkv_macro"
+
+    def source_code(self):
+        return "layers.2.self_attn.qkv_proj.weight -> Q, K, V, fused_qkv, num_heads = 8, num_key_value_groups = 2"
+
+    def expected(self):
+        return [
+            'layers.2.self_attn.qkv_proj.weight -> Q0,Q1,Q2,Q3,K0,V0,Q4,Q5,Q6,Q7,K1,V1, axis=1',
+            'Q0,Q1,Q2,Q3,Q4,Q5,Q6,Q7 -> Q, axis=1',
+            'K0,K1 -> K, axis=1',
+            'V0,V1 -> V, axis=1',
+        ]
+
+    def test(self):
+        self.start_macro_test()
+
+
+class TestFusedQKVMacro2(TestMacro):
+    def macro_name(self):
+        return "fused_qkv_macro"
+
+    def source_code(self):
+        return "Q, K, V -> layers.2.self_attn.qkv_proj.weight, fused_qkv, num_heads = 8, num_key_value_groups = 8"
+
+    def expected(self):
+        return [
+            'Q -> Q0,Q1,Q2,Q3,Q4,Q5,Q6,Q7, axis=1',
+            'K -> K0,K1,K2,K3,K4,K5,K6,K7, axis=1',
+            'V -> V0,V1,V2,V3,V4,V5,V6,V7, axis=1',
+            'Q0,K0,V0,Q1,K1,V1,Q2,K2,V2,Q3,K3,V3,Q4,K4,V4,Q5,K5,V5,Q6,K6,V6,Q7,K7,V7 -> layers.2.self_attn.qkv_proj.weight, axis=1',
+        ]
+
+    def test(self):
+        self.start_macro_test()
+
+
+class TestFusedQkvOldMacro2(TestMacro):
+    def macro_name(self):
+        return "fused_qkv_old_macro"
+
+    def source_code(self):
+        return "Q,K,V -> layers.2.self_attn.qkv_proj.weight, fused_qkv_old, num_heads = 8, num_key_value_groups = 4"
+
+    def expected(self):
+        return [
+            'Q,K,V  ->  Q.K.V.tmp, axis=1',
+            'Q.K.V.tmp -> fused_qkv_old_tmp.Q_0,fused_qkv_old_tmp.Q_1,fused_qkv_old_tmp.Q_2,fused_qkv_old_tmp.Q_3,fused_qkv_old_tmp.Q_4,fused_qkv_old_tmp.Q_5,fused_qkv_old_tmp.Q_6,fused_qkv_old_tmp.Q_7,fused_qkv_old_tmp.K_0,fused_qkv_old_tmp.K_1,fused_qkv_old_tmp.K_2,fused_qkv_old_tmp.K_3,fused_qkv_old_tmp.V_0,fused_qkv_old_tmp.V_1,fused_qkv_old_tmp.V_2,fused_qkv_old_tmp.V_3, axis=1',
+            'fused_qkv_old_tmp.Q_0,fused_qkv_old_tmp.Q_1,fused_qkv_old_tmp.K_0,fused_qkv_old_tmp.V_0,fused_qkv_old_tmp.Q_2,fused_qkv_old_tmp.Q_3,fused_qkv_old_tmp.K_1,fused_qkv_old_tmp.V_1,fused_qkv_old_tmp.Q_4,fused_qkv_old_tmp.Q_5,fused_qkv_old_tmp.K_2,fused_qkv_old_tmp.V_2,fused_qkv_old_tmp.Q_6,fused_qkv_old_tmp.Q_7,fused_qkv_old_tmp.K_3,fused_qkv_old_tmp.V_3 -> layers.2.self_attn.qkv_proj.weight, axis=1',
+        ]
+
+    def test(self):
+        self.start_macro_test()
+
+
+class TestFusedQkvOldMacro3(TestMacro):
+    def macro_name(self):
+        return "fused_qkv_old_macro"
+
+    def source_code(self):
+        return "fused_qkv_old_test_name -> q_test_name ,k_test_name, v_test_name, fused_qkv_old, num_heads = 8, num_key_value_groups = 4 "
+
+    def expected(self):
+        return [
+            'fused_qkv_old_test_name -> fused_qkv_old_tmp.Q_0,fused_qkv_old_tmp.Q_1,fused_qkv_old_tmp.Q_2,fused_qkv_old_tmp.Q_3,fused_qkv_old_tmp.K_0,fused_qkv_old_tmp.K_1,fused_qkv_old_tmp.V_0,fused_qkv_old_tmp.V_1,fused_qkv_old_tmp.Q_4,fused_qkv_old_tmp.Q_5,fused_qkv_old_tmp.Q_6,fused_qkv_old_tmp.Q_7,fused_qkv_old_tmp.K_2,fused_qkv_old_tmp.K_3,fused_qkv_old_tmp.V_2,fused_qkv_old_tmp.V_3, axis=1',
+            'fused_qkv_old_tmp.Q_0,fused_qkv_old_tmp.Q_1,fused_qkv_old_tmp.Q_2,fused_qkv_old_tmp.Q_3,fused_qkv_old_tmp.Q_4,fused_qkv_old_tmp.Q_5,fused_qkv_old_tmp.Q_6,fused_qkv_old_tmp.Q_7 -> q_test_name, axis=1',
+            'fused_qkv_old_tmp.K_0,fused_qkv_old_tmp.K_1,fused_qkv_old_tmp.K_2,fused_qkv_old_tmp.K_3 -> k_test_name, axis=1',
+            'fused_qkv_old_tmp.V_0,fused_qkv_old_tmp.V_1,fused_qkv_old_tmp.V_2,fused_qkv_old_tmp.V_3 -> v_test_name, axis=1',
+        ]
+
+    def test(self):
+        self.start_macro_test()
+
+
+class TestFusedQkvOldMacro4(TestMacro):
+    def macro_name(self):
+        return "fused_qkv_old_macro"
+
+    def source_code(self):
+        return "fused_qkv_old_test_name ->  layers.2.self_attn.qkv_proj.weight,fused_qkv_old, num_heads = 8, num_key_value_groups = 8 "
+
+    def expected(self):
+        return [
+            'fused_qkv_old_test_name -> fused_qkv_old_tmp.Q_0,fused_qkv_old_tmp.Q_1,fused_qkv_old_tmp.Q_2,fused_qkv_old_tmp.Q_3,fused_qkv_old_tmp.K_0,fused_qkv_old_tmp.K_1,fused_qkv_old_tmp.K_2,fused_qkv_old_tmp.K_3,fused_qkv_old_tmp.V_0,fused_qkv_old_tmp.V_1,fused_qkv_old_tmp.V_2,fused_qkv_old_tmp.V_3,fused_qkv_old_tmp.Q_4,fused_qkv_old_tmp.Q_5,fused_qkv_old_tmp.Q_6,fused_qkv_old_tmp.Q_7,fused_qkv_old_tmp.K_4,fused_qkv_old_tmp.K_5,fused_qkv_old_tmp.K_6,fused_qkv_old_tmp.K_7,fused_qkv_old_tmp.V_4,fused_qkv_old_tmp.V_5,fused_qkv_old_tmp.V_6,fused_qkv_old_tmp.V_7, axis=1',
+            'fused_qkv_old_tmp.Q_0,fused_qkv_old_tmp.Q_1,fused_qkv_old_tmp.K_0,fused_qkv_old_tmp.K_1,fused_qkv_old_tmp.V_0,fused_qkv_old_tmp.V_1,fused_qkv_old_tmp.Q_2,fused_qkv_old_tmp.Q_3,fused_qkv_old_tmp.K_2,fused_qkv_old_tmp.K_3,fused_qkv_old_tmp.V_2,fused_qkv_old_tmp.V_3,fused_qkv_old_tmp.Q_4,fused_qkv_old_tmp.Q_5,fused_qkv_old_tmp.K_4,fused_qkv_old_tmp.K_5,fused_qkv_old_tmp.V_4,fused_qkv_old_tmp.V_5,fused_qkv_old_tmp.Q_6,fused_qkv_old_tmp.Q_7,fused_qkv_old_tmp.K_6,fused_qkv_old_tmp.K_7,fused_qkv_old_tmp.V_6,fused_qkv_old_tmp.V_7 -> layers.2.self_attn.qkv_proj.weight, axis=1',
+        ]
+
+    def test(self):
+        self.start_macro_test()
+
+
+class TestFusedQkvOldMacro5(TestMacro):
+    def macro_name(self):
+        return "fused_qkv_old_macro"
+
+    def source_code(self):
+        return "layers.2.self_attn.qkv_proj.bias -> layers.2.self_attn.qkv_proj.bias, fused_qkv_old, num_heads = 8, num_key_value_groups = 4, axis = 0"
+
+    def expected(self):
+        return [
+            'layers.2.self_attn.qkv_proj.bias -> fused_qkv_old_tmp.Q_0,fused_qkv_old_tmp.Q_1,fused_qkv_old_tmp.Q_2,fused_qkv_old_tmp.Q_3,fused_qkv_old_tmp.K_0,fused_qkv_old_tmp.K_1,fused_qkv_old_tmp.V_0,fused_qkv_old_tmp.V_1,fused_qkv_old_tmp.Q_4,fused_qkv_old_tmp.Q_5,fused_qkv_old_tmp.Q_6,fused_qkv_old_tmp.Q_7,fused_qkv_old_tmp.K_2,fused_qkv_old_tmp.K_3,fused_qkv_old_tmp.V_2,fused_qkv_old_tmp.V_3, axis=0',
+            'fused_qkv_old_tmp.Q_0,fused_qkv_old_tmp.Q_1,fused_qkv_old_tmp.K_0,fused_qkv_old_tmp.V_0,fused_qkv_old_tmp.Q_2,fused_qkv_old_tmp.Q_3,fused_qkv_old_tmp.K_1,fused_qkv_old_tmp.V_1,fused_qkv_old_tmp.Q_4,fused_qkv_old_tmp.Q_5,fused_qkv_old_tmp.K_2,fused_qkv_old_tmp.V_2,fused_qkv_old_tmp.Q_6,fused_qkv_old_tmp.Q_7,fused_qkv_old_tmp.K_3,fused_qkv_old_tmp.V_3 -> layers.2.self_attn.qkv_proj.bias, axis=0',
+        ]
+
+    def test(self):
+        self.start_macro_test()
+
+
+class TestFusedQkvOldMacro6(TestMacro):
+    def macro_name(self):
+        return "fused_qkv_old_macro"
+
+    def source_code(self):
+        return [
+            "fused_qkv_old_test_name ->  A_TEST_NAME,fused_qkv_old, num_heads = 8, num_key_value_groups = 8 ",
+            "A_TEST_NAME ->  layers.2.self_attn.qkv_proj.weight",
+        ]
+
+    def expected(self):
+        return [
+            'fused_qkv_old_test_name -> fused_qkv_old_tmp.Q_0,fused_qkv_old_tmp.Q_1,fused_qkv_old_tmp.Q_2,fused_qkv_old_tmp.Q_3,fused_qkv_old_tmp.K_0,fused_qkv_old_tmp.K_1,fused_qkv_old_tmp.K_2,fused_qkv_old_tmp.K_3,fused_qkv_old_tmp.V_0,fused_qkv_old_tmp.V_1,fused_qkv_old_tmp.V_2,fused_qkv_old_tmp.V_3,fused_qkv_old_tmp.Q_4,fused_qkv_old_tmp.Q_5,fused_qkv_old_tmp.Q_6,fused_qkv_old_tmp.Q_7,fused_qkv_old_tmp.K_4,fused_qkv_old_tmp.K_5,fused_qkv_old_tmp.K_6,fused_qkv_old_tmp.K_7,fused_qkv_old_tmp.V_4,fused_qkv_old_tmp.V_5,fused_qkv_old_tmp.V_6,fused_qkv_old_tmp.V_7, axis=1',
+            'fused_qkv_old_tmp.Q_0,fused_qkv_old_tmp.Q_1,fused_qkv_old_tmp.K_0,fused_qkv_old_tmp.K_1,fused_qkv_old_tmp.V_0,fused_qkv_old_tmp.V_1,fused_qkv_old_tmp.Q_2,fused_qkv_old_tmp.Q_3,fused_qkv_old_tmp.K_2,fused_qkv_old_tmp.K_3,fused_qkv_old_tmp.V_2,fused_qkv_old_tmp.V_3,fused_qkv_old_tmp.Q_4,fused_qkv_old_tmp.Q_5,fused_qkv_old_tmp.K_4,fused_qkv_old_tmp.K_5,fused_qkv_old_tmp.V_4,fused_qkv_old_tmp.V_5,fused_qkv_old_tmp.Q_6,fused_qkv_old_tmp.Q_7,fused_qkv_old_tmp.K_6,fused_qkv_old_tmp.K_7,fused_qkv_old_tmp.V_6,fused_qkv_old_tmp.V_7 -> A_TEST_NAME, axis=1',
+            'A_TEST_NAME ->  layers.2.self_attn.qkv_proj.weight',
+        ]
+
+    def test(self):
+        self.start_macro_test(aoa_config_reverse=True)
+
+
+class TestFusedFfnMacro(TestMacro):
+    def macro_name(self):
+        return "fused_ffn_macro"
+
+    def source_code(self):
+        return "layers.2.mlp.gate_up_fused_proj.weight -> layers.2.mlp.gate_up_fused_proj.weight, fused_ffn"
+
+    def expected(self):
+        return [
+            'layers.2.mlp.gate_up_fused_proj.weight  -> fused_ffn_tmp.GATE_0,fused_ffn_tmp.GATE_1,fused_ffn_tmp.UP_0,fused_ffn_tmp.UP_1,fused_ffn_tmp.GATE_2,fused_ffn_tmp.GATE_3,fused_ffn_tmp.UP_2,fused_ffn_tmp.UP_3, axis=1',
+            'fused_ffn_tmp.GATE_0,fused_ffn_tmp.UP_0,fused_ffn_tmp.GATE_1,fused_ffn_tmp.UP_1,fused_ffn_tmp.GATE_2,fused_ffn_tmp.UP_2,fused_ffn_tmp.GATE_3,fused_ffn_tmp.UP_3 -> layers.2.mlp.gate_up_fused_proj.weight, axis=1',
+        ]
+
+    def test(self):
+        self.start_macro_test()
+
+
+class TestFusedFfnMacro2(TestMacro):
+    def macro_name(self):
+        return "fused_ffn_macro"
+
+    def source_code(self):
+        return "layers.1.mlp.gate_up_fused_proj.weight -> layers.1.mlp.gate_proj.weight,layers.1.mlp.up_proj.weight, fused_ffn "
+
+    def expected(self):
+        return [
+            'layers.1.mlp.gate_up_fused_proj.weight  -> fused_ffn_tmp.GATE_0,fused_ffn_tmp.UP_0,fused_ffn_tmp.GATE_1,fused_ffn_tmp.UP_1, axis=1',
+            'fused_ffn_tmp.GATE_0,fused_ffn_tmp.GATE_1 -> layers.1.mlp.gate_proj.weight, axis=1',
+            'fused_ffn_tmp.UP_0,fused_ffn_tmp.UP_1 -> layers.1.mlp.up_proj.weight, axis=1',
+        ]
+
+    def test(self):
+        self.start_macro_test()
+
+
+class TestFusedFfnMacro3(TestMacro):
+    def macro_name(self):
+        return "fused_ffn_macro"
+
+    def source_code(self):
+        return "layers.1.mlp.gate_up_fused_proj.weight -> layers.1.mlp.gate_proj.weight,layers.1.mlp.up_proj.weight, fused_ffn "
+
+    def expected(self):
+        return [
+            'layers.1.mlp.gate_up_fused_proj.weight  -> fused_ffn_tmp.GATE_0,fused_ffn_tmp.UP_0,fused_ffn_tmp.GATE_1,fused_ffn_tmp.UP_1, axis=1',
+            'fused_ffn_tmp.GATE_0,fused_ffn_tmp.GATE_1 -> layers.1.mlp.gate_proj.weight, axis=1',
+            'fused_ffn_tmp.UP_0,fused_ffn_tmp.UP_1 -> layers.1.mlp.up_proj.weight, axis=1',
+        ]
+
+    def test(self):
+        self.start_macro_test()
+
+
+class TestFusedFfnMacro4(TestMacro):
+    def macro_name(self):
+        return "fused_ffn_macro"
+
+    def source_code(self):
+        return "layers.2.mlp.gate_up_fused_proj.bias -> layers.2.mlp.gate_up_fused_proj.bias, fused_ffn, axis=0"
+
+    def expected(self):
+        return [
+            'layers.2.mlp.gate_up_fused_proj.bias  -> fused_ffn_tmp.GATE_0,fused_ffn_tmp.GATE_1,fused_ffn_tmp.UP_0,fused_ffn_tmp.UP_1,fused_ffn_tmp.GATE_2,fused_ffn_tmp.GATE_3,fused_ffn_tmp.UP_2,fused_ffn_tmp.UP_3, axis=0',
+            'fused_ffn_tmp.GATE_0,fused_ffn_tmp.UP_0,fused_ffn_tmp.GATE_1,fused_ffn_tmp.UP_1,fused_ffn_tmp.GATE_2,fused_ffn_tmp.UP_2,fused_ffn_tmp.GATE_3,fused_ffn_tmp.UP_3 -> layers.2.mlp.gate_up_fused_proj.bias, axis=0',
+        ]
+
+    def test(self):
+        self.start_macro_test()
+
+
+class TestFusedFfnMacro5(TestMacro):
+    def macro_name(self):
+        return "fused_ffn_macro"
+
+    def source_code(self):
+        return [
+            "layers.1.mlp.gate_up_fused_proj.weight_test1 -> A_TEST_NAME, fused_ffn ",
+            "A_TEST_NAME -> layers.1.mlp.gate_up_fused_proj.weight_test2",
+        ]
+
+    def expected(self):
+        return [
+            'layers.1.mlp.gate_up_fused_proj.weight_test1  -> fused_ffn_tmp.GATE_0,fused_ffn_tmp.GATE_1,fused_ffn_tmp.UP_0,fused_ffn_tmp.UP_1,fused_ffn_tmp.GATE_2,fused_ffn_tmp.GATE_3,fused_ffn_tmp.UP_2,fused_ffn_tmp.UP_3, axis=1',
+            'fused_ffn_tmp.GATE_0,fused_ffn_tmp.UP_0,fused_ffn_tmp.GATE_1,fused_ffn_tmp.UP_1,fused_ffn_tmp.GATE_2,fused_ffn_tmp.UP_2,fused_ffn_tmp.GATE_3,fused_ffn_tmp.UP_3 -> A_TEST_NAME, axis=1',
+            'A_TEST_NAME -> layers.1.mlp.gate_up_fused_proj.weight_test2',
+        ]
+
+    def test(self):
+        self.start_macro_test(aoa_config_reverse=True)
+
+
+class TestLayerIdOffsetMacro(TestMacro):
+    def macro_name(self):
+        return "layer_id_offset_macro"
+
+    def source_code(self):
+        return "layers.$LAYER_ID_OFFSET.experts.0.weight -> layers.$LAYER_ID_OFFSET.experts.0.weight, axis = 1"
+
+    def expected(self):
+        return [
+            'layers.1.experts.0.weight->layers.0.experts.0.weight,axis=1\n',
+            'layers.2.experts.0.weight->layers.1.experts.0.weight,axis=1\n',
+        ]
+
+    def test(self):
+        self.start_macro_test()
+
+
+class TestIdMacroCase0(TestMacro):
+    def macro_name(self):
+        return "id_macro"
+
+    def source_code(self):
+        return "layers.$LAYER_ID.qkv_proj.weight->layers.$LAYER_ID.q_proj.weight,layer.$LAYER_ID.k_proj.weight,layer.$LAYER_ID.v_proj.weight, fused_qkv_old, num_heads = 8, num_key_value_groups = 4\n"
+
+    def expected(self):
+        return [
+            'layers.0.qkv_proj.weight->layers.0.q_proj.weight,layer.0.k_proj.weight,layer.0.v_proj.weight,fused_qkv_old,num_heads=8,num_key_value_groups=4\n',
+        ]
+
+    def test(self):
+        self.start_macro_test()
+
+
+class TestIdMacroCase1(TestMacro):
+    def macro_name(self):
+        return "id_macro"
+
+    def source_code(self):
+        return "layers.5.experts.$EXPERT_ID.up_gate_proj.weight -> layers.5.experts.$EXPERT_ID.gate_proj.weight, layers.5.experts.$EXPERT_ID.up_proj.weight, fused_ffn"
+
+    def expected(self):
+        return [
+            'layers.5.experts.0.up_gate_proj.weight->layers.5.experts.0.gate_proj.weight,layers.5.experts.0.up_proj.weight,fused_ffn\n',
+            'layers.5.experts.1.up_gate_proj.weight->layers.5.experts.1.gate_proj.weight,layers.5.experts.1.up_proj.weight,fused_ffn\n',
+        ]
+
+    def test(self):
+        self.start_macro_test()
+
+
+class TestIdMacroCase2(TestMacro):
+    def macro_name(self):
+        return "id_macro"
+
+    def source_code(self):
+        return "layers.$LAYER_ID.experts.$EXPERT_ID.up_gate_proj.weight -> layers.$LAYER_ID.experts.$EXPERT_ID.gate_proj.weight, fused_ffn"
+
+    def expected(self):
+        return [
+            'layers.3.experts.0.up_gate_proj.weight->layers.3.experts.0.gate_proj.weight,fused_ffn\n',
+            'layers.5.experts.0.up_gate_proj.weight->layers.5.experts.0.gate_proj.weight,fused_ffn\n',
+            'layers.3.experts.1.up_gate_proj.weight->layers.3.experts.1.gate_proj.weight,fused_ffn\n',
+            'layers.5.experts.1.up_gate_proj.weight->layers.5.experts.1.gate_proj.weight,fused_ffn\n',
+        ]
+
+    def test(self):
+        self.start_macro_test()
+
+
+class TestIdMacroCase3(TestMacro):
+    def macro_name(self):
+        return "id_macro"
+
+    def source_code(self):
+        return "layers.$LAYER_ID.experts.$EXPERT_ID.up_gate_proj.weight^T -> layers.$LAYER_ID.experts.$EXPERT_ID.gate_proj.weight, fused_ffn"
+
+    def expected(self):
+        return [
+            'layers.3.experts.0.up_gate_proj.weight^T->layers.3.experts.0.gate_proj.weight,fused_ffn\n',
+            'layers.5.experts.0.up_gate_proj.weight^T->layers.5.experts.0.gate_proj.weight,fused_ffn\n',
+            'layers.3.experts.1.up_gate_proj.weight^T->layers.3.experts.1.gate_proj.weight,fused_ffn\n',
+            'layers.5.experts.1.up_gate_proj.weight^T->layers.5.experts.1.gate_proj.weight,fused_ffn\n',
+        ]
+
+    def test(self):
+        self.start_macro_test()
+
+
+if __name__ == "__main__":
+    unittest.main()

test/flex_checkpoint/test_model_full_param.py

@@ -0,0 +1,328 @@
+# 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 unittest
+
+import collective.test_communication_api_base as test_base
+
+import paddle
+from paddle import nn
+
+TEST_CONFIGS = {
+    "2_card_tests": [
+        {
+            "world_size": 2,
+            "tp": 2,
+            "dp": 1,
+            "sharding_degree": 1,
+            "has_bias": "True",
+        },
+        {
+            "world_size": 2,
+            "tp": 2,
+            "dp": 1,
+            "sharding_degree": 1,
+            "has_bias": "True",
+        },
+        {
+            "world_size": 2,
+            "tp": 2,
+            "dp": 1,
+            "sharding_degree": 1,
+            "has_bias": "False",
+        },
+        {
+            "world_size": 2,
+            "tp": 2,
+            "dp": 1,
+            "sharding_degree": 1,
+            "has_bias": "False",
+        },
+        {
+            "world_size": 2,
+            "tp": 2,
+            "dp": 1,
+            "sharding_degree": 1,
+            "has_bias": "False",
+        },
+        {
+            "test_type": "layer",
+            "layer_type": "ColumnSequenceParallelLinear",
+            "world_size": 2,
+            "tp": 2,
+            "dp": 1,
+            "sharding_degree": 1,
+            "has_bias": "True",
+        },
+        {
+            "world_size": 2,
+            "tp": 2,
+            "dp": 1,
+            "sharding_degree": 1,
+            "has_bias": "True",
+        },
+        {
+            "world_size": 2,
+            "tp": 2,
+            "sharding_degree": 1,
+            "has_bias": "False",
+        },
+        {
+            "world_size": 2,
+            "tp": 1,
+            "sharding_degree": 2,
+            "has_bias": "False",
+        },
+        {
+            "world_size": 2,
+            "tp": 1,
+            "sharding_degree": 2,
+            "has_bias": "False",
+        },
+        {
+            "world_size": 2,
+            "tp": 2,
+            "sharding_degree": 1,
+            "has_bias": "True",
+            "master_weight": "True",
+        },
+        {
+            "world_size": 2,
+            "tp": 1,
+            "sharding_degree": 2,
+            "has_bias": "True",
+            "master_weight": "True",
+        },
+        {
+            "world_size": 2,
+            "tp": 1,
+            "sharding_degree": 2,
+            "has_bias": "True",
+            "master_weight": "True",
+        },
+    ],
+    "4_card_tests": [
+        {
+            "world_size": 4,
+            "tp": 4,
+            "dp": 1,
+            "sharding_degree": 1,
+            "has_bias": "True",
+        },
+        {
+            "world_size": 4,
+            "tp": 4,
+            "dp": 1,
+            "sharding_degree": 1,
+            "has_bias": "True",
+        },
+        {
+            "world_size": 4,
+            "tp": 2,
+            "dp": 2,
+            "sharding_degree": 1,
+            "has_bias": "True",
+        },
+        {
+            "world_size": 4,
+            "tp": 2,
+            "dp": 2,
+            "sharding_degree": 1,
+            "has_bias": "True",
+        },
+    ],
+    "4_card_hv_group_tests": [
+        {
+            "world_size": 4,
+            "tp": 2,
+            "pp": 2,
+            "sharding_degree": 1,
+            "has_bias": "True",
+            "test_using_hv_group": 1,
+        },
+    ],
+    "2_card_hv_group_tests": [
+        {
+            "world_size": 2,
+            "tp": 2,
+            "pp": 1,
+            "sharding_degree": 1,
+            "has_bias": "True",
+            "test_using_hv_group": 1,
+        },
+    ],
+    "sharding3_with_convert2cpu_tests": [
+        {
+            "world_size": 2,
+            "tp": 1,
+            "pp": 1,
+            "sharding_degree": 2,
+            "has_bias": "True",
+        },
+    ],
+}
+
+
+class TestFullParamWith2Devices(test_base.CommunicationTestDistBase):
+    def setUp(self):
+        super().setUp(num_of_devices=2, timeout=240)
+
+    def test_full_param(self):
+        for config in TEST_CONFIGS["2_card_tests"]:
+            envs = {k: str(v) for k, v in config.items()}
+            envs["test_using_hv_group"] = "0"
+            self.run_test_case(
+                "model_full_param_logic.py",
+                user_defined_envs=envs,
+            )
+
+
+class TestFullParamWith4Devices(test_base.CommunicationTestDistBase):
+    def setUp(self):
+        super().setUp(num_of_devices=4, timeout=240)
+
+    def test_full_param(self):
+        for config in TEST_CONFIGS["4_card_tests"]:
+            envs = {k: str(v) for k, v in config.items()}
+            envs["test_using_hv_group"] = "0"
+            self.run_test_case(
+                "model_full_param_logic.py",
+                user_defined_envs=envs,
+            )
+
+
+class TestFullParamWithSingleDevices(unittest.TestCase):
+    class SimpleMLP(nn.Layer):
+        def __init__(self, hidden_size=100, has_bias=False):
+            super().__init__()
+            self.embedding = nn.Embedding(24, hidden_size)
+            self.linear1 = nn.Linear(
+                hidden_size, hidden_size, bias_attr=has_bias
+            )
+            self.linear2 = nn.Linear(
+                hidden_size, hidden_size, bias_attr=has_bias
+            )
+            self.llm_head = nn.Linear(hidden_size, 24, bias_attr=False)
+
+        def forward(self, x):
+            x = self.embedding(x)
+            x = self.linear1(x)
+            x = self.linear2(x)
+            x = self.llm_head(x)
+            return x
+
+    def test_full_param(self):
+        self.batch_size = 2
+        self.hidden_size = 32
+        self.has_bias = True
+        model = self.SimpleMLP(
+            hidden_size=self.hidden_size, has_bias=self.has_bias
+        )
+        model = paddle.amp.decorate(
+            models=model, optimizers=None, level="O2", dtype="float16"
+        )
+        model.train()
+        model_state_dict = model.state_dict()
+
+        for k, v in model_state_dict.items():
+            ones = paddle.ones_like(v)
+            paddle.assign(ones, v)
+            if k == "linear1.weight":
+                zeros = paddle.zeros_like(v)
+                paddle.assign(zeros, v)
+
+        aoa_config = {
+            "aoa_statements": [
+                "linear1.weight, linear2.weight -> fused_weight, axis=1"
+                "embedding.weight -> embedding.weight, dtype = 'float32'"
+            ]
+        }
+
+        full_param_iter = model.full(aoa_config)
+        full_param = dict(full_param_iter)
+
+        param_shape = {
+            # "linear1.weight" : [32,32],
+            # "linear2.weight" : [32, 32],
+            "embedding.weight": [24, 32],
+            "linear1.bias": [32],
+            "linear2.bias": [32],
+            "llm_head.weight": [24, 32],
+            "fused_weight": [32, 64],
+        }
+
+        for name, shape in param_shape.items():
+            if name == "fused_weight":
+                continue
+            if not self.has_bias:
+                if ".bias" in name:
+                    continue
+            assert name in full_param.keys()
+            tensor = full_param[name]
+            answer = paddle.ones_like(tensor)
+            assert tensor._md5sum() == answer._md5sum()
+            if name == "embedding.weight":
+                assert tensor.dtype == paddle.float32
+        assert "fused_weight" in full_param.keys()
+        ones = paddle.ones([32, 32], 'float16')
+        zeros = paddle.zeros([32, 32], 'float16')
+        answer = paddle.concat([zeros, ones], axis=1)
+        assert full_param["fused_weight"]._md5sum() == answer._md5sum()
+
+
+class TestFullParamHVGroupWith2Devices(test_base.CommunicationTestDistBase):
+    def setUp(self):
+        super().setUp(num_of_devices=2, timeout=240)
+
+    def test_full_param(self):
+        for config in TEST_CONFIGS["2_card_hv_group_tests"]:
+            envs = {k: str(v) for k, v in config.items()}
+            envs["test_using_hv_group"] = "1"
+            self.run_test_case(
+                "model_full_param_logic.py",
+                user_defined_envs=envs,
+            )
+
+
+class TestFullParamHVGroupWith4Devices(test_base.CommunicationTestDistBase):
+    def setUp(self):
+        super().setUp(num_of_devices=4, timeout=240)
+
+    def test_full_param(self):
+        for config in TEST_CONFIGS["4_card_hv_group_tests"]:
+            envs = {k: str(v) for k, v in config.items()}
+            envs["test_using_hv_group"] = "1"
+            self.run_test_case(
+                "model_full_param_logic.py",
+                user_defined_envs=envs,
+            )
+
+
+class TestFullParamWithSharding3(test_base.CommunicationTestDistBase):
+    def setUp(self):
+        super().setUp(num_of_devices=2, timeout=240)
+
+    def test_full_param(self):
+        for config in TEST_CONFIGS["sharding3_with_convert2cpu_tests"]:
+            envs = {k: str(v) for k, v in config.items()}
+            envs["test_using_hv_group"] = "0"
+            envs["test_with_sharding3"] = "1"
+            self.run_test_case(
+                "model_full_param_logic.py",
+                user_defined_envs=envs,
+            )
+
+
+if __name__ == "__main__":
+    unittest.main()

test/flex_checkpoint/test_sharded_state_dict.py

@@ -0,0 +1,235 @@
+# 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 unittest
+
+import collective.test_communication_api_base as test_base
+
+TEST_CONFIGS = {
+    "2_card_tests": [
+        {
+            "test_type": "layer",
+            "layer_type": "ColumnParallelLinear",
+            "world_size": 2,
+            "tp": 2,
+            "dp": 1,
+            "sharding_degree": 1,
+            "has_bias": "True",
+        },
+        {
+            "test_type": "layer",
+            "layer_type": "RowParallelLinear",
+            "world_size": 2,
+            "tp": 2,
+            "dp": 1,
+            "sharding_degree": 1,
+            "has_bias": "True",
+        },
+        {
+            "test_type": "layer",
+            "layer_type": "VocabParallelEmbedding",
+            "world_size": 2,
+            "tp": 2,
+            "dp": 1,
+            "sharding_degree": 1,
+            "has_bias": "False",
+        },
+        {
+            "test_type": "layer",
+            "layer_type": "ColumnParallelLinear",
+            "world_size": 2,
+            "tp": 2,
+            "dp": 1,
+            "sharding_degree": 1,
+            "has_bias": "False",
+        },
+        {
+            "test_type": "layer",
+            "layer_type": "RowParallelLinear",
+            "world_size": 2,
+            "tp": 2,
+            "dp": 1,
+            "sharding_degree": 1,
+            "has_bias": "False",
+        },
+        {
+            "test_type": "layer",
+            "layer_type": "ColumnSequenceParallelLinear",
+            "world_size": 2,
+            "tp": 2,
+            "dp": 1,
+            "sharding_degree": 1,
+            "has_bias": "True",
+        },
+        {
+            "test_type": "layer",
+            "layer_type": "RowSequenceParallelLinear",
+            "world_size": 2,
+            "tp": 2,
+            "dp": 1,
+            "sharding_degree": 1,
+            "has_bias": "True",
+        },
+        {
+            "test_type": "optimizer",
+            "layer_type": "AdamW",
+            "world_size": 2,
+            "tp": 2,
+            "sharding_degree": 1,
+            "has_bias": "False",
+        },
+        {
+            "test_type": "optimizer",
+            "layer_type": "DygraphShardingOptimizer",
+            "world_size": 2,
+            "tp": 1,
+            "sharding_degree": 2,
+            "has_bias": "False",
+        },
+        {
+            "test_type": "optimizer",
+            "layer_type": "DygraphShardingOptimizerV2",
+            "world_size": 2,
+            "tp": 1,
+            "sharding_degree": 2,
+            "has_bias": "False",
+        },
+        {
+            "test_type": "optimizer",
+            "layer_type": "AdamW",
+            "world_size": 2,
+            "tp": 2,
+            "sharding_degree": 1,
+            "has_bias": "True",
+            "master_weight": "True",
+        },
+        {
+            "test_type": "optimizer",
+            "layer_type": "DygraphShardingOptimizer",
+            "world_size": 2,
+            "tp": 1,
+            "sharding_degree": 2,
+            "has_bias": "True",
+            "master_weight": "True",
+        },
+        {
+            "test_type": "optimizer",
+            "layer_type": "DygraphShardingOptimizerV2",
+            "world_size": 2,
+            "tp": 1,
+            "sharding_degree": 2,
+            "has_bias": "True",
+            "master_weight": "True",
+        },
+        {
+            "test_type": "optimizer",
+            "layer_type": "GroupShardedOptimizerStage2",
+            "world_size": 2,
+            "tp": 1,
+            "sharding_degree": 2,
+            "has_bias": "True",
+            "master_weight": "True",
+        },
+        {
+            "test_type": "optimizer",
+            "layer_type": "GroupShardedStage3",
+            "world_size": 2,
+            "tp": 1,
+            "sharding_degree": 2,
+            "has_bias": "True",
+            "master_weight": "True",
+        },
+    ],
+    "4_card_tests": [
+        {
+            "test_type": "layer",
+            "layer_type": "ColumnParallelLinear",
+            "world_size": 4,
+            "tp": 4,
+            "dp": 1,
+            "sharding_degree": 1,
+            "has_bias": "True",
+        },
+        {
+            "test_type": "layer",
+            "layer_type": "RowParallelLinear",
+            "world_size": 4,
+            "tp": 4,
+            "dp": 1,
+            "sharding_degree": 1,
+            "has_bias": "True",
+        },
+        {
+            "test_type": "layer",
+            "layer_type": "ColumnParallelLinear",
+            "world_size": 4,
+            "tp": 2,
+            "dp": 2,
+            "sharding_degree": 1,
+            "has_bias": "True",
+        },
+        {
+            "test_type": "layer",
+            "layer_type": "RowParallelLinear",
+            "world_size": 4,
+            "tp": 2,
+            "dp": 2,
+            "sharding_degree": 1,
+            "has_bias": "True",
+        },
+    ],
+}
+
+
+class TestParallelLayersWith2Devices(test_base.CommunicationTestDistBase):
+    def setUp(self):
+        super().setUp(num_of_devices=2, timeout=240)
+
+    def test_metadata(self):
+        for config in TEST_CONFIGS["2_card_tests"]:
+            envs = {k: str(v) for k, v in config.items()}
+            self.run_test_case(
+                "sharded_state_dict_logic.py",
+                user_defined_envs=envs,
+            )
+
+
+class TestParallelLayersWith4Devices(test_base.CommunicationTestDistBase):
+    def setUp(self):
+        super().setUp(num_of_devices=4, timeout=240)
+
+    def test_metadata(self):
+        for config in TEST_CONFIGS["4_card_tests"]:
+            envs = {k: str(v) for k, v in config.items()}
+            self.run_test_case(
+                "sharded_state_dict_logic.py",
+                user_defined_envs=envs,
+            )
+
+
+class TestMergeShardedAOA(test_base.CommunicationTestDistBase):
+    def setUp(self):
+        super().setUp(num_of_devices=2, timeout=120)
+
+    def test_merge_sharded(self):
+        config = TEST_CONFIGS["2_card_tests"][0]
+        envs = {k: str(v) for k, v in config.items()}
+        self.run_test_case(
+            "merge_sharded_state_dict.py",
+            user_defined_envs=envs,
+        )
+
+
+if __name__ == "__main__":
+    unittest.main()

test/flex_checkpoint/test_strategy_conversion.py

@@ -0,0 +1,289 @@
+# 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 argparse
+import logging
+import os
+import subprocess
+import sys
+import tempfile
+import unittest
+
+
+def p_str_to_dict(p_str):
+    """Parses a strategy string like 'd2·t2' into a config dictionary."""
+    config = {"tp": 1, "dp": 1, "pp": 1, "ep": 1}
+    parts = p_str.split('·')
+    for part in parts:
+        if part.startswith('d'):
+            config['dp'] = int(part[1:])
+        elif part.startswith('t'):
+            config['tp'] = int(part[1:])
+        elif part.startswith('p'):
+            config['pp'] = int(part[1:])
+        elif part.startswith('e'):
+            config['ep'] = int(part[1:])
+
+    if config['ep'] > 1 and config['dp'] < config['ep']:
+        config['dp'] = config['ep']
+
+    config["num_cards"] = config["tp"] * config["dp"] * config["pp"]
+    if p_str in ["d1", "t1", "p1", "e1"]:
+        config["num_cards"] = 1
+
+    return config
+
+
+TEST_CASES = [
+    {
+        "id": "B1_d2_to_d4",
+        "src": p_str_to_dict("d2"),
+        "tgt": p_str_to_dict("d4"),
+        "gpu_num": 4,
+    },
+    {
+        "id": "B2_t2_to_t4",
+        "src": p_str_to_dict("t2"),
+        "tgt": p_str_to_dict("t4"),
+        "gpu_num": 4,
+    },
+    {
+        "id": "B3_p2_to_p4",
+        "src": p_str_to_dict("p2"),
+        "tgt": p_str_to_dict("p4"),
+        "gpu_num": 4,
+    },
+    {
+        "id": "B4_e2_to_e4",
+        "src": p_str_to_dict("e2"),
+        "tgt": p_str_to_dict("e4"),
+        "model_type": "moe",
+        "gpu_num": 4,
+    },
+    # Case 5 (pp2 -> tp4)
+    {
+        "id": "X5_pp2_to_tp4",
+        "src": p_str_to_dict("p2"),
+        "tgt": p_str_to_dict("t4"),
+        "gpu_num": 4,
+    },
+    # Case 6 (tp2 -> pp2)
+    {
+        "id": "X6_tp2_to_pp2",
+        "src": p_str_to_dict("t2"),
+        "tgt": p_str_to_dict("p2"),
+        "gpu_num": 2,
+    },
+    # Case 7 (dp4 -> tp2·dp2)
+    {
+        "id": "X7_dp4_to_tp2dp2",
+        "src": p_str_to_dict("d4"),
+        "tgt": p_str_to_dict("t2·d2"),
+        "gpu_num": 4,
+    },
+    # Case 8 (dp2 -> pp2)
+    {
+        "id": "X8_dp2_to_pp2",
+        "src": p_str_to_dict("d2"),
+        "tgt": p_str_to_dict("p2"),
+        "gpu_num": 2,
+    },
+    # Case 9 (dp2 -> ep2)
+    {
+        "id": "X9_dp2_to_ep2",
+        "src": p_str_to_dict("d2"),
+        "tgt": p_str_to_dict("e2"),
+        "model_type": "moe",
+        "gpu_num": 2,
+    },
+    # Case 10 (ep2 -> tp2)
+    {
+        "id": "X10_ep2_to_tp2",
+        "src": p_str_to_dict("e2"),
+        "tgt": p_str_to_dict("t2"),
+        "model_type": "moe",
+        "gpu_num": 2,
+    },
+    # Case 11 (tp2 -> ep2)
+    {
+        "id": "X11_tp2_to_ep2",
+        "src": p_str_to_dict("t2"),
+        "tgt": p_str_to_dict("e2"),
+        "model_type": "moe",
+        "gpu_num": 2,
+    },
+    {
+        "id": "M12_dp2tp2_to_tp4",
+        "src": p_str_to_dict("d2·t2"),
+        "tgt": p_str_to_dict("t4"),
+        "gpu_num": 4,
+    },
+    {
+        "id": "M13_dp2tp2_to_pp4",
+        "src": p_str_to_dict("d2·t2"),
+        "tgt": p_str_to_dict("p4"),
+        "gpu_num": 4,
+    },
+    {
+        "id": "M14_dp2pp2_to_tp4",
+        "src": p_str_to_dict("d2·p2"),
+        "tgt": p_str_to_dict("t4"),
+        "gpu_num": 4,
+    },
+    {
+        "id": "M15_tp2pp2_to_dp4",
+        "src": p_str_to_dict("t2·p2"),
+        "tgt": p_str_to_dict("d4"),
+        "gpu_num": 4,
+    },
+    {
+        "id": "M16_tp2pp2_to_dp2tp2",
+        "src": p_str_to_dict("t2·p2"),
+        "tgt": p_str_to_dict("d2·t2"),
+        "gpu_num": 4,
+    },
+    {
+        "id": "M17_dp2ep2_to_dp4",
+        "src": p_str_to_dict("d2·e2"),
+        "tgt": p_str_to_dict("d4"),
+        "model_type": "moe",
+        "gpu_num": 4,
+    },
+    {
+        "id": "M18_tp2ep2_to_tp4",
+        "src": p_str_to_dict("t2·e2"),
+        "tgt": p_str_to_dict("t4"),
+        "model_type": "moe",
+        "gpu_num": 4,
+    },
+    # Case 19 (dp2·tp2 -> pp2)
+    {
+        "id": "M19_dp2tp2_to_pp2",
+        "src": p_str_to_dict("d2·t2"),
+        "tgt": p_str_to_dict("p2"),
+        "gpu_num": 4,
+    },
+    # E1 (e2->e4) is covered by B4
+    {
+        "id": "E2_dp2ep2_to_tp2ep2",
+        "src": p_str_to_dict("d2·e2"),
+        "tgt": p_str_to_dict("t2·e2"),
+        "model_type": "moe",
+        "gpu_num": 4,
+    },
+]
+
+
+class TestStrategyConversion(unittest.TestCase):
+    def _run_workflow(self, case, logic_script="strategy_conversion_engine.py"):
+        import paddle
+
+        if case["gpu_num"] > paddle.device.cuda.device_count():
+            self.skipTest("number of GPUs is not enough")
+
+        case_id = case['id']
+        src_config = case['src']
+        tgt_config = case['tgt']
+
+        src_gpus_count = src_config.pop("num_cards")
+        tgt_gpus_count = tgt_config.pop("num_cards")
+        src_gpus = ",".join(map(str, range(src_gpus_count)))
+        tgt_gpus = ",".join(map(str, range(tgt_gpus_count)))
+
+        with tempfile.TemporaryDirectory() as tmpdir:
+            src_ckpt_path = os.path.join(tmpdir, "src_ckpt")
+            tgt_ckpt_path = os.path.join(tmpdir, "tgt_ckpt")
+
+            def config_to_args(config, prefix):
+                return [
+                    f"--{prefix}_{k}={v}"
+                    for k, v in config.items()
+                    if not k.startswith('s_')
+                ]
+
+            common_args = config_to_args(src_config, "src") + config_to_args(
+                tgt_config, "tgt"
+            )
+            if "model_type" in case:
+                common_args.append(f"--model_type={case['model_type']}")
+            path_args = [
+                f"--src_ckpt_path={src_ckpt_path}",
+                f"--tgt_ckpt_path={tgt_ckpt_path}",
+            ]
+            base_cmd = [
+                sys.executable,
+                "-m",
+                "paddle.distributed.launch",
+                "--log_dir",
+                os.path.join(tmpdir, "logs"),
+            ]
+
+            steps = ["save_source", "convert", "verify"]
+            gpus_per_step = [src_gpus, tgt_gpus, src_gpus]
+
+            for i, step_name in enumerate(steps):
+                cmd = [
+                    *base_cmd,
+                    f"--gpus={gpus_per_step[i]}",
+                    logic_script,
+                    f"--step={step_name}",
+                    *common_args,
+                    *path_args,
+                ]
+                process = subprocess.run(
+                    cmd, capture_output=True, text=True, check=False
+                )
+
+                self.assertEqual(
+                    process.returncode,
+                    0,
+                    f"Step '{step_name}' FAILED for case '{case_id}'!\n"
+                    f"STDOUT:\n{process.stdout}\nSTDERR:\n{process.stderr}",
+                )
+
+
+def _create_test_method(case):
+    def test_method(self):
+        self._run_workflow(case)
+
+    return test_method
+
+
+for case_info in TEST_CASES:
+    test_name = f"test_{case_info['id']}"
+    test_func = _create_test_method(case_info)
+    setattr(TestStrategyConversion, test_name, test_func)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        '--list_tests',
+        action='store_true',
+        help='List all test case names that unittest can discover and exit.',
+    )
+    args, unknown = parser.parse_known_args()
+
+    if args.list_tests:
+        for case in TEST_CASES:
+            module_name = os.path.splitext(os.path.basename(__file__))[0]
+            logging.basicConfig(
+                stream=sys.stdout, level=logging.INFO, format="%(message)s"
+            )
+            logging.info(
+                f"{module_name}.TestStrategyConversion.test_{case['id']}"
+            )
+        sys.exit(0)
+
+    unittest.main(argv=[sys.argv[0]], *unknown)