[REQUEST] parameter sharding, gradient sharding, and optimizer state sharding with various sharding factors like Zero++

Is your feature request related to a problem? Please describe.
Zero++ only reconstruct parameter sharding. However, when there are a large number of Gpus, it is also necessary to aggregate gradient sharding and optimizer state sharding according to different communication bottlenecks. When there are thousands of Gpus, splitting the optimizer state into N parts is not necessary and only increases the communication overhead.We can deal with gradient sharding and optimizer state sharding like Zero++ deals with parameter sharding.

As shown in the figure, you can refer to this paper for details https://arxiv.org/html/2401.09149v1

As a result, this new method of sharding may reach a very high MFU.

Describe the solution you'd like
Base on Zero++, make gradient sharding and optimizer state sharding with various sharding factors N. For example, optimizer state sharding is N, gradient sharding maybe N/8.

Additional context
Here is the code implementation, it could be easily migrated into deepspeed.
https://github.com/InternLM/InternEvo

Start with the linked InternEvo implementation and the referenced paper, then compare its approach with DeepSpeed's existing Zero++ behavior. The requested result is configurable sharding factors for gradient and optimizer state sharding, such as N and N/8, with the communication and MFU goals described in the issue.