NNX: fix checkpointing in the training loop

- Convert nnx.State to pure dict for checkpoint saving
- Restore pure dict back to nnx.State after loading

Author: ecnal-cienet

src/maxtext/common/checkpointing.py

@@ -590,8 +590,13 @@ def map_to_pspec(data):
         )
         ocp.type_handlers.register_type_handler(jax.Array, array_handler, override=True)
 
-      restore_args = jax.tree_util.tree_map(map_to_pspec, abstract_unboxed_pre_state)
-      checkpoint_args = ocp.args.PyTreeRestore(item=abstract_unboxed_pre_state, restore_args=restore_args)
+      # Convert nnx.State to pure dict to match how checkpoints are saved for NNX
+      restore_target = abstract_unboxed_pre_state
+      if isinstance(abstract_unboxed_pre_state, nnx.State):
+        restore_target = abstract_unboxed_pre_state.to_pure_dict()
+
+      restore_args = jax.tree_util.tree_map(map_to_pspec, restore_target)
+      checkpoint_args = ocp.args.PyTreeRestore(item=restore_target, restore_args=restore_args)
 
       match (checkpoint_manager, dataset_type, data_iterator):
         # Case 1: Matches if 'checkpoint_manager' is an instance of either EmergencyCheckpointManager
@@ -717,15 +722,35 @@ def save_params_to_path(checkpoint_dir, params, use_ocdbt=True, use_zarr3=True):
   print(f"Quantized params checkpoint saved at: {checkpoint_dir}")
 
 
-def maybe_save_checkpoint(checkpoint_manager, state, config, data_iterator, step=None):
-  """Save checkpoint if checkpointing is enabled."""
+def maybe_save_checkpoint(checkpoint_manager, state, config, data_iterator, step=None, force=False):
+  """Save checkpoint if checkpointing is enabled.
+
+  Args:
+    checkpoint_manager: The checkpoint manager.
+    state: The training state to save.
+    config: The config object.
+    data_iterator: The data iterator.
+    step: The step number. If None, extracts from state (for Linen TrainState).
+    force: If True, force save the checkpoint regardless of checkpoint_period.
+  """
   if checkpoint_manager is None:
     return
 
   # Determine the effective step for saving a checkpoint.
   # If 'step' is not provided, this call is for a potential final checkpoint
   # and use the last completed step from the state.
-  actual_step = (int(state.step) - 1) if step is None else int(step)
+  if step is not None:
+    actual_step = int(step)
+  else:
+    if config.pure_nnx:
+      actual_step = int(state.optimizer.step) - 1
+    else:
+      # Linen TrainState has .step attribute
+      actual_step = int(state.step) - 1
+
+  if config.pure_nnx:
+    # Convert nnx.State to dict.
+    state = state.to_pure_dict()
 
   # Determine if a checkpoint save should be forced, overriding the usual `config.checkpoint_period` logic.
   # This occurs if this function was called:

src/maxtext/utils/maxtext_utils.py

@@ -1168,7 +1168,10 @@ def setup_initial_state(
         # The update of data_iterator state happens in place, no need to assign explicitly
         state = restored["items"]
 
-      # TODO: For NNX, convert the pure dict to nnx.State.
+      # For NNX, convert the pure dict to nnx.State using the abstract state as template
+      if config.pure_nnx:
+        nnx.replace_by_pure_dict(unboxed_abstract_state, state)
+        state = unboxed_abstract_state
     else:
       init_state_partial = init_state_fn
       init_state_partial.__name__ = "initialize_state"

tests/integration/checkpointing_test.py

@@ -93,6 +93,7 @@ def get_checkpointing_command(run_date, hardware, steps, metrics_file, attention
           f"dataset_type={dataset_type}",
           "async_checkpointing=False",
           f"attention={attention_type}",
+          "profiler=''",
       ]
       + model_params
       + pathways_command
@@ -135,19 +136,19 @@ def run_checkpointing(hardware, attention_type):
   # Determine dataset path/pattern depending on decoupled mode.
   gcsfuse_pattern = "/tmp/gcsfuse/array-record/c4/en/3.0.1/c4-train.array_record*"
   local_decoupled_root = os.path.join(
-      MAXTEXT_PKG_DIR, "..", "tests", "assets", "local_datasets", "c4_en_dataset_minimal", "c4", "en", "3.0.1"
+      MAXTEXT_PKG_DIR, "..", "..", "tests", "assets", "local_datasets", "c4_en_dataset_minimal", "c4", "en", "3.0.1"
   )
   local_pattern = os.path.join(local_decoupled_root, "c4-train.array_record*")
   selected_pattern = gcsfuse_pattern
   dataset_path = "/tmp/gcsfuse"
 
-  if is_decoupled():
+  if not glob.glob(gcsfuse_pattern):
     # Prefer local minimal dataset if gcsfuse data absent
-    if not glob.glob(gcsfuse_pattern) and glob.glob(local_pattern):
+    if glob.glob(local_pattern):
       selected_pattern = local_pattern
-      dataset_path = os.path.join(MAXTEXT_PKG_DIR, "..", "tests", "assets", "local_datasets")
-    elif not glob.glob(gcsfuse_pattern) and not glob.glob(local_pattern):
-      pytest.skip("No grain ArrayRecord shards found for checkpointing test in decoupled mode.")
+      dataset_path = os.path.join(MAXTEXT_PKG_DIR, "..", "..", "tests", "assets", "local_datasets")
+    else:
+      pytest.skip("No grain ArrayRecord shards found for checkpointing test.")
 
   grain_command = [
       "grain_worker_count=0",

tests/integration/generate_param_only_checkpoint_test.py

@@ -54,6 +54,8 @@ def run_e2e_test_flow(hardware, model_config, attention_type="autoselected", sta
       f"attention={attention_type}",
       "max_target_length=128",
       "per_device_batch_size=1",
+      "profiler=''",
+      "pure_nnx=False",
   ] + model_config
 
   pathways_command = []
@@ -72,6 +74,7 @@ def run_e2e_test_flow(hardware, model_config, attention_type="autoselected", sta
             dataset_type="tfds",
             dataset_path=dataset_path,
         )
+        + ["pure_nnx=False"]
     )
     state_path = f"{base_output_directory}/runner_{run_date}/checkpoints/0/items"
 

tests/integration/standalone_dl_ckpt_test.py

@@ -89,6 +89,8 @@ def test_standalone_checkpointer(self):
             "async_checkpointing=False",
             "enable_goodput_recording=False",
             "skip_jax_distributed_system=True",
+            "pure_nnx=False",
+            "enable_nnx=False",
         )
     )
     # restore at 50 and checkpoint at 100
@@ -110,6 +112,8 @@ def test_standalone_checkpointer(self):
             "async_checkpointing=False",
             "enable_goodput_recording=False",
             "skip_jax_distributed_system=True",
+            "pure_nnx=False",
+            "enable_nnx=False",
         )
     )