[zero] refactor low level zero for shard evenly

[Gy-Lu]

📌 Checklist before creating the PR

• I have created an issue for this PR for traceability
• The title follows the standard format: [doc/gemini/tensor/...]: A concise description
• I have added relevant tags if possible for us to better distinguish different PRs

🚨 Issue number

Link this PR to your issue with words like fixed to automatically close the linked issue upon merge

e.g. fixed #1234, closed #1234, resolved #1234

#3954

📝 What does this PR do?

Summarize your work here.
if you have any plots/diagrams/screenshots/tables, please attach them here.

This PR has refactored low level zero for load balancing.

💥 Checklist before requesting a review

• I have linked my PR to an issue (instruction)
• My issue clearly describes the problem/feature/proposal, with diagrams/charts/table/code if possible
• I have performed a self-review of my code
• I have added thorough tests.
• I have added docstrings for all the functions/methods I implemented

⭐️ Do you enjoy contributing to Colossal-AI?

• 🌝 Yes, I do.
• 🌚 No, I don't.

Tell us more if you don't enjoy contributing to Colossal-AI.

[Gy-Lu]

The design is in #3954
And code polish is still on its way :<

[kurisusnowdeng]

The design is in #3954 And code polish is still on its way :<

This version is distributing each master parameter across devices instead of distributing the param list. Correct?

Just be aware of tensor precision during communication in terms of efficiency.

[Gy-Lu]

The design is in #3954 And code polish is still on its way :<

This version is distributing each master parameter across devices instead of distributing the param list. Correct?

Just be aware of tensor precision during communication in terms of efficiency.

Right.
Now the dtype of communication is the gradients' default dtype. It seems add an arg is better.

[github-actions]

The code coverage for the changed files is 83%.

Click me to view the complete report

Name                                                       Stmts   Miss  Cover
------------------------------------------------------------------------------
colossalai/zero/low_level/_utils.py                          125     66    47%
colossalai/zero/low_level/bookkeeping/bucket_store.py         51      0   100%
colossalai/zero/low_level/bookkeeping/gradient_store.py       31      1    97%
colossalai/zero/low_level/bookkeeping/parameter_store.py      16      0   100%
colossalai/zero/low_level/low_level_optim.py                 285     20    93%
tests/test_zero/test_low_level/test_grad_acc.py               90     28    69%
tests/test_zero/test_low_level/test_zero1_2.py               102      1    99%
------------------------------------------------------------------------------
TOTAL                                                        700    116    83%

[github-actions]

The code coverage for the changed files is 75%.

Click me to view the complete report

Name                                                           Stmts   Miss  Cover
----------------------------------------------------------------------------------
colossalai/zero/low_level/_utils.py                              125     69    45%
colossalai/zero/low_level/bookkeeping/bucket_store.py             51      4    92%
colossalai/zero/low_level/bookkeeping/gradient_store.py           31      2    94%
colossalai/zero/low_level/bookkeeping/parameter_store.py          16      1    94%
colossalai/zero/low_level/low_level_optim.py                     285     59    79%
tests/test_booster/test_plugin/test_low_level_zero_plugin.py      59      6    90%
----------------------------------------------------------------------------------
TOTAL                                                            567    141    75%

[github-actions]

The code coverage for the changed files is 84%.

Click me to view the complete report

Name                                                           Stmts   Miss  Cover
----------------------------------------------------------------------------------
colossalai/zero/low_level/_utils.py                              125     66    47%
colossalai/zero/low_level/bookkeeping/bucket_store.py             51      0   100%
colossalai/zero/low_level/bookkeeping/gradient_store.py           31      1    97%
colossalai/zero/low_level/bookkeeping/parameter_store.py          16      0   100%
colossalai/zero/low_level/low_level_optim.py                     285     21    93%
tests/test_booster/test_plugin/test_low_level_zero_plugin.py      59      6    90%
tests/test_zero/test_low_level/test_grad_acc.py                   90     28    69%
tests/test_zero/test_low_level/test_zero1_2.py                   102      1    99%
----------------------------------------------------------------------------------
TOTAL                                                            759    123    84%