Fix bug in bfloat16 optimizer related to checkpointing

[okoge-kaz]

Overview

Fixed conditional branching where errors occurred.

-        elif self.bfloat16_enabled() and not self.zero_optimization():
+        elif self.bfloat16_enabled() and hasattr(self.optimizer, "bf16_groups"):

Error Detail

I encountered the following error while training with zero 1 optimization and bf16 using Megatron-DeepSpeed.

.env/lib/python3.11/site-packages/deepspeed/runtime/engine.py", line 3288, in _get_zero_param_shapes
    ) == 2 else self.optimizer.fp16_groups
                ^^^^^^^^^^^^^^^^^^^^^^^^^^
                ^^^^^^^^^^^^^^^^^^^^^^^^^^
    ) == 2 else self.optimizer.fp16_groups
    ) == 2 else self.optimizer.fp16_groups
AttributeError: 'BF16_Optimizer' object has no attribute 'fp16_groups'
    ) == 2 else self.optimizer.fp16_groups
AttributeError: 'BF16_Optimizer' object has no attribute 'fp16_groups'. Did you mean: 'bf16_groups'?
                ^^^^^^^^^^^^^^^^^^^^^^^^^^
                ^^^^^^^^^^^^^^^^^^^^^^^^^^
                ^^^^^^^^^^^^^^^^^^^^^^^^^^
AttributeError: 'BF16_Optimizer' object has no attr

and also when loading the checkpoint, I got another error.

.env/lib/python3.11/site-packages/deepspeed/runtime/engine.py", line 2873, in _load_zero_checkpoint
    self.optimizer.load_state_dict(state_dict_list=zero_sd_list,
TypeError: BF16_Optimizer.load_state_dict() got an unexpected keyword argument 'load_serial'

Related Pull Request

#3759

Related Issue

#4272

[okoge-kaz]

@tjruwase

Thanks for your review.
I have also confirmed that checkpoints can be saved & loaded with bf16 in Megatron-DeepSpeed!

(A100 40GB x 16(2node))

[tjruwase]

@okoge-kaz, thanks for sharing your results.

[tjruwase]

Fix #4272

[okoge-kaz]

@tjruwase All test passed!
May I ask for a merge?