DeepSpeed

Fixes for training models with bf16 + freshly initialized optimizer via `load_module_only`

#4141

Merged