Add the DistributedFusedLamb optimizer

[sneaxiy]

PR types

New features

PR changes

OPs

Describe

Add hybrid parallel DistributedFusedLamb optimizer.

• At the beginning, all workers have the whole parameters, the local gradients of the whole parameters, the partial moments.
• The local gradients of the whole parameter are reduce-scattered. Each worker have the partial reduce-scattered gradients.
• Each worker calculates the partial trust ratio div tensor using the partial reduce-scattered gradients.
• If the global norm clip is needed, each worker calculates the local gradient square L2-norm value using the partial reduce-scattered gradients and then calls ncclAllReduce to get the global gradient square L2-norm.
• Each worker calculates the square L2-norm value of the whole parameter.
• Each worker calculates the local square L2-norm value of the partial trust ratio div tensor, and then calls ncclAllReduce to get the global square L2-norm value.
• Each worker updates partial parameter, and call the ncclAllGather to get the whole updated parameter.

[paddle-bot-old]

Thanks for your contribution!
Please wait for the result of CI firstly. See Paddle CI Manual for details.

[paddle-bot-old]

Sorry to inform you that b9a7f57's CIs have passed for more than 7 days. To prevent PR conflicts, you need to re-run all CIs manually.

[limin2021]

LGTM.

[Aurelius84]

LGTM for dtype registerar

[zhiqiu]

LGTM

[zhiqiu]

Not important, but it seems params_grads is not used in _add_dynamic_loss_scaling in this case.

[sneaxiy]

Yes. But not quite important.

[XiaoguangHu01]

LGTM

[TCChenlong]

LGTM

[XieYunshen]

LGTM for set_tests_properties(test_distributed_fused_lamb_op_with_clip PROPERTIES TIMEOUT 120) set_tests_properties(test_distributed_fused_lamb_op_without_clip PROPERTIES TIMEOUT 120)