[REQUEST] High CPU memory usage at initialization

Is your feature request related to a problem? Please describe.

When I initiate a model on 4 GPUs with the DeepSpeed stage3, the CPU memory usage is very high. For instance, a model with 3B parameters takes 70GB of memory at its peak.
Because our server has a relatively small CPU memory size, there will be a lot of memory swapping when trying to initiate an even larger model. The initialization becomes unbearable slow and eventually exceeds the 30min timeout limit. RuntimeError: Timed out initializing process group in store based barrier on rank: 0, for key: store_based_barrier_key:1 (world_size=4, worker_count=1, timeout=0:30:00)

Describe the solution you'd like

I noticed there is a PR by @stas00 to address this issue. Do you plan to merge this PR soon?
https://github.com/bigscience-workshop/Megatron-DeepSpeed/pull/248

Describe alternatives you've considered
I could beg my advisor to buy larger memory chips, but I would prefer a software solution to the hardware solution.

The issue concerns DeepSpeed ZeRO stage 3 initialization and references PR #248; start by reviewing that pull request and the stage 3 initialization path. Done means reducing peak CPU memory during initialization enough to avoid swapping and the reported 30-minute process-group timeout.