[AutoParallel] add use custom op test in pipeline mode (#10653)

* add use rms norm ci test in pp mode

* fix loss base

* fix base

Author: waliwali777

scripts/distribute/ci_case_auto.sh

@@ -472,64 +472,67 @@ function llama_dygraph_auto_bs8_fp32_DP2-MP2-PP2() {
     task_name="llama_auto_bs8_dp2mp2pp2"
     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" \
-        --input_dir "./data" \
-        --output_dir $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 1 \
-        --per_device_eval_batch_size 4 \
-        --gradient_accumulation_steps 4 \
-        --use_flash_attention 0 \
-        --use_fused_rms_norm 0 \
-        --fp16 0 \
-        --fp16_opt_level "O2" \
-        --scale_loss 1024 \
-        --pipeline_parallel_degree 2 \
-        --tensor_parallel_degree 2 \
-        --sharding_parallel_degree 1 \
-        --learning_rate 0.0001 \
-        --min_learning_rate 0.00001 \
-        --max_steps 10 \
-        --save_steps 5000 \
-        --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
-    loss=`cat $case_log_dir/workerlog.0 | 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=9.3513937
-    if [ $IS_A100 -ne 0 ];then
-        loss_base=9.39356422
-    fi
-    ips_base=-1
-    mem_base=-1
-    check_result $FUNCNAME ${loss_base} ${loss} ${ips_base} ${ips} ${mem_base} ${mem}
+    for use_fused_rms_norm in "1" "0"; do
+        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" \
+            --input_dir "./data" \
+            --output_dir $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 1 \
+            --per_device_eval_batch_size 4 \
+            --gradient_accumulation_steps 4 \
+            --use_flash_attention 0 \
+            --use_fused_rms_norm ${use_fused_rms_norm} \
+            --fp16 0 \
+            --fp16_opt_level "O2" \
+            --scale_loss 1024 \
+            --pipeline_parallel_degree 2 \
+            --tensor_parallel_degree 2 \
+            --sharding_parallel_degree 1 \
+            --learning_rate 0.0001 \
+            --min_learning_rate 0.00001 \
+            --max_steps 10 \
+            --save_steps 5000 \
+            --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
+        loss=`cat $case_log_dir/workerlog.0 | grep 'global_step: 10' | awk -F 'loss: ' '{print $2}' | awk -F ',' '{print $1}'`
+        ips=-1
+        mem=-1
+        echo "use_fused_rms_norm=$use_fused_rms_norm  result: loss=$loss ips=$ips mem=$mem"
+        loss_base=9.3513937
+        if [ $IS_A100 -ne 0 ];then
+            loss_base=9.39356422
+        fi
+        ips_base=-1
+        mem_base=-1
+        check_result $FUNCNAME ${loss_base} ${loss} ${ips_base} ${ips} ${mem_base} ${mem}
+    done
     echo "=========== $FUNCNAME run  end ==========="
 }