Unify the NNX and Linen restore paths in load_state_if_possible

src/maxtext/checkpoint_conversion/utils/load_dynamic.py

@@ -251,9 +251,12 @@ def tensor_getter(key):
   return {"params": restored_params}
 
 
-def load_safetensors_dynamic_state(path, abstract_unboxed_pre_state, maxtext_config):
+def load_safetensors_dynamic_state(path, abstract_params, maxtext_config):
   """Main entry point to dynamically build and load safetensors into MaxText format.
 
+  `abstract_params` is the weights of the target state -- Linen's `params` collection, or the
+  NNX params state -- not the full train state; the HF param mappings name weights only.
+
   Splits execution into:
   1. Deriving Mappings
   2. Loading Sharded arrays directly to TPUs
@@ -349,11 +352,7 @@ def load_safetensors_dynamic_state(path, abstract_unboxed_pre_state, maxtext_con
   param_map_mt_to_hf, hook_fn_map_mt = get_hf_config_and_mappings(maxtext_config)
   max_logging.log(f"[1/3] Mappings derived in {time.time() - t_total:.2f}s")
 
-  target_tree = (
-      abstract_unboxed_pre_state.to_pure_dict()
-      if isinstance(abstract_unboxed_pre_state, nnx.State)
-      else abstract_unboxed_pre_state.params
-  )
+  target_tree = abstract_params.to_pure_dict() if isinstance(abstract_params, nnx.State) else abstract_params
 
   t1 = time.time()
   hf_state = load_sharded_hf_state(path)

src/maxtext/common/checkpointing.py

@@ -16,6 +16,7 @@
 """Create an Orbax CheckpointManager with specified (Async or not) Checkpointer."""
 
 import datetime
+import importlib
 import time
 from typing import Any
 
@@ -157,33 +158,58 @@ def _load_linen_checkpoint_into_nnx(
   restore_args = ocp.checkpoint_utils.construct_restore_args(linen_abstract)
   restored = ocp.args.PyTreeRestore(item=linen_abstract, restore_args=restore_args, partial_restore=True)
   restored = ckptr.restore(epath.Path(path), args=restored)
-  _raise_on_weight_mismatch(*_expected_and_restored_params(abstract_nnx_state, restored))
-  return _linen_items_to_nnx(restored, abstract_nnx_state)
+  return _restored_linen_to_nnx(restored, abstract_nnx_state)
 
 
-def _restore_emergency_linen_checkpoint_into_nnx(
-    checkpoint_manager,
-    step,
-    abstract_nnx_state,
-    map_to_pspec,
-):
-  """Restores an emergency Linen-layout checkpoint into an NNX state.
+def _restored_linen_to_nnx(restored_linen, abstract_nnx_state):
+  """Reshapes a restored Linen-layout tree into the NNX state.
 
-  The `nnx_aux` subtree is stored inside `items`, so an emergency checkpoint
-  carries it too; it's restored when present and otherwise kept at its fresh
-  init value. A genuinely-missing weight raises.
+  Raises if the checkpoint is missing a weight. Every NNX restore path ends here: the load
+  itself is the Linen one, since pure_nnx reads and writes the Linen on-disk layout.
   """
