[zero] add state dict for low level zero

[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

closed #4154

📝 What does this PR do?

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

Provide interface of low level zero's state_dict and load_state_dict.

💥 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.

[github-actions]

The code coverage for the changed files is 87%.

Click me to view the complete report

Name                                                           Stmts   Miss  Cover
----------------------------------------------------------------------------------
colossalai/booster/booster.py                                     64     10    84%
colossalai/booster/plugin/gemini_plugin.py                       110      9    92%
colossalai/booster/plugin/low_level_zero_plugin.py                96      9    91%
colossalai/booster/plugin/plugin_base.py                          40     10    75%
colossalai/booster/plugin/torch_ddp_plugin.py                     66      2    97%
colossalai/booster/plugin/torch_fsdp_plugin.py                    97     13    87%
colossalai/zero/low_level/_utils.py                              119     66    45%
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                     279     19    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                   92     31    66%
tests/test_zero/test_low_level/test_zero1_2.py                   102      1    99%
tests/test_zero/test_low_level/test_zero_ckpt.py                  69      5    93%
tests/test_zero/test_low_level/test_zero_init.py                  40      4    90%
tests/test_zero/test_low_level/test_zero_tp.py                    66      1    98%
----------------------------------------------------------------------------------
TOTAL                                                           1397    187    87%

[Gy-Lu]

The current implementation of load_state_dict would change the passing state_dict.
Not sure if it would introduce some bugs under some situation.

[github-actions]

The code coverage for the changed files is 78%.

Click me to view the complete report

Name                                                           Stmts   Miss  Cover
----------------------------------------------------------------------------------
colossalai/booster/booster.py                                     64     17    73%
colossalai/booster/plugin/gemini_plugin.py                       110     60    45%
colossalai/booster/plugin/low_level_zero_plugin.py                96      9    91%
colossalai/booster/plugin/plugin_base.py                          40     10    75%
colossalai/booster/plugin/torch_ddp_plugin.py                     66     26    61%
colossalai/booster/plugin/torch_fsdp_plugin.py                    97     48    51%
colossalai/zero/low_level/_utils.py                              119     66    45%
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                     283     19    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                   92     31    66%
tests/test_zero/test_low_level/test_zero1_2.py                   102      1    99%
tests/test_zero/test_low_level/test_zero_ckpt.py                  69      5    93%
tests/test_zero/test_low_level/test_zero_init.py                  40      4    90%
tests/test_zero/test_low_level/test_zero_tp.py                    66      1    98%
----------------------------------------------------------------------------------
TOTAL                                                           1401    304    78%