Fix bug in bfloat16 optimizer related to checkpointing (#4434)

* fix: bf16 optimizer if condition * fix: unexpected keyword argument 'load_serial' * fix: add load_serial arg to bf16_optimizer * style: fix indentation --------- Co-authored-by: Olatunji Ruwase <[email protected]>
microsoft · Oct 7, 2023 · ce6b3e0 · ce6b3e0
1 parent c4d4679
commit ce6b3e0
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 2 deletions.
diff --git a/deepspeed/runtime/bf16_optimizer.py b/deepspeed/runtime/bf16_optimizer.py
@@ -365,7 +365,8 @@ def load_state_dict(self,
                         state_dict_list,
                         checkpoint_folder,
                         load_optimizer_states=True,
-                        load_from_fp32_weights=False):
+                        load_from_fp32_weights=False,
+                        load_serial=None):
         if checkpoint_folder:
             self._load_universal_checkpoint(checkpoint_folder, load_optimizer_states, load_from_fp32_weights)
         else:

diff --git a/deepspeed/runtime/engine.py b/deepspeed/runtime/engine.py
@@ -3287,7 +3287,7 @@ def _get_zero_param_shapes(self):
         # if we don't use it, we get parameters ordered incorrectly
         if hasattr(self.optimizer, "round_robin_bit16_groups"):
             bit16_groups = self.optimizer.round_robin_bit16_groups
-        elif self.bfloat16_enabled() and not self.zero_optimization():
+        elif self.bfloat16_enabled() and hasattr(self.optimizer, "bf16_groups"):
             bit16_groups = self.optimizer.bf16_groups
         else:
             bit16_groups = self.optimizer.bit16_groups if self.zero_optimization_stage(