[low CPU RAM] allocate on gpu directly, stagger checkpoint load/save

[stas00]

As we are dealing with CPU RAM < GPU RAM this PR is trying

1. to init the model on GPU by allocating the model on gpu directly
   Code is courtesy of @jeffra with a demo at https://gist.github.com/jeffra/ec8a0b762e58a64d19ad1417250dc600

The idea was to remove deepspeed.zero.Init as supposedly it wasn't really doing anything under zero1, but when I did that the tp test started hanging. So I put it back and double inserted AllocateOnGPU in a nested fashion. deepspeed.zero.Init actually does a whole bunch of things regardless of ZeRO stage.

2. stagger save/load_checkpoints in groups defined by --stagger_checkpoint_save_load_group_size - if not set normal all processes at once is performed.

This works on load_checkpoint, but hangs on save_checkpoint - I think it deadlocks since save_checkpoint probably calls a barrier too.

---

3. while at it sped up several tests by making the test model smaller and training for 50x shorter.

[stas00]

Is there any time where we don't want to allocate directly on GPU for pretrain?

[stas00]

For posterity: this works most of the time, except when it doesn't. In the complex Megatron-Deepspeed framework using this PR lead to all kinds of very weird problems often on CUDA level.

We eventually found that the problem wasn't the shortage of CPU memory but some issue in apex's FusedAdam which is described here and a workaround has been merged: #249
The merged workaround is a hack and needs more investigation.