fix

Author: epwalsh

src/olmo_core/distributed/checkpoint/__init__.py

@@ -204,7 +204,9 @@ def load_model_and_optim_state(
                 log.info(
                     f"Mapping current key '{current_key}' to key '{original_key}' in checkpoint"
                 )
-                state_dict[original_key] = state_dict.pop(current_key)
+                current_root, current_key = current_key.split(".", 1)
+                original_root, original_key = original_key.split(".", 1)
+                state_dict[original_root][original_key] = state_dict[current_root].pop(current_key)
 
     dist_cp.load(
         state_dict,