【FlexCheckpoint】Upgrade some macros and optimize load_state_dict communication (#75282)

* upgrad macros and load_state_dict comm task

fix

fix

support 0-d tensor

fix

balance save and fix

* fix test

Author: xingmingyyj

python/paddle/distributed/fleet/meta_optimizers/dygraph_optimizer/dygraph_sharding_optimizer.py

@@ -1343,9 +1343,13 @@ def _create_sharded_weight(
         master_weights = optim_state_dict.pop("master_weights", None)
         optim_state_dict.pop("LR_Scheduler", None)
 
-        static_to_struct = {
-            v.local_tensor.name: k for k, v in model_sharded_state_dict.items()
-        }
+        static_to_struct = {}
+        model_sharded_state_dict = dict(
+            sorted(model_sharded_state_dict.items())
+        )
+        for k, v in model_sharded_state_dict.items():
+            if v.local_tensor.name not in static_to_struct:
+                static_to_struct[v.local_tensor.name] = k
 
         sharded_state = {}
 

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

@@ -89,13 +89,13 @@ def get_num_hidden_layers(
             )
         prefix, suffix = name_with_layer_id.split(layer_id_macro_tag, 1)
         pattern = re.compile(rf"{re.escape(prefix)}(\d+){re.escape(suffix)}")
-        max_layer = 0
+        match_layer_id = set()
         for key in self.get_all_dst_state_keys():
             match = pattern.fullmatch(key)
             if match:
                 layer_num = int(match.group(1))
-                max_layer = max(max_layer, layer_num)
-        return max_layer + 1
+                match_layer_id.add(layer_num)
+        return match_layer_id
 
     def get_src_state_shard_num(self, src_state_key: str) -> int:
         if src_state_key not in self.source_state_shard_info:

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

@@ -58,6 +58,20 @@ def register_macro(self, name, func, priority):
 ]
 
 
+def extract_axis_and_clean_tokens(tokens):
+    axis = 1
+    for idx, tkn in enumerate(tokens):
+        if tkn.value == "axis" and idx + 2 < len(tokens):
+            axis = int(tokens[idx + 2].value)
+            end_idx = idx + 3
+            if end_idx < len(tokens) - 1:
+                assert tokens[end_idx].value == ","
+                end_idx += 1
+            tokens = tokens[:idx] + tokens[end_idx:]
+            break
+    return axis, tokens
+
+
 # star_macro must be called after layer_id_macro
 @macro(name='star_macro', priority=3)
 def star_macro(tokens, expression, context):
@@ -119,12 +133,14 @@ def layer_id_macro(tokens, expression, context):
     )
     assert name_with_layer_id, "No $LAYER_ID found in NAME tokens"
 
-    num_layers = context.get_num_hidden_layers(
+    match_layer_id = context.get_num_hidden_layers(
         name_with_layer_id, LAYER_ID_MACRO_TAG
     )
     expanded_expressions = []
 
-    for layer_id in range(num_layers):
+    match_layer_id = sorted(match_layer_id)
+
+    for layer_id in match_layer_id:
         expr = ""
         for token in tokens:
             if token.type == TokenType.IDENTIFIER:
@@ -181,6 +197,8 @@ def fused_qkv_old_macro(tokens, expression, context):
     if not any(tkn.value == FUSED_QKV_OLD_TAG for tkn in tokens):
         return expression
 
+    axis, tokens = extract_axis_and_clean_tokens(tokens)
+
     attn_head_num = None
     num_key_value_groups = None
     fused_qkv_old_pos = None
@@ -263,10 +281,14 @@ def gen_expr(tp_degree, num_heads, tp_rank, comp):
                 for c, n in head_config
             ]
             if idx == 0:
-                mapping = f"{qkv_weight_name} -> {','.join(qkv_parts)}, axis=1"
+                mapping = (
+                    f"{qkv_weight_name} -> {','.join(qkv_parts)}, axis={axis}"
+                )
                 results.append(mapping)
             elif qkv_weight_name is not None:
-                mapping = f"{','.join(qkv_parts)} -> {qkv_weight_name}, axis=1"
+                mapping = (
+                    f"{','.join(qkv_parts)} -> {qkv_weight_name}, axis={axis}"
+                )
                 results.append(mapping)
 
         if fused_qkv_old_pos > 4:
@@ -275,7 +297,7 @@ def _generate_expr(prefix, count, target_name):
                 elements = ",".join(
                     f"fused_qkv_old_tmp.{prefix}_{i}" for i in range(count)
                 )
-                return f"{elements} -> {target_name}, axis=1"
+                return f"{elements} -> {target_name}, axis={axis}"
 
             q_name = tokens[2].value
             k_name = tokens[4].value
@@ -292,7 +314,7 @@ def _generate_expr(prefix, count, target_name):
 
         fused_qkv_tmp_name = f"{q_name}.{k_name}.{v_name}.tmp"
         results.append(
-            f"{q_name},{k_name},{v_name}  ->  {fused_qkv_tmp_name}, axis=1"
+            f"{q_name},{k_name},{v_name}  ->  {fused_qkv_tmp_name}, axis={axis}"
         )
         dst_state_shard_num = context.get_dst_state_shard_num(
             dst_qkv_weight_name
@@ -324,9 +346,13 @@ def gen_expr(tp_degree, num_heads, tp_rank, comp):
                 for c, n in head_config
             ]
             if idx == 0:
