[Dist][Pass] ZeRO optimization

[zhen-jia]

Description

This PR performs all ZeRO related optimizations. There are two main optimizations and two passes are impacted. May changes are listed below:

1. Modify the partition gradient pass: after modification, the reduce_scatters are grouped and issued together for ZeRO-2. ZeRO-1 is remain the same
2. Implement a new pass, i.e., group_allgather. This pass detect the computation pattern in ZeRO and group the corresponding ops. Including cast and allgather. The group_allgather updates the parameter in-placed. So no extra add zero will be launched.
3. A new field is added to distributed context: group_bucket_size. This is an hyper parameter to control the group size in group communications. The default value is 5000000000, which is in consistent with DeepSpeed. It can be set to 0 or 1 so as to not group any communications at all.

Checklist

• PR's title starts with a category (e.g. [BUGFIX], [MODEL], [TUTORIAL], [FEATURE], [DOC], etc)
• Changes are complete (i.e. I finished coding on this PR)
• All changes have test coverage
• Code is well-documented

cc @awslabs/raf-reviewer

[comaniac]

Overall LGTM. One insight I got from this PR is that we really need an ANF pattern matcher...

Other than the comments I left, there are two points

1. Please be ware of [Dist] Refactor dist context & communicator and add inter-process communication with file #20 that changes dist_context and will result in conflicts (please also help review that PR and get it merged).
2. @hzfan @zachzzc @whbldhwj @hgt312 @yzhliu please help review this PR as well.

[zhen-jia]

Got it, Thanks! Will modify accordingly.

Overall LGTM. One insight I got from this PR is that we really need an ANF pattern matcher...

Other than the comments I left, there are two points

1. Please be ware of [Dist] Refactor dist context & communicator and add inter-process communication with file #20 that changes dist_context and will result in conflicts (please also help review that PR and get it merged).
2. @hzfan @zachzzc @whbldhwj @hgt312 @yzhliu please help review this PR as well.

[comaniac]

LGTM

[comaniac]

Thanks @zhen-jia