[AutoParallel] GradientClipByGlobalNorm: Add align mode branch in autotrainer (#10960)

* Add test branch for align mode

* fix test bug

* rerun CI

* add the warning to warn the max_grad_norm is not 0.0

Author: zty-king

paddlenlp/trainer/auto_trainer.py

@@ -16,6 +16,7 @@
 import os
 import random
 import time
+import types
 from typing import Any, Dict, Optional, Union
 
 import numpy as np
@@ -24,6 +25,7 @@
 import paddle.distributed.auto_parallel.intermediate.parallelize as parallelize
 import paddle.nn as nn
 from paddle.distributed import fleet
+from paddle.distributed.auto_parallel._utils import _patch_grads_for_step
 from paddle.profiler.utils import switch_job_schedule_profiler
 from tqdm.auto import tqdm
 
@@ -518,6 +520,18 @@ def _inner_training_loop(
             npu_accelerate_plugin(self.optimizer)
 
         model, dist_loader = self._wrap_for_auto(model, train_dataloader)
+
+        if (
+            dist.in_auto_parallel_align_mode()
+        ):  # When in auto parallel align mode, patching the optimizer step function
+
+            orig_step = (
+                self.optimizer.step.__func__ if hasattr(self.optimizer.step, "__func__") else self.optimizer.step
+            )
+            decorator = _patch_grads_for_step(amp_master_grad=self.args.amp_master_grad)
+            new_step = decorator(orig_step)
+            self.optimizer.__dict__["step"] = types.MethodType(new_step, self.optimizer)
+
         train_dataloader = dist_loader()
         if resume_from_checkpoint is not None:
             self._load_from_checkpoint(resume_from_checkpoint)

paddlenlp/trainer/training_args.py

@@ -1109,7 +1109,12 @@ class TrainingArguments:
     def __post_init__(self):
         world_size = paddle.distributed.get_world_size()
         if in_auto_parallel_align_mode():
-            self.max_grad_norm = 0.0
+            # self.max_grad_norm = 0.0
+            # The current auto_hybrid_pp has aligned the handling of ClipGradByGlobalNorm with the original dygraph semi-auto parallel and dynamic manual-parallel modes and can correctly handle grad_clip, so it is no longer necessary to set max_grad_norm=0.0.
+            if self.max_grad_norm != 0.0:
+                warnings.warn(
+                    "max_grad_norm is not 0.0,We will execute ClipGradByGlobalNorm,if you want to disable it,please set max_grad_norm=0.0"
+                )
             os.environ["FLAGS_max_inplace_grad_add"] = "65536"
             os.environ["FLAGS_embedding_deterministic"] = "1"
             os.environ["FLAGS_cudnn_deterministic"] = "1"

scripts/distribute/ci_case_auto.sh

@@ -935,7 +935,7 @@ function llama_dygraph_auto_bs8_fp16_DP2-MP2-PP2_hybrid_pp() {
             --data_impl "mmap" \
             --enable_auto_parallel 1 \
             --to_static 0 \
-            --max_grad_norm 0.0 \
+            --max_grad_norm 1.0 \
             >>${log_path}/$FUNCNAME 2>&1
         loss=`cat $case_log_dir/workerlog.0 | grep 'global_step: 10' | awk -F 'loss: ' '{print $2}' | awk -F ',' '{print $1}'`
         ips=-1
@@ -1003,7 +1003,7 @@ function llama_dygraph_auto_bs8_fp16_DP2-MP2-PP2_hybrid_pp() {
             --data_impl "mmap" \
             --enable_auto_parallel 1 \
             --to_static 0 \
-            --max_grad_norm 0.0 \
+            --max_grad_norm 1.0 \
             --resume_from_checkpoint "${case_out_dir}/checkpoint-9" \
             >>${log_path}/$FUNCNAME 2>&1
         pp_resume_from_hybrid_ckpt_loss=`cat $auto_case_log_dir/workerlog.0 | grep 'global_step: 10' | awk -F 'loss: ' '{print $2}' | awk -F ',' '{print $1}'`
@@ -1012,6 +1012,69 @@ function llama_dygraph_auto_bs8_fp16_DP2-MP2-PP2_hybrid_pp() {
         echo "pp_resume from hybrid ckpt result: loss=$pp_resume_from_hybrid_ckpt_loss ips=$pp_resume_from_hybrid_ckpt_ips mem=$pp_resume_from_hybrid_ckpt_mem"
 
         check_result $FUNCNAME ${loss} ${pp_resume_from_hybrid_ckpt_loss} ${ips} ${pp_resume_from_hybrid_ckpt_ips} ${mem} ${pp_resume_from_hybrid_ckpt_mem}
+        
+        echo "=========== $FUNCNAME run dygraph auto hybrid pp in align mode ==========="
+        export FLAGS_enable_auto_parallel_align_mode=1
+        task_name="llama_auto_bs8_fp16_dp2mp2pp2_hybrid_pp_in_align_mode"
+        align_mode_case_out_dir="output/$task_name"
+        align_mode_case_log_dir="output/$task_name""_log"
+        rm -rf $align_mode_case_out_dir
+        rm -rf $align_mode_case_log_dir
+
+        python -u -m paddle.distributed.launch --gpus "0,1,2,3,4,5,6,7" --log_dir $align_mode_case_log_dir run_pretrain_auto.py \
+            --model_type "llama_pp" \
+            --model_name_or_path "facebook/llama-7b" \
+            --tokenizer_name_or_path "facebook/llama-7b" \
+            --input_dir "./data" \
+            --output_dir $align_mode_case_out_dir \
+            --split 949,50,1 \
+            --max_seq_length 2048 \
+            --hidden_size 1024 \
+            --intermediate_size 3072 \
+            --num_hidden_layers 8 \
+            --num_attention_heads 32 \
+            --per_device_train_batch_size 4 \
+            --per_device_eval_batch_size 4 \
+            --n_microbatch 4 \
+            --gradient_accumulation_steps 1 \
+            --use_flash_attention 1 \
+            --use_fused_rms_norm 0 \
+            --fp16 1 \
+            --fp16_opt_level "O2" \
+            --amp_master_grad 1 \
+            --scale_loss 1024 \
+            --pipeline_parallel_degree 2 \
+            --pipeline_schedule_mode "FThenB" \
+            --tensor_parallel_degree 2 \
+            --sharding_parallel_degree 1 \
+            --learning_rate 0.0001 \
+            --min_learning_rate 0.00001 \
+            --max_steps 10 \
+            --save_steps 20 \
+            --weight_decay 0.01 \
+            --warmup_ratio 0.01 \
+            --logging_steps 1 \
+            --dataloader_num_workers 1 \
+            --sharding "" \
+            --eval_steps 1000000 \
+            --disable_tqdm true \
+            --continue_training 0 \
+            --recompute 0 \
+            --do_train \
+            --do_eval \
+            --device "gpu" \
+            --data_impl "mmap" \
+            --enable_auto_parallel 1 \
+            --to_static 0 \
+            --max_grad_norm 1.0 \
+            >>${log_path}/$FUNCNAME 2>&1
+        align_mode_loss=`cat $align_mode_case_log_dir/workerlog.0 | grep 'global_step: 10' | awk -F 'loss: ' '{print $2}' | awk -F ',' '{print $1}'`
+        align_mode_ips=-1
+        align_mode_mem=-1
+        echo "result: loss=$align_mode_loss ips=$align_mode_ips mem=$align_mode_mem"
+
+        check_result $FUNCNAME ${loss} ${align_mode_loss} ${ips} ${align_mode_ips} ${mem} ${align_mode_mem}
+        
         echo "=========== $FUNCNAME run  end ==========="
     fi
 }