-                mapping = f"{qkv_weight_name} -> {','.join(qkv_parts)}, axis=1"
+                mapping = (
+                    f"{qkv_weight_name} -> {','.join(qkv_parts)}, axis={axis}"
+                )
             else:
-                mapping = f"{','.join(qkv_parts)} -> {qkv_weight_name}, axis=1"
+                mapping = (
+                    f"{','.join(qkv_parts)} -> {qkv_weight_name}, axis={axis}"
+                )
             results.append(mapping)
     else:
         raise ValueError(
@@ -340,6 +366,9 @@ def fused_ffn_macro(tokens, expression, context):
     FUSED_FFN_TAG = "fused_ffn"
     if not any(tkn.value == FUSED_FFN_TAG for tkn in tokens):
         return expression
+
+    axis, tokens = extract_axis_and_clean_tokens(tokens)
+
     rarrow_pos = None
     fused_ffn_pos = None
     for idx, token in enumerate(tokens):
@@ -388,19 +417,19 @@ def gen_expr(tp_degree, splited_num, tp_rank, comp):
             ]
             if idx == 0:
                 results.append(
-                    f"{ffn_weight_name}  -> {','.join(ffn_parts)}, axis=1"
+                    f"{ffn_weight_name}  -> {','.join(ffn_parts)}, axis={axis}"
                 )
             elif ffn_weight_name is not None:
                 results.append(
-                    f"{','.join(ffn_parts)} -> {ffn_weight_name}, axis=1"
+                    f"{','.join(ffn_parts)} -> {ffn_weight_name}, axis={axis}"
                 )
         if fused_ffn_pos > 4:
 
             def _generate_expr(prefix, count, target_name):
                 elements = ",".join(
                     f"fused_ffn_tmp.{prefix}_{i}" for i in range(count)
                 )
-                return f"{elements} -> {target_name}, axis=1"
+                return f"{elements} -> {target_name}, axis={axis}"
 
             gate_name = tokens[2].value
             up_name = tokens[4].value
@@ -415,7 +444,7 @@ def _generate_expr(prefix, count, target_name):
 
         fused_gate_up_tmp_name = f"{gate_name}.{up_name}.tmp"
         results.append(
-            f"{gate_name},{up_name}  ->  {fused_gate_up_tmp_name}, axis=1"
+            f"{gate_name},{up_name}  ->  {fused_gate_up_tmp_name}, axis={axis}"
         )
         dst_state_shard_num = context.get_dst_state_shard_num(
             dst_ffn_weight_name
@@ -445,11 +474,11 @@ def gen_expr(tp_degree, splited_num, tp_rank, comp):
             ]
             if idx == 0:
                 results.append(
-                    f"{ffn_weight_name}  -> {','.join(ffn_parts)}, axis=1"
+                    f"{ffn_weight_name}  -> {','.join(ffn_parts)}, axis={axis}"
                 )
             else:
                 results.append(
-                    f"{','.join(ffn_parts)} -> {ffn_weight_name}, axis=1"
+                    f"{','.join(ffn_parts)} -> {ffn_weight_name}, axis={axis}"
                 )
     else:
         raise ValueError(f"Unsupported fused_ffn macro format: {expression}.")
@@ -508,6 +537,8 @@ def fused_qkv(tokens, expression, context):
     if not any(tkn.value == FUSED_QKV_TAG for tkn in tokens):
         return expression
 
+    axis, tokens = extract_axis_and_clean_tokens(tokens)
+
     attn_head_num = num_heads = None
     num_key_value_groups = None
     fused_qkv_pos = None
@@ -566,12 +597,12 @@ def make_names(base, n):
             fused_qkv_order.append(k_names[g])
             fused_qkv_order.append(v_names[g])
         results.append(
-            f"{fused_qkv_var} -> {','.join(fused_qkv_order)}, axis=1"
+            f"{fused_qkv_var} -> {','.join(fused_qkv_order)}, axis={axis}"
         )
 
-        results.append(f"{','.join(q_names)} -> {q_var}, axis=1")
-        results.append(f"{','.join(k_names)} -> {k_var}, axis=1")
-        results.append(f"{','.join(v_names)} -> {v_var}, axis=1")
+        results.append(f"{','.join(q_names)} -> {q_var}, axis={axis}")
+        results.append(f"{','.join(k_names)} -> {k_var}, axis={axis}")
+        results.append(f"{','.join(v_names)} -> {v_var}, axis={axis}")
 
         return results
 
@@ -585,9 +616,9 @@ def make_names(base, n):
         k_names = make_names(k_var, num_key_value_groups)
         v_names = make_names(v_var, num_key_value_groups)
 
-        results.append(f"{q_var} -> {','.join(q_names)}, axis=1")
-        results.append(f"{k_var} -> {','.join(k_names)}, axis=1")
-        results.append(f"{v_var} -> {','.join(v_names)}, axis=1")
+        results.append(f"{q_var} -> {','.join(q_names)}, axis={axis}")
+        results.append(f"{k_var} -> {','.join(k_names)}, axis={axis}")
+        results.append(f"{v_var} -> {','.join(v_names)}, axis={axis}")
 
         fused_qkv_order = []
         for g in range(num_key_value_groups):
@@ -597,7 +628,7 @@ def make_names(base, n):
             fused_qkv_order.append(k_names[g])
             fused_qkv_order.append(v_names[g])
         results.append(
-            f"{','.join(fused_qkv_order)} -> {fused_qkv_var}, axis=1"
+            f"{','.join(fused_qkv_order)} -> {fused_qkv_var}, axis={axis}"
         )
         return results