[AutoParallel] Optimize Reshard [part1: optimize remove_no_need_vars] #59912

AndSonder · 2023-12-11T14:50:48Z

PR types

Performance optimization

PR changes

Others

Description

静态图模式下分析得，_remove_var 中 sync_with_cpp 消耗了大量的时间，耗时图如下所示：

经过分析可得，_remove_var 中每次删除var都会进行同步，这是没有必要的，我们只需要在全部的var删除之后做一次同步即可。经模型验证，该优化不会影响模型精度，模型精度可与优化前对齐。

本地测试优化后模型实际Run之前的耗时，优化幅度达 1.37 倍。测试环境：本地四卡 1080Ti机器，PaddleNLP Llama2 7b模型, 静态图模式（hack config: config.num_hidden_layers = 12）

优化前后 cprofiler 结果如下, <-- !!! 指向的是被优化的项：

优化前：
   ncalls  tottime  percall  cumtime  percall filename:lineno(function)
        1    0.000    0.000   32.811   32.811 run_pretrain_auto.py:414(main)
        1    0.000    0.000   24.414   24.414 engine.py:1465(prepare)
        1    0.000    0.000   24.414   24.414 engine.py:562(_prepare_program)
        1    0.000    0.000   17.678   17.678 engine.py:762(_parallel)
        1    0.000    0.000   17.674   17.674 parallelizer_v2.py:65(parallel)
        1    0.000    0.000   12.932   12.932 reshard.py:2795(reshard) <--!!!
     1625    9.704    0.006   11.517    0.007 framework.py:4568(_sync_with_cpp) 
     1272    0.016    0.000    8.683    0.007 framework.py:4331(_remove_var)
        1    0.000    0.000    8.238    8.238 reshard.py:904(remove_no_need_in_main)
        1    0.006    0.006    8.207    8.207 reshard.py:846(remove_no_need_vars)
7795/6419    0.019    0.000    4.753    0.001 decorator.py:229(fun)
7795/6419    0.010    0.000    4.575    0.001 wrapped_decorator.py:23(__impl__)
        1    0.033    0.033    4.137    4.137 reshard.py:2414(_reshard_input)
      261    0.376    0.001    3.927    0.015 reshard.py:1782(parse_op_desc)
        1    0.000    0.000    3.390    3.390 engine.py:575(_build)
        1    0.000    0.000    3.248    3.248 helper.py:222(build_program)
        3    0.000    0.000    3.175    1.058 program_translator.py:912(concrete_program)
        3    0.000    0.000    3.175    1.058 program_translator.py:942(concrete_program_specify_input_spec)
        3    0.000    0.000    3.175    1.058 program_translator.py:825(get_concrete_program)
        3    0.000    0.000    3.174    1.058 program_translator.py:1637(__getitem__)
优化后：
ncalls  tottime  percall  cumtime  percall filename:lineno(function)
        1    0.000    0.000   23.889   23.889 run_pretrain_auto.py:414(main)
        1    0.000    0.000   15.531   15.531 engine.py:1465(prepare)
        1    0.000    0.000   15.531   15.531 engine.py:562(_prepare_program)
        1    0.000    0.000    9.015    9.015 engine.py:762(_parallel)
        1    0.000    0.000    9.015    9.015 parallelizer_v2.py:65(parallel)
7795/6419    0.019    0.000    4.790    0.001 decorator.py:229(fun)
7795/6419    0.010    0.000    4.610    0.001 wrapped_decorator.py:23(__impl__)
        1    0.000    0.000    4.311    4.311 reshard.py:2797(reshard) <--!!!
        1    0.033    0.033    4.184    4.184 reshard.py:2416(_reshard_input)
      261    0.388    0.001    3.974    0.015 reshard.py:1784(parse_op_desc)
        1    0.000    0.000    3.306    3.306 engine.py:575(_build)
        1    0.000    0.000    3.198    3.198 helper.py:222(build_program)
        3    0.000    0.000    3.125    1.042 program_translator.py:912(concrete_program)
        3    0.000    0.000    3.125    1.042 program_translator.py:942(concrete_program_specify_input_spec)
        3    0.000    0.000    3.125    1.042 program_translator.py:825(get_concrete_program)
        3    0.000    0.000    3.125    1.042 program_translator.py:1637(__getitem__)
        1    0.000    0.000    3.125    3.125 program_translator.py:1558(_build_once)
  264/151    0.001    0.000    3.032    0.020 base.py:66(__impl__)
        1    0.000    0.000    3.015    3.015 program_translator.py:1276(from_func_spec)
        1    0.000    0.000    3.013    3.013 parallel.py:943(init_parallel_env)

优化后的耗时可视化如下：

脚本启动命令：

