Support AutoEP with ZeRO-3 zero.Init source modules (#8060)
This PR enables ZeRO3 support for AutoEP-managed MoE layers by
partitioning expert parameters over expert replica groups while router
and replicated parameters use the global data-parallel group.
With ZeRO3 enable, AutoEP preserves global data-parallel gradient
averaging for AutoEP expert parameters while reducing them over expert
replica groups. ZeRO parameters are gathered before AutoEP reads router
or expert tensors when replacing MoE modules created under
`deepspeed.zero.Init()`.
---------
Signed-off-by: Masahiro Tanaka <mtanaka@anyscale.com>
Co-authored-by: Ma, Guokai <guokai.ma@gmail.com>
