Merge pull request #3188 from AI-Hypercomputer:hengtaoguo-nnx-logits

PiperOrigin-RevId: 907034928

Author: Google-ML-Automation

src/maxtext/utils/model_creation_utils.py

@@ -31,7 +31,7 @@
 import numpy as np
 from jax.sharding import Mesh
 from maxtext.configs import pyconfig
-from maxtext.common.common_types import MODEL_MODE_TRAIN
+from maxtext.common.common_types import MODEL_MODE_AUTOREGRESSIVE, MODEL_MODE_TRAIN
 from maxtext.layers import quantizations
 from maxtext.models import models
 from maxtext.utils import max_logging
@@ -580,20 +580,21 @@ def _adjust_target_for_moe_fusion(target, meta_tree, is_nnx):
               }
           }
         else:
-          # structure of nnx checkpoint: {'decoder': {'value': ...}}
+          # NNX checkpoint: {'decoder': {'value': ...}}, or NNX-RL with extra 'base' nesting.
+          # Restore only nnx.Param — RNG variable shapes may differ between checkpoint and model.
           target_for_restore = jax.tree.map(
               lambda v: {"value": v.value},
               sharded_state,
               is_leaf=lambda n: isinstance(n, nnx.Variable),
           )
-          target_for_restore = _adjust_target_for_moe_fusion(target_for_restore, metadata.item_metadata.tree, True)
-          item_to_restore = target_for_restore
-          base_restore_args = ocp.checkpoint_utils.construct_restore_args(target_for_restore)
+          has_base_key = "base" in metadata.item_metadata.tree
+          meta_tree_for_params = metadata.item_metadata.tree.get("base", metadata.item_metadata.tree)
+          target_for_restore = _adjust_target_for_moe_fusion(target_for_restore, meta_tree_for_params, True)
+          item_to_restore = {"base": target_for_restore} if has_base_key else target_for_restore
           restore_args = _fix_restore_args_for_shape_mismatch(
-              base_restore_args,
-              metadata.item_metadata.tree,
-              mesh,
+              ocp.checkpoint_utils.construct_restore_args(target_for_restore), meta_tree_for_params, mesh
           )
+          restore_args = {"base": restore_args} if has_base_key else restore_args
 
         restored = ckptr.restore(
             epath.Path(config.load_parameters_path),
@@ -603,9 +604,10 @@ def _adjust_target_for_moe_fusion(target, meta_tree, is_nnx):
         )
 
         if is_nnx_checkpoint:
+          restored_root = restored["base"] if has_base_key else restored
           checkpoint = jax.tree.map(
               lambda v: v["value"],
-              restored,
+              restored_root,
               is_leaf=lambda x: isinstance(x, dict) and "value" in x and not isinstance(x.get("value"), dict),
           )
         else:
@@ -656,6 +658,13 @@ def _fuse_moe_weights(ckpt_tree, model_arrays_tree):
           # This prevents the replicated intermediate copies from persisting until function return.
           del restored
 
+          def _filter_to_model_keys(ckpt, model):
+            """Recursively keep only keys present in model, dropping checkpoint-only fields (e.g. to_nnx__rngs)."""
+            if not hasattr(ckpt, "items") or not hasattr(model, "items"):
+              return ckpt
+            return {k: _filter_to_model_keys(ckpt[k], model[k]) for k in model if k in ckpt}
+
+          checkpoint = _filter_to_model_keys(checkpoint, model_arrays)
           checkpoint = jax.tree.map(_expand_checkpoint_to_model_shapes, checkpoint, model_arrays)
           nnx.update(model, checkpoint)
 
@@ -672,3 +681,44 @@ def _fuse_moe_weights(ckpt_tree, model_arrays_tree):
       return model
     else:
       return model, mesh
+
+
+def setup_decode_state_from_nnx(model, config, rng, mesh):
+  """Setup decode state by loading an NNX or NNX-RL checkpoint into a linen TrainState.
+
+  Calls from_pretrained (which handles NNX and NNX-RL 'base'-nested checkpoints and
+  applies mesh sharding internally), then extracts nnx.Param values into a plain dict
+  for the linen TrainState. For linen checkpoints, use maxtext_utils.setup_decode_state instead.
+
+  Args:
+    model: the flax linen model to initialize
+    config: config object
+    rng: jax.prng key
+    mesh: jax.devices() mesh
+
+  Returns:
+    state: linen TrainState with params loaded from the NNX checkpoint
+    state_mesh_annotations: the mesh annotations for the state
+  """
+  init_state_fn = partial(maxtext_utils.init_initial_state, model, None, config, False, rng)
+  _, state_mesh_annotations, _ = maxtext_utils.get_abstract_state(config, mesh, init_state_fn, False)
+
+  # Load the NNX model; from_pretrained handles sharding via jax.jit(out_shardings=...).
+  nnx_model = from_pretrained(config, mesh=mesh, model_mode=MODEL_MODE_AUTOREGRESSIVE)
+
+  # Extract nnx.Param values, converting the State pytree to a plain nested dict.
+  def _state_to_dict(tree):
+    if isinstance(tree, nnx.Variable):
+      return tree.value
+    if hasattr(tree, "items") and not isinstance(tree, jax.Array):
+      return {k: _state_to_dict(v) for k, v in tree.items()}
+    return tree
+
+  nnx_param_state = nnx.state(nnx_model, nnx.Param)
+  raw_params = _state_to_dict(nnx_param_state)
+  del nnx_model, nnx_param_state  # free memory
+
+  params = {"params": raw_params}
+
+  state = maxtext_utils.init_decode_state(model.apply, params)
+  return state, state_mesh_annotations

