

TITLE
[DeepSpeed] ZeRO Stage 3

COMMENTS
2

REACTIONS
+1: 1
-1: 0
laugh: 0
hooray: 1
heart: 0
rocket: 0
eyes: 1

BODY
This PR implements DeepSpeed ZeR0 3 integration:

* [x] removes the "wind-down" of the deepspeed setup at the end of train, since zero3 can't do inference w/o this setup - we will have some other ways to reclaim memory for the no longer needed optimizer in the future.
* [x] adds initial support for eval w/o train - more work will be done in the future
* [x] to support `predict_with_generate` extends `generate` and its 5 beam search ways to support a new `synced_gpus` flag which is needed by ZeRO stage3 - under ZeRO3 parallelization if this gpu finished before max_length was reached, it must continue running forward so that the other gpus who may have not finished their generate yet, can complete as they rely on this gpu to received its slice of params. Currently deployed for DeepSpeed - but may need to do the same for fairscale elsewhere.
* [x] because now we are forced to run all gpus in sync, the `generate` logic is also now equipped with a stopping early mechanism that is synchronized across all participating gpus
* [x] reworks how pretrained model is loaded - `from_pretrained` is now zero3 aware and does a whole lot to efficiently preload massive models
* [x] since `state_dict` is fake under `zero3` it can't be saved and used - so care is taken to either not save the bogus model, or the weights get reconsolidated before saving if `stage3_gather_fp16_weights_on_model_save` is enabled
* [x] adds new DeepSpeed configuration docs and basic tuning recommendation
* [x] adds lots of new tests, now testing zero2 and zero3 separately
* [x] fixes a disappearing std stream problem in an older test using a workaround  
* [x] a new DS feature: `deepspeed.zero.register_external_parameter(self, self.layer1.weight)` - haven't needed it so far - need to find which models may need this feature. this is needed for when a layer accesses weights of another layer. but most our models don't do that. so just documented this for now.


DeepSpeed PRs that need to be merged and a new release is made 0.3.14:
* [x] https://github.com/microsoft/DeepSpeed/pull/881 (memory issue)
* [x] https://github.com/microsoft/DeepSpeed/pull/884 support lists of params in `GatheredParameters` (needed for pretrained model load) 
* [x] https://github.com/microsoft/DeepSpeed/pull/892 script to extract consolidated fp32 weights for zero2 and zero3
* [x] https://github.com/microsoft/DeepSpeed/pull/893 save the consolidated fp16 weights under zero3
* [x]  https://github.com/microsoft/DeepSpeed/pull/896 leak memory fix needed for tests
* [x] `deepspeed==0.3.14` released

May be:
* [x] https://github.com/microsoft/DeepSpeed/pull/882 - this needs to be changed to be optional - won't be efficient by default


Future PR TODO:

* [ ] make loading and resuming more efficient - gotta find a way to not preload the model from weights when we are resuming from a checkpoint. and of course not init the weights. Currently we do it 3 times! A huge overhead for big models.
