DeepSpeed

estimate_zero2_model_states_mem_needs: fixing memory estiamtion

#5099

Merged