set step count in opt state for each param

microsoft · Mar 13, 2024 · db79711 · db79711
1 parent 2943e6a
commit db79711
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/deepspeed/runtime/engine.py b/deepspeed/runtime/engine.py
@@ -2977,7 +2977,7 @@ def set_step(d):
         for group in optimizer.param_groups:
             set_step(group)
             for p in group['params']:
-                if p in state and len(state[p]) > 0 and 'step' in state[p]:
+                if p in state and len(state[p]) > 0:
                     set_step(state[p])
 
     def _get_mp_rank_zero_checkpoint_names(self, load_dir, tag, mp_rank, dp_world_size, bf16_mode):