DeepSpeed

Fix bug in bfloat16 optimizer related to checkpointing

#4434

Merged