move dsv3 model from paddlenlp-dsv3-sft

Author: Difers

examples/run_finetune.py

@@ -30,6 +30,10 @@
     get_last_checkpoint,
     set_seed,
 )
+from paddleformers.trainer.trainer_callback import (
+    MoECorrectionBiasAdjustCallback,
+    MoeExpertsGradScaleCallback,
+)
 from paddleformers.transformers import (
     AutoConfig,
     AutoModelForCausalLM,
@@ -86,6 +90,19 @@
 ]
 
 
+def mock_offload_optimizer():
+    """
+    mock offload optimizer
+    """
+    try:
+        from paddleformers.trainer.utils.offload_optimizer import hack_offload_optimizer
+
+        hack_offload_optimizer()
+        logger.warning("hack_offload_optimizer called.")
+    except ImportError:
+        logger.warning("hack_offload_optimizer is not imported")
+
+
 def main():
     parser = PdArgumentParser((ModelConfig, DataConfig, SFTConfig))
     if len(sys.argv) >= 2 and sys.argv[1].endswith(".json"):
@@ -97,9 +114,18 @@ def main():
     else:
         model_args, data_args, training_args = parser.parse_args_into_dataclasses()
 
+    if training_args.tensorwise_offload_optimizer:
+        mock_offload_optimizer()
+
     training_args.print_config(model_args, "Model")
     training_args.print_config(data_args, "Data")
 
+    if training_args.pre_alloc_memory > 0:
+        memory_size = int(training_args.pre_alloc_memory * 1024 * 1024 * 1024)
+        x = paddle.empty([memory_size], dtype=paddle.uint8)
+        logger.info(f"pre_alloc_memory size {x.shape}")
+        del x
+
     # Setup GPU & distributed training
     paddle.set_device(training_args.device)
     set_seed(seed=training_args.seed)
@@ -171,6 +197,7 @@ def main():
     model_config.max_sequence_length = training_args.max_seq_len
     model_config.num_nextn_predict_layers = model_args.num_nextn_predict_layers
     model_config._attn_implementation = model_args.attn_impl
+    model_config.gradient_accumulation_steps = training_args.gradient_accumulation_steps
     logger.info(f"Final model config: {model_config}")
     logger.info("Creating model")
 
@@ -181,6 +208,11 @@ def main():
 
         model_class = AutoModelForCausalLMPipe
 
+    model_config.using_flex_token = model_args.using_flex_token
+    model_config.using_fake_gate = model_args.using_fake_gate
+    model_config.moe_subbatch_token_num = model_args.moe_subbatch_token_num
+    model_config.aux_loss_alpha = model_args.aux_loss_alpha
+
     if model_args.continue_training and not training_args.autotuner_benchmark:
         model = model_class.from_pretrained(
             model_args.model_name_or_path,
@@ -309,6 +341,17 @@ def neft_post_hook(module, input, output):
         training_args.logging_strategy = IntervalStrategy.STEPS
         training_args.logging_steps = int(training_args.max_steps / training_args.num_train_epochs)
 
+    callbacks = []
+
+    if getattr(model_config, "topk_method", None) == "noaux_tc":
+        # deepseek_v3 finetune do not update the bias, so set lr to 0.0
+        callbacks += [MoECorrectionBiasAdjustCallback(lr=0.0)]
+
+    if training_args.use_expert_parallel:
+        callbacks += [MoeExpertsGradScaleCallback(training_args)]
+
+    print("callbacks:", callbacks, flush=True)
+
     trainer = SFTTrainer(
         model=model,
         args=training_args,
@@ -319,6 +362,7 @@ def neft_post_hook(module, input, output):
         data_collator=data_collator,
         do_generation=data_args.eval_with_do_generation,
         data_args=data_args,
+        callbacks=callbacks,
     )
     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)