【Flexcheckpoint】add_get_var_mapping_chain_macro (PaddlePaddle#76013)

* add_get_var_mapping_chain_macro

* add note

* fix the bug input_vars and resolve_mapping_chain

* fix the code style

* fit the dtype assert bug

* fix the bug

* fix the merge_sharded_state_dict bug

Author: zty-king

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

@@ -90,6 +90,8 @@ def __init__(
     ) -> None:
         self.source_state_shard_info = source_state_shard_info
         self.destination_state_shard_info = destination_state_shard_info
+        self.left_var_to_right_var_mapping = {}
+        self.right_var_from_left_var_mapping = {}
 
     def get_all_dst_state_keys(self):
         dst_state_keys = set()
@@ -135,11 +137,16 @@ def get_src_state_shard_num(self, src_state_key: str) -> int:
             "AOA notions apply only to the model state, but are automatically propagated to the optimizer state."
         )
 
+        # Only need to parse the model state key for optimizer state shard num, because the optimizer state slice info is completely consistent with the model state slice info.
+        resolved_model_state_key = self.resolve_mapping_chain(
+            model_state_key, reverse=True
+        )
+
         state_keys = [
-            model_state_key,
-            f"{model_state_key}.w_0",
-            f"{model_state_key}.moment1_0",
-            f"{model_state_key}.moment2_0",
+            resolved_model_state_key,
+            f"{resolved_model_state_key}.w_0",
+            f"{resolved_model_state_key}.moment1_0",
+            f"{resolved_model_state_key}.moment2_0",
         ]
 
         shard_nums = {
@@ -168,7 +175,6 @@ def get_dst_state_shard_num(self, dst_state_key: str) -> int:
         if self.destination_state_shard_info is None:
             # Default `dst_state_shard_num=1` if `destination_state_shard_info` is missing.
             return 1
-
         model_state_key, opt_state_name = split_optimizer_state_key(
             dst_state_key
         )
@@ -177,11 +183,16 @@ def get_dst_state_shard_num(self, dst_state_key: str) -> int:
             "AOA notions apply only to the model state, but are automatically propagated to the optimizer state."
         )
 
+        # Only need to parse the model state key for optimizer state shard num, because the optimizer state slice info is completely consistent with the model state slice info.
+        resolved_model_state_key = self.resolve_mapping_chain(
+            model_state_key, reverse=False
+        )
+
         state_keys = [
-            model_state_key,
-            f"{model_state_key}.w_0",
-            f"{model_state_key}.moment1_0",
-            f"{model_state_key}.moment2_0",
+            resolved_model_state_key,
+            f"{resolved_model_state_key}.w_0",
+            f"{resolved_model_state_key}.moment1_0",
+            f"{resolved_model_state_key}.moment2_0",
         ]
 
         shard_nums = {
@@ -206,6 +217,44 @@ def get_dst_state_shard_num(self, dst_state_key: str) -> int:
             )
         return shard_nums.pop()
 
