[Auto Parallel] Add tensor_fusion and overlap in auto dy sharding

[Xing-lil]

PR Category

Auto Parallel

PR Types

New features

Description

1. inplace_master_grad

• In-place param.main_grad replaces the old master_grad in auto dy.
• param.main_grad will use inplace add_ to save or cast grad to fp32 and store them in param.main_grad.
• Enable by setting export Flags_enable_inplace_master_grad=1.

2. tensor_fusion

• tensor_fusion groups params and grads into continuous param_storage and grad_storage.
• grad_storage is used for grad's reduce_scatter comm.
• param_storage is used for param's all_gather comm.
• Supports non-uniform partitioning of params and grads across GPUs.
• Each step requires get non-uniform params and grads from param_storage and grad_storage using view_slice.
• Non-uniform grad_chip requires call all_reduce manually to collect global_norm_var.
• Enable by setting export FLAGS_enable_tensor_fusion=1.

3. sharding_overlap

• Overlap reduce_scatter comm for grads with grad computation in bwd.
• Overlap all_gather comm for params with opt computation.
• Enable by setting export FLAGS_enable_tensor_fusion=1.

Note: non-uniform tensor_fusion changes the order of add in grad_chip, introducing some loss diff.
Convergence results on llama7b, 1NC8, sharding8, 50,000 steps.

【TODO】Add strategy config in auto-dy, like hand-dy (feelt.init(strategy)) and auto-static (to_static(strategy)).

Pcard-70448

[paddle-bot]

你的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.

[liym27]

sharding_overlap will not be set through FLAGS in the future and better to configure it with Strategy

[Xing-lil]

Thx, will fix it later

[liym27]

LGTM