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