NNX in logits check, workaround

this is still broadcast among devices

update

update

Author: hengtaoguo

src/maxtext/utils/maxtext_utils.py

@@ -1087,6 +1087,58 @@ def setup_decode_state(model, config, rng, mesh, checkpoint_manager):
   return state, state_mesh_annotations
 
 
+def setup_decode_state_from_nnx(model, config, rng, mesh):
+  """Setup decode state by loading params from an NNX checkpoint.
+
+  Args:
+    model: the flax model to initialize
+    config: config object
+    rng: jax.prng key
+    mesh: jax.devices() mesh
+
+  Returns:
+    state: state with decode params loaded from the NNX checkpoint
+    state_mesh_annotations: the mesh annotations for the state
+  """
+  from maxtext.checkpoint_conversion.utils.utils import load_orbax_checkpoint  # pylint: disable=import-outside-toplevel
+
+  def unwrap_nnx_values(tree):
+    if isinstance(tree, dict):
+      filtered_tree = {k: v for k, v in tree.items() if 'to_nnx__rngs' not in k}
+      if len(filtered_tree) == 1 and 'value' in filtered_tree and isinstance(filtered_tree['value'], (jax.Array, np.ndarray)):
+        return filtered_tree['value']
+      return {k: unwrap_nnx_values(v) for k, v in filtered_tree.items()}
+    return tree
+
+  def to_numpy(tree):
+    def convert_leaf(x):
+      if isinstance(x, jax.Array):
+        return np.array(x)
+      return x
+    return jax.tree_util.tree_map(convert_leaf, tree)
+
+  def reshard_to_match(target_tree, source_tree):
+    def copy_sharding(target_leaf, source_leaf):
+      if isinstance(source_leaf, (jax.Array, np.ndarray)):
+        if isinstance(target_leaf, jax.Array) and hasattr(target_leaf, 'sharding'):
+          return jax.device_put(source_leaf, target_leaf.sharding)
+        return source_leaf
+      return source_leaf
+    return jax.tree_util.tree_map(copy_sharding, target_tree, source_tree)
+
+  unboxed_abstract_state, state_mesh_annotations, _ = get_abstract_state(model, None, config, rng, mesh, False)
+
+  loaded_params = load_orbax_checkpoint(config)
+  loaded_params = unwrap_nnx_values(loaded_params)
+  # NNX-RL checkpoints may have an extra 'base' nesting level, while NNX-SFT may not
+  loaded_params = {'params': loaded_params['base'] if 'base' in loaded_params else loaded_params}
+  loaded_params = to_numpy(loaded_params)
+  params = reshard_to_match(unboxed_abstract_state.params, loaded_params)
+
+  state = init_decode_state(model.apply, params)
+  return state, state_mesh_annotations
+
+
 def setup_training_state(model, data_iterator, tx, config, rng, mesh, checkpoint_manager):
   is_training = True
   return setup_initial_state(

src/maxtext/utils/model_creation_utils.py

@@ -183,6 +183,16 @@ def create_sharded_state():
         metadata = ckptr.metadata(config.load_parameters_path)
 
         is_nnx_checkpoint = True
+        has_base_key = False
+        
+        # Helper function to enable padding/truncation for mesh mismatches
+        def create_flexible_restore_args(tree):
+          return jax.tree.map(
+              lambda x: ocp.type_handlers.ArrayRestoreArgs(sharding=x.sharding) if hasattr(x, "sharding") else None,
+              tree,
+              is_leaf=lambda x: hasattr(x, "sharding"),
+          )
+        
         if (
             "params" in metadata.item_metadata.tree.keys()
             and "params" in metadata.item_metadata.tree.get("params", {}).keys()
@@ -196,7 +206,19 @@ def create_sharded_state():
           )
 
           item_to_restore = {"params": {"params": target_for_restore}}
-          restore_args = {"params": {"params": ocp.checkpoint_utils.construct_restore_args(target_for_restore)}}
+          # restore_args = {"params": {"params": ocp.checkpoint_utils.construct_restore_args(target_for_restore)}}
+          restore_args = {"params": {"params": create_flexible_restore_args(target_for_restore)}}
+        elif "base" in metadata.item_metadata.tree.keys():
+          # structure of nnx-rl checkpoint: {'base': {'decoder': {..., 'value': ...}}}
+          has_base_key = True
+          target_for_restore = jax.tree.map(
+              lambda v: {"value": v.value},
+              sharded_state,
+              is_leaf=lambda n: isinstance(n, nnx.Variable),
+          )
+          item_to_restore = {"base": target_for_restore}
+          # restore_args = {"base": ocp.checkpoint_utils.construct_restore_args(target_for_restore)}
+          restore_args = {"base": create_flexible_restore_args(target_for_restore)}
         else:
           # structure of nnx checkpoint: {'decoder': {'value': ...}}
           target_for_restore = jax.tree.map(
@@ -205,7 +227,8 @@ def create_sharded_state():
               is_leaf=lambda n: isinstance(n, nnx.Variable),
           )
           item_to_restore = target_for_restore
-          restore_args = ocp.checkpoint_utils.construct_restore_args(target_for_restore)
+          # restore_args = ocp.checkpoint_utils.construct_restore_args(target_for_restore)
+          restore_args = create_flexible_restore_args(target_for_restore)
 
         restored = ckptr.restore(
             epath.Path(config.load_parameters_path),
@@ -215,6 +238,10 @@ def create_sharded_state():
         )
 
         if is_nnx_checkpoint:
+          # Unwrap 'base' key if present (NNX-RL format)
+          if has_base_key:
+            restored = restored.get("base", restored)
+          
           checkpoint = jax.tree.map(
               lambda v: v["value"],
               restored,

tests/utils/forward_pass_logit_checker.py

@@ -425,7 +425,10 @@ def main(config, test_args):  # pylint: disable=W0621
     mesh = jax.sharding.Mesh(devices_array, config.mesh_axes)
     quant = quantizations.configure_quantization(config)
     maxtext_model = models.transformer_as_linen(config, mesh, quant=quant, model_mode=MODEL_MODE_TRAIN)
-    maxtext_state, _ = maxtext_utils.setup_decode_state(maxtext_model, config, rng1, mesh, None)
+    if test_args.ckpt_type == "linen":
+      maxtext_state, _ = maxtext_utils.setup_decode_state(maxtext_model, config, rng1, mesh, None)
+    else:
+      maxtext_state, _ = maxtext_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 = []
@@ -537,6 +540,34 @@ def main(config, test_args):  # pylint: disable=W0621
   test_args, remaining_args = parser.parse_known_args()
   # Reconstruct model_args (script name + the args MaxText needs)
   model_args = [sys.argv[0]] + remaining_args
+  parser.add_argument(
+      "--ckpt_type",
+      type=str,
+      required=False,
+      default="linen",
+      choices=["linen", "nnx"],
+      help="Checkpoint format to load: 'linen' (default) or 'nnx'.",
+  )
+  test_args, _ = parser.parse_known_args()
+
+  # Remove args defined in this test file to avoid error from pyconfig
+  model_args = sys.argv
+  to_remove_args = [
+      "--atol",
+      "--rtol",
+      "--token_size",
+      "--max_kl_div",
+      "--golden_logits_path",
+      "--hf_model_path",
+      "--run_hf_model",
+      "--output_logits_path",
+      "--gcs_output_logits_path",
+      "--clip_logits_epsilon",
+      "--skip_first_token",
+      "--ckpt_type",
+  ]
+  for arg in to_remove_args:
+    model_args = [s for s in model_args if not s.startswith(arg)]
 
   cfg = pyconfig.initialize(model_args)
   assert (