soft coding.

Author: wtmlon

paddleformers/cli/train/sft/workflow.py

@@ -99,7 +99,7 @@ def extract_layer_idx(text):
         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))
+    jackpot = set(range(config.num_hidden_layers, config.num_hidden_layers + config.mtp_num_layers))
     for name, param in model.state_dict().items():
         layer_idx = extract_layer_idx(name)
         is_mtp = layer_idx in jackpot
@@ -288,6 +288,22 @@ def run_sft(
     LlmMetaConfig.set_llm_config(model_config, training_args)
     model_config.use_fast_layer_norm = model_args.use_fast_layer_norm
 
+    # autoregressive mtp training
+    activate_autoregressive_mtp_training = False
+    if model_config.mtp_num_layers > 1:
+        activate_autoregressive_mtp_training = True
+        tmp = model_config.mtp_num_layers
+        model_config.mtp_num_layers = model_config.num_nextn_predict_layers
+        model_config.num_nextn_predict_layers = tmp
+
+        tmp = training_args.mtp_num_layers
+        training_args.mtp_num_layers = training_args.num_nextn_predict_layers
+        training_args.num_nextn_predict_layers = tmp
+
+        logger.info(
+            f"MTP args changing for autoregressive mtp training, mtp_num_layers: {model_config.mtp_num_layers}, num_nextn_predict_layers: {model_config.num_nextn_predict_layers}!!"
+        )
+
     # Config for model using dropout, such as GPT.
     if hasattr(model_config, "hidden_dropout_prob"):
         model_config.hidden_dropout_prob = finetuning_args.hidden_dropout_prob
@@ -699,10 +715,14 @@ def fetch_and_serialize(generator, dtype):
         data_args=data_args,
         callbacks=callbacks,
     )
+    freeze_param_except_mtp(model, model_config)
+
+    if activate_autoregressive_mtp_training:
+        # activate autoregressive mtp training
+        freeze_param_except_mtp(model, model_config)
     trainable_parameters = [
         p for p in model.parameters() if not p.stop_gradient or ("quantization_linear" in p.name and "w_1" in p.name)
     ]
-    freeze_param_except_mtp(model, model_config)
     trainer.set_optimizer_grouped_parameters(trainable_parameters)
 
     # Train

paddleformers/trainer/training_args.py

@@ -1212,6 +1212,10 @@ class TrainingArguments:
         metadata={"help": "pre allocate memory size GB"},
     )
     num_nextn_predict_layers: int = field(default=0, metadata={"help": "Number of nextn predict layers."})
+    mtp_distillation_loss: bool = field(default=False, metadata={"help": "Whether to use distillation MTP loss."})
+    mtp_num_layers: int = field(
+        default=0, metadata={"help": "Whether to use Autoregressive MTP Training, activate if > 1."}
+    )
     profile: bool = field(default=False, metadata={"help": "Enable nsys profiling."})
     profile_step_start: int = field(default=10, metadata={"help": "Step to start nsys profiling."})
     profile_step_end: int = field(default=12, metadata={"help": "Step to end nsys profiling."})

paddleformers/transformers/configuration_utils.py

@@ -413,7 +413,9 @@ class LlmMetaConfig:
     ]
 
     mtp_attributes = [
-        ("num_nextn_predict_layers", int, 0, "Number of nextn predict layers."),
+        # ("num_nextn_predict_layers", int, 0, "Number of nextn predict layers."),
+        ("mtp_distillation_loss", bool, False, "Whether to use distillation MTP loss."),
+        ("mtp_num_layers", int, 0, "Whether to use Autoregressive MTP Training, activate if > 1."),
         (
             "mtp_loss_scaling_factor",
             float,