Unify the NNX and Linen restore paths in load_state_if_possible

The NNX path had two early-return blocks that each duplicated the Linen restore flow, one
for the standard manager and one for the emergency managers. Any change to restore logic
had to be made in three places, and they had already drifted: the NNX grain branch dropped
the iterator element that the Linen branch returns.

pure_nnx saves in the Linen on-disk layout, so both paths restore the same tree. Convert
the NNX abstract to that layout going in, run the existing match statement unchanged, and
reshape what comes back on the way out. Linen behavior is unchanged. The weights-only
branch's if/else folds into `_abstract_params`, which uses nnx.split_state rather than
nnx.split -- the State-native API, without the unused GraphDef.

That leaves the safetensors paths, which had the same duplication problem and three bugs
behind it:

The dynamic safetensors loader received the whole abstract state rather than the weights.
For an NNX state that is fatal before any weight is read: the optimizer's opt_state is
keyed by integer chain indices, and the loader flattens its target with a "." separator,
so the load died in flatten_dict with `TypeError: sequence item 2: expected str instance,
int found`. The loader now receives the weights, the way the weights-only Orbax load
already does, and its result is unwrapped back into the NNX params state.

A weight that no mapping covered came back as the unmaterialized leaf it went in as and
reached the model as an untrained init value. Orbax does not flag this -- a shape conflict
it rejects itself, but a weight the checkpoint simply does not carry it leaves as a
placeholder. The weights-only Orbax load already makes that check; the dynamic load now
makes it too, for both state types.

`checkpoint_conversion_fn` arrives from config as a dotted string but was called directly,
so it raised TypeError for any value that was set. It is now imported, and resolved before
the checkpoint is read so a bad config fails immediately.

Full-state safetensors also handles NNX now: the conversion fn produces MaxText's on-disk
layout, which is what pure_nnx reads, so it goes through the same reshape as every other
NNX restore.

Author: ecnal-cienet

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

@@ -15,6 +15,7 @@
 
 """Create an Orbax CheckpointManager with specified (Async or not) Checkpointer."""
 
+import importlib
 import time
 from typing import Any
 
@@ -158,33 +159,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 set(tree) == {"params"} 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(
@@ -227,6 +253,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)
@@ -238,7 +266,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}")
@@ -547,49 +579,12 @@ 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
+      # pure_nnx saves in the Linen on-disk layout, so both paths restore the same tree through
+      # the same manager calls below: convert the NNX abstract going in, reshape it coming out.
+      is_nnx = isinstance(abstract_unboxed_pre_state, nnx.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,
@@ -609,6 +604,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,
@@ -625,30 +622,44 @@ 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)
+    is_nnx = isinstance(params, nnx.State)
+    # 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,