-  max_logging.log(f"Restoring emergency Linen-layout checkpoint into NNX state at step {step}")
-  linen_abstract = train_state_nnx.to_checkpoint_dict(abstract_nnx_state)
-  restore_args = jax.tree_util.tree_map(map_to_pspec, linen_abstract)
-  checkpoint_args = ocp.args.PyTreeRestore(
-      item=linen_abstract,
-      restore_args=restore_args,
-      partial_restore=True,
-  )
-  restored = checkpoint_manager.restore(step, args=Composite(state=checkpoint_args)).state
-  _raise_on_weight_mismatch(*_expected_and_restored_params(abstract_nnx_state, restored))
-  return _linen_items_to_nnx(restored, abstract_nnx_state)
+  _raise_on_weight_mismatch(*_expected_and_restored_params(abstract_nnx_state, restored_linen))
+  return _linen_items_to_nnx(restored_linen, abstract_nnx_state)
+
+
+def _abstract_params(abstract_unboxed_pre_state):
+  """Returns the state's weights: the NNX Param subtree, or Linen's `params` collection."""
+  if isinstance(abstract_unboxed_pre_state, nnx.State):
+    return nnx.split_state(abstract_unboxed_pre_state.model, nnx.Param, ...)[0]
+  return abstract_unboxed_pre_state.params
+
+
+def _bare_weights(tree):
+  """Strips the Flax `params` collection wrapper so weights compare at the same depth.
+
+  A Linen params tree is the collection, an NNX one the bare weights; the dynamic
+  safetensors loader always returns the collection.
+  """
+  return tree["params"] if isinstance(tree, dict) and len(tree) == 1 and "params" in tree else tree
+
+
+def _resolve_conversion_fn(checkpoint_conversion_fn):
+  """Returns `checkpoint_conversion_fn` as a callable.
+
+  Config carries it as a dotted string ("my_pkg.my_module.my_fn"), so it has to be imported
+  before it can be called. A callable is used as is.
+  """
+  if checkpoint_conversion_fn is None:
+    raise ValueError(
+        "source_checkpoint_layout='safetensors' needs `checkpoint_conversion_fn` to map the "
+        "checkpoint's weights onto the model's, e.g. checkpoint_conversion_fn=my_pkg.my_module.my_fn."
+    )
+  if callable(checkpoint_conversion_fn):
+    return checkpoint_conversion_fn
+  module_name, _, fn_name = str(checkpoint_conversion_fn).rpartition(".")
+  if not module_name:
+    raise ValueError(f"`checkpoint_conversion_fn` must be a dotted path to a function, got {checkpoint_conversion_fn!r}.")
+  try:
+    fn = getattr(importlib.import_module(module_name), fn_name, None)
+  except ImportError as e:
+    raise ValueError(f"Could not import `checkpoint_conversion_fn` {checkpoint_conversion_fn!r}: {e}") from e
+  if not callable(fn):
+    raise ValueError(f"`checkpoint_conversion_fn` {checkpoint_conversion_fn!r} is not a function.")
+  return fn
 
 
 def _load_full_state_from_path(
@@ -226,6 +252,8 @@ def _load_full_state_from_path(
       with context:
         return ocp_v1.load_pytree(path, abstract_unboxed_pre_state)
     elif source_checkpoint_layout == "safetensors":
+      # Resolved first, so a bad config fails before the weights are read.
+      conversion_fn = _resolve_conversion_fn(checkpoint_conversion_fn)
       context = ocp_v1.Context(checkpoint_layout=ocp_v1.options.CheckpointLayout.SAFETENSORS)
       with context:
         metadata = ocp_v1.pytree_metadata(path)
@@ -237,7 +265,11 @@ def combine_sharding(sds, shardings):
 
         sharded_abstract_state = jax.tree.map(combine_sharding, simple_abstract_state, shardings)
         pre_transformed_state = ocp_v1.load_pytree(path, sharded_abstract_state)
-      state = checkpoint_conversion_fn(pre_transformed_state)
+      state = conversion_fn(pre_transformed_state)
+      # The conversion fn returns MaxText's on-disk (Linen) layout, which is what pure_nnx reads,
+      # so NNX needs the same reshape as every other restore. An NNX state passes through.
+      if isinstance(abstract_unboxed_pre_state, nnx.State) and not isinstance(state, nnx.State):
+        state = _restored_linen_to_nnx(state, abstract_unboxed_pre_state)
       return state
     else:
       raise ocp_v1.errors.InvalidLayoutError(f"Unknown checkpoint layout: {source_checkpoint_layout}")
@@ -406,6 +438,11 @@ def load_state_if_possible(
      set.
   """
 
+  # pure_nnx saves in the Linen on-disk layout, so every branch below loads the same tree Linen
+  # does: the NNX abstract is converted to that layout going in, and what comes back is reshaped
+  # into the NNX state on the way out.
+  is_nnx = isinstance(abstract_unboxed_pre_state, nnx.State)
+
   if checkpoint_manager is not None:
     max_logging.log("checkpoint manager exists so trying to load this run's existing checkpoint")
 
@@ -437,49 +474,9 @@ def map_to_pspec(data):
         )
         ocp.type_handlers.register_type_handler(jax.Array, array_handler, override=True)
 
-      # pure_nnx saves in the Linen on-disk layout; restore that layout (weights +
-      # opt_state + step + nnx_aux), restoring the grain iterator in place when
-      # present, then reshape it back into the NNX state.
-      # (Emergency managers use their own restore path below.)
-      if isinstance(abstract_unboxed_pre_state, nnx.State) and not isinstance(
-          checkpoint_manager,
-          (EmergencyCheckpointManager, EmergencyReplicatorCheckpointManager),
-      ):
-        linen_abstract = train_state_nnx.to_checkpoint_dict(abstract_unboxed_pre_state)
-        restore_args = jax.tree_util.tree_map(map_to_pspec, linen_abstract)
-        checkpoint_args = ocp.args.PyTreeRestore(item=linen_abstract, restore_args=restore_args, partial_restore=True)
-        if (
-            dataset_type == "grain"
-            and data_iterator
-            and not isinstance(data_iterator, PlaceHolderDataIterator)
-            and (checkpoint_manager.directory / str(step) / "iter").exists()
-        ):
-          restored, _ = grain_utility.restore_grain_iterator(
-              checkpoint_manager, step, data_iterator, checkpoint_args, expansion_factor_real_data
-          )
-        else:
-          restored = checkpoint_manager.restore(step, args=Composite(items=checkpoint_args))
-        _raise_on_weight_mismatch(*_expected_and_restored_params(abstract_unboxed_pre_state, restored["items"]))
-        restored_nnx = _linen_items_to_nnx(restored["items"], abstract_unboxed_pre_state)
-        return ({"items": restored_nnx}, None)
-
-      if isinstance(abstract_unboxed_pre_state, nnx.State) and isinstance(
-          checkpoint_manager,
-          (EmergencyCheckpointManager, EmergencyReplicatorCheckpointManager),
-      ):
-        restored = _restore_emergency_linen_checkpoint_into_nnx(
-            checkpoint_manager,
-            step,
-            abstract_unboxed_pre_state,
-            map_to_pspec,
-        )
-        return (
-            restored,
-            None,
-        )
-
-      # Only Linen TrainState reaches here; the NNX cases returned above.
-      restore_target = abstract_unboxed_pre_state
+      restore_target = (
+          train_state_nnx.to_checkpoint_dict(abstract_unboxed_pre_state) if is_nnx else abstract_unboxed_pre_state
+      )
       restore_args = jax.tree_util.tree_map(map_to_pspec, restore_target)
       checkpoint_args = ocp.args.PyTreeRestore(
           item=restore_target,
@@ -499,6 +496,8 @@ def map_to_pspec(data):
             ),
         ):
           restored = checkpoint_manager.restore(step, args=Composite(state=checkpoint_args)).state
+          if is_nnx:
+            restored = _restored_linen_to_nnx(restored, abstract_unboxed_pre_state)
           return (
               restored,
               None,
@@ -515,30 +514,43 @@ def map_to_pspec(data):
             and not isinstance(data_iterator, PlaceHolderDataIterator)
             and (checkpoint_manager.directory / str(step) / "iter").exists()
         ):
-          return grain_utility.restore_grain_iterator(
+          restored, iterator = grain_utility.restore_grain_iterator(
               checkpoint_manager,
               step,
               data_iterator,
               checkpoint_args,
               expansion_factor_real_data,
           )
+          if is_nnx:
+            restored = {"items": _restored_linen_to_nnx(restored["items"], abstract_unboxed_pre_state)}
+          return (restored, iterator)
         # Case 3: Default/Fallback case.
         # This case acts as a wildcard ('_') and matches if none of the preceding cases were met.
         case _:
           restored = checkpoint_manager.restore(step, args=Composite(items=checkpoint_args))
+          if is_nnx:
+            restored = {"items": _restored_linen_to_nnx(restored["items"], abstract_unboxed_pre_state)}
           return (restored, None)
 
   if source_checkpoint_layout == "safetensors_dynamic":
     path = load_parameters_from_path or load_full_state_from_path
     max_logging.log(f"Dynamic On-the-Fly Formatting: Loading SafeTensors from {path}")
 
-    return load_safetensors_dynamic_state(path, abstract_unboxed_pre_state, maxtext_config)
+    # Weights-only for both paths, so the loader gets the weights rather than the whole state:
+    # the HF param mappings name weights, and an NNX state hides them under `model`.
+    params = _abstract_params(abstract_unboxed_pre_state)
+    restored, restored_params = load_safetensors_dynamic_state(path, params, maxtext_config)
+    # A weight no HF mapping covered comes back unmaterialized and would reach the model as an
+    # untrained init value. Same check the Orbax weights-only load makes.
+    _raise_on_weight_mismatch(_bare_weights(params.to_pure_dict() if is_nnx else params), _bare_weights(restored_params))
+    if is_nnx:
+      # The loader returns the Linen `params` collection; NNX holds bare weights, so unwrap it
+      # back into the params state, the shape load_params_from_path returns.
+      nnx.replace_by_pure_dict(params, restored_params["params"])
+      return restored, params
+    return restored, restored_params
   elif load_parameters_from_path != "":
-    if isinstance(abstract_unboxed_pre_state, nnx.State):
-      _, params, _ = nnx.split(abstract_unboxed_pre_state.model, nnx.Param, ...)
-    else:
-      params = abstract_unboxed_pre_state.params
-
+    params = _abstract_params(abstract_unboxed_pre_state)
     restored_params = load_params_from_path(
         load_parameters_from_path,
         params,