tests/unit/model_creation_utils_test.py

@@ -26,8 +26,9 @@
 from jax.sharding import Mesh
 from orbax import checkpoint as ocp
 
+from flax.training import train_state
 from maxtext.configs import pyconfig
-from maxtext.common.common_types import MODEL_MODE_TRAIN, MODEL_MODE_PREFILL
+from maxtext.common.common_types import MODEL_MODE_AUTOREGRESSIVE, MODEL_MODE_TRAIN, MODEL_MODE_PREFILL
 from maxtext.models import models
 from maxtext.utils import maxtext_utils
 from maxtext.utils import model_creation_utils
@@ -393,5 +394,54 @@ def test_checkpoint_load_error_raises_value_error(self, mock_ocp):
       model_creation_utils.from_pretrained(cfg, self.mesh)
 
 
+class TestSetupDecodeStateFromNnx(unittest.TestCase):
+  """Tests for setup_decode_state_from_nnx()."""
+
+  def setUp(self):
+    self.config = _make_config()
+    self.mesh = _make_mesh(self.config)
+    self.rng = jax.random.PRNGKey(0)
+
+  def test_returns_linen_train_state_and_annotations(self):
+    """Should return a linen TrainState whose params mirror the NNX model's nnx.Param values."""
+    # Build a real (small) NNX model WITHOUT any patch active so from_pretrained
+    # runs normally and produces concrete jax.Array weights.
+    real_nnx_model = model_creation_utils.from_pretrained(self.config, mesh=self.mesh)
+
+    linen_model = model_creation_utils.from_config(self.config, mesh=self.mesh, rngs=None)
+
+    # Now patch from_pretrained so setup_decode_state_from_nnx never touches a checkpoint.
+    with patch("maxtext.utils.model_creation_utils.from_pretrained", return_value=real_nnx_model) as mock_fp:
+      state, state_mesh_annotations = model_creation_utils.setup_decode_state_from_nnx(
+          linen_model, self.config, self.rng, self.mesh
+      )
+
+    # from_pretrained must have been called with the right model_mode.
+    mock_fp.assert_called_once()
+    _, call_kwargs = mock_fp.call_args
+    self.assertEqual(call_kwargs.get("model_mode"), MODEL_MODE_AUTOREGRESSIVE)
+
+    # The result should be a linen TrainState.
+    self.assertIsInstance(state, train_state.TrainState)
+
+    # Params must be nested under "params" and be non-empty concrete arrays.
+    self.assertIn("params", state.params)
+    param_leaves = jax.tree.leaves(state.params["params"])
+    self.assertGreater(len(param_leaves), 0)
+    for leaf in param_leaves:
+      self.assertIsInstance(leaf, jax.Array)
+
+    # The NNX Param values and the extracted linen params must be numerically identical.
+    nnx_param_state = nnx.state(real_nnx_model, nnx.Param)
+    nnx_leaves = jax.tree.leaves(nnx_param_state)
+    linen_leaves = jax.tree.leaves(state.params["params"])
+    self.assertEqual(len(nnx_leaves), len(linen_leaves))
+    for nnx_val, linen_val in zip(nnx_leaves, linen_leaves):
+      self.assertTrue(jnp.all(nnx_val == linen_val))
+
+    # state_mesh_annotations must be returned (non-None).
+    self.assertIsNotNone(state_mesh_annotations)
+
+
 if __name__ == "__main__":
   unittest.main()

tests/utils/forward_pass_logit_checker.py

@@ -53,6 +53,7 @@
 from maxtext.models import models
 from maxtext.utils import max_logging
 from maxtext.utils import maxtext_utils
+from maxtext.utils import model_creation_utils
 import numpy as np
 import torch
 import torch.nn.functional as F
@@ -447,7 +448,10 @@ def main(config, test_args):  # pylint: disable=W0621
     else:
       maxtext_model = models.transformer_as_linen(config, mesh, quant=quant, model_mode=MODEL_MODE_TRAIN)
       init_state_fn = functools.partial(maxtext_utils.init_initial_state, maxtext_model, None, config, False, rng1)
-    maxtext_state, _ = maxtext_utils.setup_decode_state(config, mesh, None, init_state_fn)
+    if test_args.ckpt_type == "linen":
+      maxtext_state, _ = maxtext_utils.setup_decode_state(config, mesh, None, init_state_fn)
+    else:
+      maxtext_state, _ = model_creation_utils.setup_decode_state_from_nnx(maxtext_model, config, rng1, mesh)
 
     prompts = ["I love to", "Today is a", "What is the"]
     all_data_to_save = []
@@ -554,6 +558,14 @@ def main(config, test_args):  # pylint: disable=W0621
       default=False,
       help="Skip the first token during comparison to ignore BOS/init mismatches.",
   )
+  parser.add_argument(
+      "--ckpt_type",
+      type=str,
+      required=False,
+      default="linen",
+      choices=["linen", "nnx"],
+      help="Checkpoint format to load: 'linen' (default) or 'nnx'.",
+  )
 
   # Parse known args returns the namespace AND the list of remaining arguments
   test_args, remaining_args = parser.parse_known_args()