Keep Pete's fixes

allenai · Dec 18, 2023 · 8e7b0fd · 8e7b0fd
1 parent 52bcd3f
commit 8e7b0fd
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/olmo/checkpoint.py b/olmo/checkpoint.py
@@ -668,6 +668,10 @@ def restore_checkpoint(
                 optim_state_dict_to_load = load_state_dict(
                     load_path, "optim.pt", local_cache=local_cache, map_location="cpu"
                 )
+                optim_state_dict_to_load = self._make_optim_state_dict_compatible(
+                    optim_state_dict_to_load,
+                    og_keys_to_new,
+                )
                 load_fsdp_optim_state(fsdp_model, optim, optim_state_dict_to_load)
                 del optim_state_dict_to_load