fix aoa transpose corner case (PaddlePaddle#76234)

Author: xingmingyyj

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

@@ -744,6 +744,7 @@ def _handle_aoa(
                 and src_desc.global_shape == dst_desc.global_shape
                 and src_desc.global_offset == dst_desc.global_offset
                 and src_desc.dtype == dst_desc.dtype
+                and mapping.postprocess_list is None
             ):
                 new_load_dict[idx] = ShardedWeight(
                     key=src_desc.key,

test/flex_checkpoint/load_state_dict_transpose_logic.py

@@ -43,10 +43,10 @@ def sharded_state_dict(
 
 
 class SimpleMLP(Layer):
-    def __init__(self, hidden_size=1024):
+    def __init__(self, in_features=1024, out_features=1024):
         super().__init__()
         self.linear = ColumnParallelLinear(
-            hidden_size, hidden_size * 2, has_bias=True
+            in_features, out_features, has_bias=True
         )
 
     def forward(self, x):
@@ -55,10 +55,10 @@ def forward(self, x):
 
 
 class SimpleMLPTransWeight(Layer):
-    def __init__(self, hidden_size=1024):
+    def __init__(self, in_features=1024, out_features=1024):
         super().__init__()
         self.linear = ColumnParallelLinearTransWeight(
-            hidden_size, hidden_size * 2, has_bias=True
+            in_features, out_features, has_bias=True
         )
 
     def forward(self, x):
@@ -70,6 +70,8 @@ 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()
@@ -99,5 +101,13 @@ def run_save_state_dict(self):
         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()