【AutoParallel】Support 'master_grad' in Llama in static auto-parallelism (PaddlePaddle#7658)

* add master_grad

* add llama_fp16_test

* polish

* merge develop

Author: heavyrain-lzy

paddlenlp/trainer/training_args.py

@@ -1174,9 +1174,6 @@ def is_segment_parallel_supported():
                 pipeline.micro_batch_size = self.per_device_train_batch_size
                 pipeline.schedule_mode = self.pipeline_schedule_mode
 
-                if self.amp_master_grad:
-                    warnings.warn("`amp_master_grad` is not supported NOW in AutoParallel!")
-                    self.amp_master_grad = False
                 logger.info(f"PP configs:{strategy.pipeline}, use master_grad: {self.amp_master_grad}")
 
                 if self.do_eval:
@@ -1260,6 +1257,7 @@ def is_segment_parallel_supported():
                 amp.enable = True
                 amp.dtype = "bfloat16" if self.bf16 else "float16"
                 amp.level = self.fp16_opt_level.lower()
+                amp.use_master_grad = self.amp_master_grad
                 amp.init_loss_scaling = self.scale_loss
                 amp.custom_black_list = self.amp_custom_black_list if self.amp_custom_black_list is not None else []
                 amp.custom_white_list = self.amp_custom_white_list if self.amp_custom_white_list is not None else []

scripts/distribute/ci_case_auto.sh

@@ -52,6 +52,7 @@ function llama_case_list_auto() {
     llama_static_auto_recompute_bs16_fp32_DP2-MP2-PP1
     llama_static_auto_recompute_bs16_fp32_DP2-MP2-PP2
     llama_static_auto_recompute_bs16_fp32_DP2-MP2-PP2-VPP2-Sharding2_stage2
+    llama_static_auto_recompute_bs16_fp16_DP2-MP2-PP2-VPP2-Sharding2_stage2
 }
 
 function gpt_case_list_auto_pir() {
@@ -1168,6 +1169,75 @@ function llama_static_auto_recompute_bs16_fp32_DP2-MP2-PP2-VPP2-Sharding2_stage2
     echo "=========== $FUNCNAME run  end ==========="
 }
 
+function llama_static_auto_recompute_bs16_fp16_DP2-MP2-PP2-VPP2-Sharding2_stage2() {
+    echo "=========== $FUNCNAME run begin ==========="
+    export PYTHONPATH=$root_path/:$PYTHONPATH
+    export FLAGS_call_stack_level=2
+
+    task_name="llama_auto_bs16_fp16_dp2mp2pp2vpp2sharding2"
+    case_out_dir="output/$task_name"
+    case_log_dir="output/$task_name""_log"
+    rm -rf $case_out_dir
+    rm -rf $case_log_dir
+
+    python -u -m paddle.distributed.launch --gpus "0,1,2,3,4,5,6,7" --log_dir $case_log_dir run_pretrain_auto.py \
+        --model_type "llama" \
+        --model_name_or_path "facebook/llama-7b" \
+        --tokenizer_name_or_path "facebook/llama-7b" \
+        --hidden_size 1024 \
+        --intermediate_size 3072 \
+        --num_hidden_layers 8 \
+        --num_attention_heads 32 \
+        --input_dir "./data" \
+        --output_dir $case_out_dir \
+        --split 949,50,1 \
+        --max_seq_length 2048 \
+        --per_device_train_batch_size 1 \
+        --per_device_eval_batch_size 8 \
+        --gradient_accumulation_steps 8 \
+        --use_flash_attention 0 \
+        --use_fused_rms_norm 0 \
+        --fp16 1 \
+        --fp16_opt_level "O2"  \
+        --amp_master_grad 1 \
+        --scale_loss 1024 \
+        --tensor_parallel_degree 2 \
+        --pipeline_parallel_degree 2 \
+        --virtual_pp_degree 2 \
+        --pipeline_schedule_mode "VPP" \
+        --sharding_parallel_degree 2 \
+        --sharding "stage2" \
+        --learning_rate 0.0001 \
+        --min_learning_rate 0.00001 \
+        --max_steps 10 \
+        --save_steps 5000 \
+        --weight_decay 0.01 \
+        --warmup_ratio 0.01 \
+        --max_grad_norm 1.0 \
+        --logging_steps 1 \
+        --dataloader_num_workers 1 \
+        --eval_steps 1000 \
+        --report_to "visualdl" \
+        --disable_tqdm true \
+        --continue_training 0 \
+        --recompute 1 \
+        --do_train \
+        --do_eval \
+        --device "gpu" \
+        --data_impl "mmap" \
+        --parallel_mode "auto" \
+        >>${log_path}/$FUNCNAME 2>&1
+    loss=`cat $case_log_dir/workerlog.3 | grep 'global_step: 10' | awk -F 'loss: ' '{print $2}' | awk -F ',' '{print $1}'`
+    ips=-1
+    mem=-1
+    echo "result: loss=$loss ips=$ips mem=$mem"
+    loss_base=10.0859375
+    ips_base=-1
+    mem_base=-1
+    check_result $FUNCNAME ${loss_base} ${loss} ${ips_base} ${ips} ${mem_base} ${mem}
+    echo "=========== $FUNCNAME run  end ==========="
+}
+
 function llama_dygraph_auto_bs4_fp32_DP2-MP2-PP2() {
     echo "=========== $FUNCNAME run begin ==========="
     export PYTHONPATH=$root_path/:$PYTHONPATH
@@ -1233,7 +1303,6 @@ function llama_dygraph_auto_bs4_fp32_DP2-MP2-PP2() {
     check_result $FUNCNAME ${loss_base} ${loss} ${ips_base} ${ips} ${mem_base} ${mem}
     echo "=========== $FUNCNAME run  end ==========="
 }
-
 ############ case end ############
 
 function check_result() {

scripts/distribute/run_ci.sh

@@ -22,11 +22,13 @@ export case_list=()
 
 target_lists_for_gpt=(
     "model_zoo/gpt-3"
+    "scripts/distribute"
 )
 
 target_lists_for_llama=(
     "llm/llama/auto_parallel"
     "paddlenlp/transformers/llama/modeling_auto.py"
+    "scripts/distribute"
 )
 
 target_path_for_ci_scripts="scripts/distribute"