Fix loading universal checkpoint for BF16_Optimizer

PR#5104 (Remove optimizer step on initialization) breaks loading universal checkpoint for BF16_Optimizer. This is since universal checkpoint attempts to load the optimizer states into lp._hp_mapping.optim_state dictionary before they are initialized (by step). As a workaround for loading universal checkpoint, perform step and init hp params optimizer's states before loading from universal checkpoint files. Signed-off-by: Moshe Island <[email protected]>
microsoft · Feb 29, 2024 · 9bc33ea · 9bc33ea
1 parent aed599b
commit 9bc33ea
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 1 deletion.
diff --git a/deepspeed/runtime/bf16_optimizer.py b/deepspeed/runtime/bf16_optimizer.py
@@ -445,6 +445,8 @@ def _load_legacy_checkpoint(self, state_dict_list, load_optimizer_states=True, l
             self._link_all_hp_params()
 
     def _load_universal_checkpoint(self, checkpoint_folder, load_optimizer_states, load_from_fp32_weights):
+        self.optimizer.step()
+        self._lazy_init_hp_params_optimizer_state()
         self._load_hp_checkpoint_state(checkpoint_folder)
 
     @property

diff --git a/deepspeed/runtime/engine.py b/deepspeed/runtime/engine.py
@@ -2785,7 +2785,7 @@ def load_checkpoint(self,
         if self.load_universal_checkpoint():
             self.optimizer.update_lp_params()
             if load_zero_checkpoint:
-                self.update_optimizer_step(step=client_states['iteration'] + 1)
+                self.update_optimizer_step(step=client_states['iteration'])
 
         return load_path, client_states