From db797112996e0912d3736386d9e9928c1affe663 Mon Sep 17 00:00:00 2001 From: Masahiro Tanaka Date: Wed, 13 Mar 2024 06:33:16 +0000 Subject: [PATCH] set step count in opt state for each param --- deepspeed/runtime/engine.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/deepspeed/runtime/engine.py b/deepspeed/runtime/engine.py index 67ed066f7486..174e699c5202 100644 --- 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):