causal mtp [hard coding].

Author: wtmlon

paddleformers/cli/train/sft/workflow.py

@@ -17,6 +17,7 @@
 import gc
 import math
 import os
+import re
 from dataclasses import fields
 from functools import partial
 
@@ -85,6 +86,29 @@
 )
 
 
+def frozen_param_expect_mtp(model, config):
+    def extract_layer_idx(text):
+        match = re.search(r"model.layers.(-?\d+\.?\d*)", text)
+        if match:
+            num_str = match.group(1)
+            # 区分整数和小数返回（避免123.0这种冗余浮点数）
+            if "." in num_str:
+                return float(num_str)
+            else:
+                return int(num_str)
+        return None
+
+    # not sure can work on all model
+    jackpot = set(range(config.num_hidden_layers, config.num_hidden_layers + config.num_nextn_predict_layers))
+    for name, param in model.state_dict().items():
+        layer_idx = extract_layer_idx(name)
+        is_mtp = layer_idx in jackpot
+        if not is_mtp:
+            param.stop_gradient = True
+        else:
+            param.stop_gradient = False
+
+
 def create_pretrained_dataset(training_args, data_args, model_args):
     assert data_args.input_dir is not None and len(data_args.input_dir.split()) > 1
 
@@ -653,6 +677,8 @@ def neft_post_hook(module, input, output):
         callbacks += [FP8QuantWeightCallback()]
 
     print("callbacks:", callbacks, flush=True)
+    # print("ddd: ", model); exit()
+
     trainer = SFTTrainer(
         model=model,
         args=training_args,
@@ -665,6 +691,7 @@ def neft_post_hook(module, input, output):
         data_args=data_args,
         callbacks=callbacks,
     )
+    frozen_param_expect_mtp(model, model_config)
     trainable_parameters = [p for p in model.parameters() if not p.stop_gradient]
     trainer.set_optimizer_grouped_parameters(trainable_parameters)
 

paddleformers/trainer/trainer.py

@@ -3147,6 +3147,8 @@ def new_global_norm_func(
                 global_norm_var_not_dist,
                 *args,
             ):
+                print("WE DO NOT CAL GRAD NORM FOR NOW")
+                return
                 if len(args) > 0:
                     global_norm_func(global_norm_var_dist, global_norm_var_not_dist, *args)
                     global_norm_var_dist_moe, global_norm_var_not_dist_moe = args
@@ -3157,6 +3159,8 @@ def new_global_norm_func(
                         + global_norm_var_not_dist_moe
                     )
                 else:
+                    print("global_norm_var_dist: ", global_norm_var_dist)
+                    print("global_norm_var_not_dist: ", global_norm_var_not_dist)
                     global_norm_func(global_norm_var_dist, global_norm_var_not_dist)
                     global_norm_var_fp32 = paddle.sqrt(global_norm_var_dist + global_norm_var_not_dist)
                 training_logs["global_norm"] = global_norm_var_fp32.item()

paddleformers/transformers/glm4_moe/modeling.py

@@ -911,6 +911,7 @@ def _gen_aoa_config(cls, config: Glm4MoeConfig):
         ]
 
         num_nextn_predict_layers = config.num_nextn_predict_layers if config.num_nextn_predict_layers else 0
+        num_nextn_predict_layers = 1
 
         for layer_idx in reversed(range(num_hidden_layers, num_hidden_layers + num_nextn_predict_layers)):
             layer_idx_offset = layer_idx + num_head_empty_layers
@@ -1057,6 +1058,7 @@ def _gen_inv_aoa_config(cls, config: Glm4MoeConfig):
         ]
 
         num_nextn_predict_layers = config.num_nextn_predict_layers if config.num_nextn_predict_layers else 0
+        num_nextn_predict_layers = 1
 
         for layer_idx in reversed(range(num_hidden_layers, num_hidden_layers + num_nextn_predict_layers)):
             layer_idx_offset = layer_idx + num_head_empty_layers