Allow unexpected keys in restore_item during partial restoration.

PiperOrigin-RevId: 823229099

Author: BlaziusMaximus

checkpoint/orbax/checkpoint/_src/checkpointers/checkpointer_test_utils.py

@@ -767,30 +767,6 @@ def test_partial_restore_with_omission(self):
           )
           test_utils.assert_tree_equal(self, expected, restored)
 
-      with self.subTest('extra_leaf'):
-        with self.checkpointer(
-            PyTreeCheckpointHandler()
-        ) as restore_checkpointer:
-          reference_item = {
-              'a': 0,
-              'c': {
-                  'a': 0,
-              },
-              'z': 0,
-          }
-          with self.assertRaisesRegex(
-              ValueError,
-              'Missing keys were found in the user-provided restore item.',
-          ):
-            restore_checkpointer.restore(
-                directory,
-                args=pytree_checkpoint_handler.PyTreeRestoreArgs(
-                    item=reference_item,
-                    restore_args=self.pytree_restore_args,
-                    partial_restore=True,
-                ),
-            )
-
     def test_restore_logs_read_event(self):
       """Tests that restore logs a read event to DM Sawmill log."""
       with self.checkpointer(PyTreeCheckpointHandler()) as checkpointer:

checkpoint/orbax/checkpoint/_src/handlers/base_pytree_checkpoint_handler.py

@@ -833,14 +833,10 @@ def _partial_restore_with_omission(
           is_leaf=tree_utils.is_empty_or_leaf,
       )
 
-    try:
-      value_metadata_tree = tree_structure_utils.tree_trim(
-          serialized_item, value_metadata_tree, strict=True
-      )
-    except ValueError as e:
-      raise ValueError(
-          'Missing keys were found in the user-provided restore item.'
-      ) from e
+    value_metadata_tree = tree_structure_utils.tree_trim(
+        serialized_item, value_metadata_tree, strict=False
+    )
+    value_metadata_tree = value_metadata_tree.unsafe_structure
 
     if restore_args is not None:
       restore_args = tree_structure_utils.tree_trim(
@@ -850,26 +846,37 @@ def _partial_restore_with_omission(
     return value_metadata_tree, restore_args
 
   def _partial_restore_with_placeholders(
-      self,
-      item: PyTree,
-      value_metadata_tree: PyTree,
-  ):
+      self, item: PyTree, value_metadata_tree: PyTree
+  ) -> PyTree:
     """Restores leaves from `item`, except for those marked as placeholders."""
     serialized_item = tree_utils.serialize_tree(item, keep_empty_nodes=True)
-    diff = tree_structure_utils.tree_difference(
-        serialized_item,
-        value_metadata_tree,
-        is_leaf=tree_utils.is_empty_or_leaf,
-        leaves_equal=lambda a, b: True,
+    diff = (
+        tree_structure_utils.tree_difference(
+            serialized_item,
+            value_metadata_tree,
+            is_leaf=tree_utils.is_empty_or_leaf,
+            leaves_equal=lambda a, b: True,
+        )
+        or {}
     )
-    if diff is not None:
-      formatted_diff = tree_structure_utils.format_tree_diff(
-          diff, source_label='Item', target_label='Metadata'
-      )
-      raise ValueError(
-          'User-provided restore item and on-disk value metadata tree'
-          f' structures do not match:\n{formatted_diff}'
-      )
+    for keypath, value_diff in tree_utils.to_flat_dict(
+        diff, is_leaf=lambda x: isinstance(x, tree_structure_utils.Diff)
+    ).items():
+      if value_diff.lhs is PLACEHOLDER and value_diff.rhs is None:
+        parent = value_metadata_tree
+        for key in keypath[:-1]:
+          parent = parent[key]
+        parent[keypath[-1]] = PLACEHOLDER
+      else:
+        formatted_diff = tree_structure_utils.format_tree_diff(
+            diff, source_label='Item', target_label='Metadata'
+        )
+        raise ValueError(
+            'User-provided restore item and on-disk value metadata tree'
+            f' structures do not match:\n{formatted_diff}\nIf this mismatch is'
+            ' intentional, pass `partial_restore=True` to only restore'
+            ' parameters found in `item`.'
+        )
     return jax.tree.map(
         lambda v, i: PLACEHOLDER if type_handlers.is_placeholder(i) else v,
         value_metadata_tree,
@@ -1021,7 +1028,9 @@ class TrainState:
         )
         raise ValueError(
             'User-provided restore item and on-disk value metadata tree'
-            f' structures do not match:\n{formatted_diff}'
+            f' structures do not match:\n{formatted_diff}\nIf this mismatch is'
+            ' intentional, pass `partial_restore=True` to only restore'
+            ' parameters found in `item`.'
         )
     restore_args = _fill_missing_save_or_restore_args(
         item, restore_args, mode='restore'
@@ -1043,6 +1052,14 @@ class TrainState:
         )
     )
 
+    if args.partial_restore:
+      restored_item = jax.tree.map(
+          lambda r, i: i if r is type_handlers.PLACEHOLDER else r,
+          restored_item,
+          item,
+          is_leaf=tree_utils.is_empty_or_leaf,
+      )
+
     if logging.vlog_is_on(1):
       logging.vlog(1, 'param_infos: %s', param_infos)
       logging.vlog(1, 'checkpoint_restore_args: %s', restore_args)

checkpoint/orbax/checkpoint/checkpoint_manager_test.py

@@ -3907,30 +3907,6 @@ def test_partial_restore_with_omission(self):
         )
         test_utils.assert_tree_equal(self, expected, restored)
 
-    with self.subTest('extra_leaf'):
-      with CheckpointManager(directory) as restore_manager:
-        reference_item = {
-            'a': 0,
-            # Omit 'b'
-            'c': {
-                'a': 0,
-                # Omit 'e'
-            },
-            'z': 0,
-        }
-        with self.assertRaisesRegex(
-            ValueError,
-            'Missing keys were found in the user-provided restore item.',
-        ):
-          restore_manager.restore(
-              0,
-              args=args.PyTreeRestore(
-                  reference_item,
-                  restore_args=self.pytree_restore_args,
-                  partial_restore=True,
-              ),
-          )
-
 
 if __name__ == '__main__':
   multiprocess_test.main()