+    def resolve_mapping_chain(self, key: str, reverse: bool = False) -> str:
+        """
+        Recursively resolve the mapping chain, find the final leaf node
+
+        Args:
+            key: The key to be resolved
+            reverse: False use left_var_to_right_var_mapping，True use right_var_from_left_var_mapping
+
+        For example:
+        - reverse=False: temp_var -> dst_key
+        - reverse=True: temp_var -> src_key
+        """
+        visited = set()  # avoid infinite loop
+        current_key = key
+
+        if reverse:
+            mapping_dict = self.right_var_from_left_var_mapping
+        else:
+            mapping_dict = self.left_var_to_right_var_mapping
+
+        while current_key in mapping_dict:
+            assert current_key not in visited, (
+                "Infinite loop detected in resolve_mapping_chain,which means the start key is not src_key or the end key is not dst_key, the aoa_config is error"
+            )
+            visited.add(current_key)
+            if reverse and current_key in self.get_all_src_state_keys():
+                break
+            elif not reverse and current_key in self.get_all_dst_state_keys():
+                break
+
+            mapped_vars = mapping_dict[current_key]
+            if mapped_vars and len(mapped_vars) > 0:
+                current_key = mapped_vars[0]
+            else:
+                break
+
+        return current_key
+
 
 class AOAEngine:
     def __init__(
@@ -248,14 +297,20 @@ def make_input_tensor(
 
     def build_input_vars(self):
         input_vars = {}
-        for key, shards in self.source_state_shard_info.items():
+        dtype = None
+        for key, shards in sorted(self.source_state_shard_info.items()):
             global_shape = shards[0].global_shape
-            dtype = shards[0].dtype
             model_state_key, opt_state_name = split_optimizer_state_key(key)
-            if opt_state_name in [".w_0", ".moment1_0", ".moment2_0", None]:
-                input_vars[model_state_key] = self.make_input_tensor(
-                    model_state_key, global_shape, dtype
-                )
+            if opt_state_name is None:
+                dtype = shards[0].dtype
+            if model_state_key in input_vars.keys() or opt_state_name in [
+                ".beta1_pow_acc_0",
+                ".beta2_pow_acc_0",
+            ]:
+                continue
+            input_vars[model_state_key] = self.make_input_tensor(
+                model_state_key, global_shape, dtype
+            )
         return input_vars
 
     def split(
@@ -652,11 +707,19 @@ def find_shard_sources(
 
         for src_key, src_slices, local_slices, pp_list in results:
             src_var = self.input_vars[src_key]
-            if src_var.dtype != target.dtype:
-                assert pp_list is not None and target.dtype in str(pp_list), (
-                    "Direct assignment of Tensors with different types is prohibited in AOA. "
-                    "If you want to achieve this functionality, please use the cast semantics provided by AOA."
-                )
+            target_model_state_key, target_opt_state_name = (
+                split_optimizer_state_key(target.key)
+            )
+            if target_opt_state_name is None:
+                if src_var.dtype != target.dtype:
+                    assert pp_list is not None and target.dtype in str(
+                        pp_list
+                    ), (
+                        "Direct assignment of Tensors with different types is prohibited in AOA. "
+                        "If you want to achieve this functionality, please use the cast semantics provided by AOA."
+                    )
+            else:
+                src_var.dtype = target.dtype
 
             src_global_shape = src_var.shape
 

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

@@ -88,12 +88,6 @@ def tokenize(self, text):
             mo = self.get_token(text, pos)
         return tokens
 
-    def apply_macros(self, expression):
-        expressions = [expression]
-        for macro in self.macros:
-            expressions = self.apply_macro(expressions, macro)
-        return expressions
-
     def apply_macro(self, expression, macro):
         if isinstance(expression, str):
             expression = [expression]
@@ -106,10 +100,24 @@ def apply_macro(self, expression, macro):
                 new_expression.extend(results)
         return new_expression
 
+    def apply_single_macro_to_all(self, expressions, macro):
+        new_expressions = []
+        for expr in expressions:
+            results = macro(self.tokenize(expr), expr, self.context)
+            if isinstance(results, str):
+                new_expressions.append(results)
+            else:
+                new_expressions.extend(results)
+        return new_expressions
+
     def all_tokens(self, expressions):
+        current_expressions = expressions
+        for macro in self.macros:
+            current_expressions = self.apply_single_macro_to_all(
+                current_expressions, macro
+            )
+
         tokens = []
-        for expr in expressions:
-            expanded_expressions = self.apply_macros(expr)
-            for e in expanded_expressions:
-                tokens.extend(self.tokenize(e))
+        for expr in current_expressions:
+            tokens.extend(self.tokenize(expr))
         return tokens

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

@@ -56,6 +56,8 @@ def register_macro(self, name, func, priority):
     'num_heads',
     'num_key_value_groups',
     'permute',
+    'dtype',
+    'fused_qkv',
 ]
 
 EXTRA_SUFFIX = [
@@ -211,7 +213,7 @@ def array_macro(tokens, expression, context):
     return new_expression
 
 
-@macro(name='fused_qkv_old_macro', priority=4)
+@macro(name='fused_qkv_old_macro', priority=6)
 def fused_qkv_old_macro(tokens, expression, context):
     FUSED_QKV_OLD_TAG = "fused_qkv_old"
     if not any(tkn.value == FUSED_QKV_OLD_TAG for tkn in tokens):
@@ -381,7 +383,7 @@ def gen_expr(tp_degree, num_heads, tp_rank, comp):
     return results
 
 
-@macro(name='fused_ffn_macro', priority=4)
+@macro(name='fused_ffn_macro', priority=6)
 def fused_ffn_macro(tokens, expression, context):
     FUSED_FFN_TAG = "fused_ffn"
     if not any(tkn.value == FUSED_FFN_TAG for tkn in tokens):
@@ -505,7 +507,7 @@ def gen_expr(tp_degree, splited_num, tp_rank, comp):
     return results
 
 
-@macro(name='transpose_macro', priority=3)
+@macro(name='transpose_macro', priority=5)
 def transpose_macro(tokens, expression, context):
     TRANSPOSE_TAG = "^T"
 
@@ -551,7 +553,7 @@ def transpose_macro(tokens, expression, context):
     return results
 
 
-@macro(name='fused_qkv_macro', priority=4)
+@macro(name='fused_qkv_macro', priority=6)
 def fused_qkv_macro(tokens, expression, context):
     FUSED_QKV_TAG = "fused_qkv"
     if not any(tkn.value == FUSED_QKV_TAG for tkn in tokens):
@@ -711,6 +713,7 @@ def find_matches(self, pattern: str) -> dict[str, list[int]]:
 _REGISTERED_PLACEHOLDERS = ['$EXPERT_ID', '$LAYER_ID']
 
 
+# TODO: need to adapt the scene of temp_layers.\$LAYER_ID.weight -> dst_layers.\$LAYER_ID.weight
 @macro(name='id_macro', priority=1)
 def id(tokens, expression, context):
     allowed_placeholders = _REGISTERED_PLACEHOLDERS
@@ -783,3 +786,43 @@ def dict_cartesian_tuples(d: dict[str, list[int]]):
         results.append(cur_statement)
 
     return results
+
+
+# This macro processes variable mappings between source and destination states,
+# but it requires that all expansion macros (layer_id_macro, expert_id_macro,
+# star_macro, array_macro, etc.) have already been executed to expand template
+# variables into concrete variable names.
+@macro(name='get_var_mapping_chain_macro', priority=4)
+def get_var_mapping_chain_macro(tokens, expression, context):
+    flag_left_var = True
+    left_var_list = []
+    right_var_list = []
+    for tkn in tokens:
+        if tkn.value in GLOBAL_ATTRIBUTE_KEYWORDS:
+            break
+        if tkn.type == TokenType.RARROW:
+            flag_left_var = False
+        if tkn.type == TokenType.IDENTIFIER:
+            extra_suffix_removed_value = tkn.value
+            for sfx in EXTRA_SUFFIX:
+                extra_suffix_removed_value = (
+                    extra_suffix_removed_value.removesuffix(sfx)
+                )
+            if flag_left_var:
+                left_var_list.append(extra_suffix_removed_value)
+            else:
+                right_var_list.append(extra_suffix_removed_value)
+    assert len(left_var_list) == 1 or len(right_var_list) == 1, (
+        "Left or right variable must have the only one element"
+    )
+    if len(left_var_list) == 1:
+        context.left_var_to_right_var_mapping[left_var_list[0]] = right_var_list
+        for right_var in right_var_list:
+            context.right_var_from_left_var_mapping[right_var] = left_var_list
+    else:
+        context.right_var_from_left_var_mapping[right_var_list[0]] = (
+            left_var_list
+        )
+        for left_var in left_var_list:
+            context.left_var_to_right_var_mapping[left_var] = right_var_list
+    return expression

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

@@ -75,7 +75,7 @@ def consume(self, expected_type=None):
         tok = self.peek()
         if expected_type and tok.type != expected_type:
             raise SyntaxError(
-                f'Expected {expected_type}, got {tok.type} at pos {tok.pos}'
+                f'Expected {expected_type}, got {tok.type} at pos {self.pos}'
             )
         self.pos += 1
         return tok

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

@@ -888,27 +888,17 @@ def load_state_dict(
                 f"{key} is not replicated!"
             )
             load_dict[key] = val
-
-        load_state_dict_impl(
-            state_dict=load_dict,
-            path=path,
-            process_group=process_group,
-            coordinator_rank=coordinator_rank,
-            unique_id=unique_id,
-            offload=offload,
-            mw_name_compatibility=mw_name_compatibility,
-            safetensors=safetensors,
-            worker_groups=worker_groups,
+        destination_state_shard_info = defaultdict(list)
+        for key, val in load_dict.items():
+            desc = build_shard_desc(val)
+            destination_state_shard_info[key].append(desc)
+    else:
+        flat_shards, nonflat_shards = _split_flat_shards(state_dict)
+        load_dict, padding_info = _unflatten_shards(flat_shards)
+        load_dict.update(nonflat_shards)
+        destination_state_shard_info = build_global_state_shard_info(
+            state_dict, process_group
         )
-        return
-
-    destination_state_shard_info = build_global_state_shard_info(
-        state_dict, process_group
-    )
-
-    flat_shards, nonflat_shards = _split_flat_shards(state_dict)
-    load_dict, padding_info = _unflatten_shards(flat_shards)
-    load_dict.update(nonflat_shards)
 
     if aoa_config is not None:
         _handle_aoa(
@@ -935,7 +925,8 @@ def load_state_dict(
             safetensors=safetensors,
             worker_groups=worker_groups,
         )
-    _finish_unflatten(flat_shards, padding_info)
+    if use_dist:
+        _finish_unflatten(flat_shards, padding_info)
 
     global _metadata_manager
     _metadata_manager.clear()