python -u -m paddle.distributed.launch \
    --gpus "0,1,2,3" \
    --log_dir "output/$task_name""_log" \
    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 "output/$task_name" \
    --split 949,50,1 \
    --max_seq_length 2048 \
    --per_device_train_batch_size 1 \
    --per_device_eval_batch_size 4 \
    --gradient_accumulation_steps 16 \
    --fp16 1 \
    --fp16_opt_level "O2"  \
    --scale_loss 1024 \
    --pipeline_parallel_degree 2 \
    --tensor_parallel_degree 2 \
    --sharding_parallel_degree 1 \
    --sharding "stage1" \
    --learning_rate 0.0001 \
    --min_learning_rate 0.00001 \
    --max_steps 30 \
    --save_steps 5000 \
    --weight_decay 0.01 \
    --warmup_ratio 0.01 \
    --max_grad_norm 1.0 \
    --logging_steps 1\
    --dataloader_num_workers 1 \
    --sharding "" \
    --eval_steps 1000 \
    --report_to "visualdl" \
    --disable_tqdm true \
    --continue_training 0\
    --recompute 1 \
    --do_train 1 \
    --do_eval 0 \
    --device "gpu" \
    --data_impl "mmap" \
    --parallel_mode "auto" \
    --fuse_attention_qkv 1 \
    --fuse_attention_ffn 1 \
    --use_fused_rope 1 \
    --use_flash_attention 1

… opt_remove_no_need_var

paddle-bot · 2023-12-11T14:50:55Z

你的PR提交成功，感谢你对开源项目的贡献!
请关注后续CI自动化测试结果，详情请参考Paddle-CI手册。
Your PR has been submitted. Thanks for your contribution!
Please wait for the result of CI firstly. See Paddle CI Manual for details.

From00 · 2023-12-12T13:08:29Z

python/paddle/distributed/auto_parallel/static/reshard.py

@@ -899,7 +899,8 @@ def remove_no_need_vars(
            for var in remove_vars:
                if var in feed_var_names:
                    continue
-                block._remove_var(var)
+                block._remove_var(var, sync=False)


这个模块中有两处类似的_remove_var操作，两处都应该修改。另外，请在PR中更新修改之后的程序调用热点图，以确认之前图中的耗时操作在本PR优化之后的对应变化。

From00

LGTM

AndSonder added 2 commits December 11, 2023 14:36

optimize remove no need var

92ffb36

Merge branch 'develop' of https://github.com/PaddlePaddle/Paddle into…

9d57ca5

… opt_remove_no_need_var

paddle-bot bot added the contributor External developers label Dec 11, 2023

From00 reviewed Dec 12, 2023

View reviewed changes

apply suggestions from code review

ce813dd

AndSonder requested a review from From00 December 14, 2023 02:47

From00 approved these changes Dec 14, 2023

View reviewed changes

From00 merged commit c0445a4 into PaddlePaddle:develop Dec 14, 2023

AndSonder changed the title ~~[AutoParallel] Optimize remove_no_need_vars~~ [AutoParallel] Optimize Reshard [part1: remove_no_need_vars] Dec 14, 2023

AndSonder mentioned this pull request Dec 14, 2023

[AutoParallel] Optimize Reshard [part2: optimize _reshard_input] #60022

Merged

AndSonder changed the title ~~[AutoParallel] Optimize Reshard [part1: remove_no_need_vars]~~ [AutoParallel] Optimize Reshard [part1: optimize remove_no_need_vars] Dec 14, 2023

AndSonder deleted the opt_remove_no_need_var branch April 23, 2024 13:56

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

[AutoParallel] Optimize Reshard [part1: optimize remove_no_need_vars] #59912

[AutoParallel] Optimize Reshard [part1: optimize remove_no_need_vars] #59912

Uh oh!

AndSonder commented Dec 11, 2023 •

edited

Loading

Uh oh!

paddle-bot bot commented Dec 11, 2023

Uh oh!

From00 Dec 12, 2023

Uh oh!

AndSonder Dec 13, 2023

Uh oh!

From00 left a comment

Uh oh!

Uh oh!

[AutoParallel] Optimize Reshard [part1: optimize remove_no_need_vars] #59912

[AutoParallel] Optimize Reshard [part1: optimize remove_no_need_vars] #59912

Uh oh!

Conversation

AndSonder commented Dec 11, 2023 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

PR types

PR changes

Description

Uh oh!

paddle-bot bot commented Dec 11, 2023

Uh oh!

From00 Dec 12, 2023

Choose a reason for hiding this comment

Uh oh!

AndSonder Dec 13, 2023

Choose a reason for hiding this comment

Uh oh!

From00 left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

AndSonder commented Dec 11, 2023 •

edited

Loading