fix(nnx): fix gpt3 paxml checkpoint conversion state map and layout

src/maxtext/checkpoint_conversion/standalone_scripts/convert_gpt3_ckpt_from_paxml.py

@@ -41,6 +41,7 @@
 import sys
 
 from flax import nnx
+from flax import serialization
 import jax
 from jax import random
 from jax.sharding import Mesh
@@ -118,6 +119,12 @@ def init_state_fn():
   )
 
   state, _, _, _, _ = maxtext_utils.setup_training_state(None, cfg, mesh, checkpoint_manager, init_state_fn)
+  if cfg.pure_nnx:
+    state = train_state_nnx.to_checkpoint_dict(state)
+    state.pop("nnx_aux", None)
+  else:
+    state = serialization.to_state_dict(state)
+
   max_logging.log("start")
   max_utils.print_mem_stats("After params initialized")
 
@@ -201,24 +208,10 @@ def init_state_fn():
       "['decoder']['decoder_norm']['bias']": (".params.lm.final_ln.bias", None),
   }
 
-  if cfg.pure_nnx:
-    # NNX state-tree paths after `nnx.split(TrainStateNNX)`. The state is a
-    # nested `nnx.State` (dict-like Mapping) with `nnx.Variable` leaves, so
-    # `jax.tree_util.keystr` produces dict-style entries (`['key']`) plus
-    # `.value` for the Variable leaf, plus `[idx]` for the optax tuple:
-    #   model params   -> ['model']<rest>.value
-    #   adam mu / nu   -> ['optimizer']['opt_state'][0]['mu' | 'nu']<rest>.value
-    #   step           -> ['optimizer']['step'].value
-    #   opt count      -> ['optimizer']['opt_state'][0]['count'].value
-    state_map = {
-        "['optimizer']['step'].value": ("step", None),
-        "['optimizer']['opt_state'][0]['count'].value": ("opt_states_0.no_prefix_0.count", None),
-    }
-  else:
-    state_map = {
-        ".step": ("step", None),
-        ".opt_state.count": ("opt_states_0.no_prefix_0.count", None),
-    }
+  state_map = {
+      "['step']": ("step", None),
+      "['opt_state']['count']": ("opt_states_0.no_prefix_0.count", None),
+  }
 
   def get_layer_prefix(keystr_pax):
     # different path format between decoder_layer variable
@@ -231,26 +224,18 @@ def get_layer_prefix(keystr_pax):
 
   for keystr_maxtext, (keystr_pax, transform_fn) in keystr_map.items():
     prefix_pax_opt_state = get_layer_prefix(keystr_pax)
-    if cfg.pure_nnx:
-      state_map[f"['model']{keystr_maxtext}.value"] = (f"mdl_vars{keystr_pax}", transform_fn)
-      state_map[f"['optimizer']['opt_state'][0]['mu']{keystr_maxtext}.value"] = (
-          f"opt_states_0.{prefix_pax_opt_state}.m{keystr_pax}",
-          transform_fn,
-      )
-      state_map[f"['optimizer']['opt_state'][0]['nu']{keystr_maxtext}.value"] = (
-          f"opt_states_0.{prefix_pax_opt_state}.v{keystr_pax}",
-          transform_fn,
-      )
-    else:
-      state_map[f".params['params']{keystr_maxtext}"] = (f"mdl_vars{keystr_pax}", transform_fn)
-      state_map[f".opt_state.mu['params']{keystr_maxtext}"] = (
-          f"opt_states_0.{prefix_pax_opt_state}.m{keystr_pax}",
-          transform_fn,
-      )
-      state_map[f".opt_state.nu['params']{keystr_maxtext}"] = (
-          f"opt_states_0.{prefix_pax_opt_state}.v{keystr_pax}",
-          transform_fn,
-      )
+    state_map[f"['params']['params']{keystr_maxtext}"] = (
+        f"mdl_vars{keystr_pax}",
+        transform_fn,
+    )
+    state_map[f"['opt_state']['mu']['params']{keystr_maxtext}"] = (
+        f"opt_states_0.{prefix_pax_opt_state}.m{keystr_pax}",
+        transform_fn,
+    )
+    state_map[f"['opt_state']['nu']['params']{keystr_maxtext}"] = (
+        f"opt_states_0.{prefix_pax_opt_state}.v{keystr_pax}",
+        transform_fn,
+    )
 
   def verify_fn(key_path, _):
     keystr = jax.tree_util.keystr(key_path)
@@ -302,9 +287,9 @@ def map_fn(key_path, value):
   max_logging.log("converted state finished")
   max_utils.print_mem_stats("converted state finished")
 
-  step_value = int(converted_state.optimizer.step.value) if cfg.pure_nnx else converted_state.step
-  if checkpointing.save_checkpoint(checkpoint_manager, step_value, converted_state):
-    max_logging.log(f"saved a checkpoint at step {step_value}")
+  step_value = int(converted_state["step"]) if isinstance(converted_state, dict) else int(converted_state.step)
+  if checkpointing.save_checkpoint(checkpoint_manager, step_value, converted_state, config=cfg):
+    max_logging.log(f"Saved a checkpoint at step {step_value}.")
   # Upon preemption, exit when and only when all ongoing saves are complete.
   if checkpointing.reached_preemption(checkpoint_manager, step_value):
     checkpointing.wait_until_finished(checkpoint_manager)
@@ -322,12 +307,22 @@ def map_fn(key_path, value):
       default="gs://mlperf-llm-public2/gpt3_spmd1x64x24_tpuv4-3072_v84_20221101/checkpoints/checkpoint_00004000",
       required=True,
   )
-  parser.add_argument("--maxtext-model-name", choices=["gpt3-175b", "gpt3-52k"], type=str, required=True)
+  parser.add_argument(
+      "--maxtext-model-name",
+      choices=["gpt3-175b", "gpt3-52k"],
+      type=str,
+      required=True,
+  )
   parser.add_argument("--base-output-directory", type=str, required=True)
   parser.add_argument("--run-name", type=str, required=True)
 
   args = parser.parse_args()
   if not args.paxml_ckpt_path.startswith("gs://"):
     raise ValueError("--paxml-ckpt-path should be a gcs path starting with gs://")
 
-  convert(args.paxml_ckpt_path, args.maxtext_model_name, args.base_output_directory, args.run_name)
+  convert(
+      args.paxml_ckpt_path,
+      args.maxtext_model_name,
+      args.base_output_directory,
+      args.run_name,
+  )

src/maxtext/common/checkpointing.py

@@ -753,18 +753,6 @@ def maybe_save_checkpoint(checkpoint_manager, state, config, data_iterator, step
     max_logging.log(f"Checkpoint for step {actual_step} already exists, skipping save.")
     return
 
-  if config.pure_nnx:
-    # Save in the Linen on-disk layout so pure_nnx and Linen checkpoints are interchangeable.
-    if config.enable_diloco:
-      # DiLoCoTrainState: persist the synchronized global model (outer params).
-      # The per-replica inner optimizer / outer-momentum state is not checkpointed.
-      step_value = state.step.get_value() if hasattr(state.step, "get_value") else state.step
-      state = train_state_nnx.to_linen_checkpoint_dict({"model": state.params, "optimizer": {"step": step_value}})
-    else:
-      # rngs/dropout/batch-stats are packed under items/nnx_aux so the RNG/dropout
-      # stream continues across resumes instead of resetting to a base key.
-      state = train_state_nnx.to_checkpoint_dict(state)
-
   try:
     checkpoint_saved = save_checkpoint(checkpoint_manager, actual_step, state, config, data_iterator, force_ckpt_save)
     if checkpoint_saved:
@@ -793,12 +781,26 @@ def maybe_save_checkpoint(checkpoint_manager, state, config, data_iterator, step
 
 def save_checkpoint(checkpoint_manager, step, state, config=None, data_iterator=None, force=False):
   """Wrapper for saving checkpoint."""
-  if config and config.enable_checkpointing:
+  if not isinstance(state, (dict, nnx.State, train_state.TrainState)):
+    if isinstance(state, train_state_nnx.TrainStateNNX):
+      state = nnx.state(state)
+    elif not isinstance(state, (dict, nnx.State)):
+      state = {}
+
+  if config and getattr(config, "pure_nnx", False) and isinstance(state, nnx.State):
+    # Save in the Linen on-disk layout so pure_nnx and Linen checkpoints are interchangeable.
+    if getattr(config, "enable_diloco", False):
+      step_value = state.step.get_value() if hasattr(state.step, "get_value") else state.step
+      state = train_state_nnx.to_linen_checkpoint_dict({"model": state.params, "optimizer": {"step": step_value}})
+    else:
+      state = train_state_nnx.to_checkpoint_dict(state)
+
+  if config and getattr(config, "enable_checkpointing", False):
     if (
         force
-        or (step % config.checkpoint_period == 0 and not config.enable_continuous_checkpointing)
+        or (step % config.checkpoint_period == 0 and not getattr(config, "enable_continuous_checkpointing", False))
         or (_uses_local_checkpoint_period(config) and step % config.local_checkpoint_period == 0)
-        or (config.enable_autocheckpoint and reached_preemption(checkpoint_manager, step))
+        or (getattr(config, "enable_autocheckpoint", False) and reached_preemption(checkpoint_manager, step))
     ):
       blocking_until_ready_start = time.time()
       max_logging.log(f"Waiting for step {step} to finish before checkpoint...")
@@ -812,7 +814,13 @@ def save_checkpoint(checkpoint_manager, step, state, config=None, data_iterator=
 
   # specify chunk_byte_size to force orbax to control maximum file size in checkpoint
   chunk_byte_size = (
-      config.checkpoint_storage_target_data_file_size_bytes if config else DEFAULT_OCDBT_TARGET_DATA_FILE_SIZE
+      getattr(
+          config,
+          "checkpoint_storage_target_data_file_size_bytes",
+          DEFAULT_OCDBT_TARGET_DATA_FILE_SIZE,
+      )
+      if config
+      else DEFAULT_OCDBT_TARGET_DATA_FILE_SIZE
   )
 
   checkpoint_args = ocp.args.PyTreeSave(
@@ -822,7 +830,11 @@ def save_checkpoint(checkpoint_manager, step, state, config=None, data_iterator=
   )
   save_args_composite = {"items": checkpoint_args}
 
-  if config and config.dataset_type == "grain" and not isinstance(data_iterator, PlaceHolderDataIterator):
+  if (
+      config
+      and getattr(config, "dataset_type", None) == "grain"
+      and not isinstance(data_iterator, PlaceHolderDataIterator)
+  ):
     if isinstance(data_iterator, RemoteIteratorWrapper):
       # Pass the wrapper directly; GrainCheckpointHandler will call save_state with the step
       save_args_composite["iter"] = grain_utility.GrainCheckpointSave(

tests/unit/train_state_nnx_checkpoint_test.py

@@ -434,13 +434,31 @@ def test_nnx_and_linen_agree_on_actual_step(self):
     )
 
   def test_nnx_state_is_saved_in_linen_layout(self):
-    """For pure_nnx=True, maybe_save_checkpoint reshapes the NNX state to the Linen on-disk layout."""
+    """For pure_nnx=True, save_checkpoint reshapes the NNX state to the Linen on-disk layout."""
     state = self._build_nnx_state(self.N_STEPS)
     self.assertIsInstance(state, nnx.State)  # precondition: NNX train_step returns an nnx.State
 
-    captured = self._invoke_maybe_save(state, pure_nnx=True)
+    config = self._config(pure_nnx=True, enable_checkpointing=True, checkpoint_period=1)
+    mgr = mock.MagicMock()
+    mgr.reached_preemption.return_value = False
+
+    captured = {}
 
-    # save_checkpoint should receive a plain dict in Linen layout, not the nnx.State.
+    def fake_save(_step, *args, **kwargs):
+      composite = kwargs.get("args")
+      if composite:
+        for key in ["items", "state"]:
+          if hasattr(composite, "_items") and key in composite._items:  # pylint: disable=protected-access
+            val = composite[key]
+            if val is not None and hasattr(val, "item"):
+              captured["state"] = val.item
+              break
+      return True
+
+    mgr.save.side_effect = fake_save
+    checkpointing.save_checkpoint(mgr, self.N_STEPS - 1, state, config=config, force=True)
+
+    # save_checkpoint should pass a plain dict in Linen layout to Orbax, not the nnx.State.
     self.assertIsInstance(captured["state"], dict)
     self.assertNotIsInstance(captured["state"], nnx.State)
     # Linen layout: {params: {params: ...}, step, opt_state}; not the NNX {model, optimizer}.
@@ -558,16 +576,12 @@ def test_maybe_save_checkpoint_allows_local_checkpoint_period(self):
         mgr.reached_preemption.return_value = False
         save_checkpoint_mock = mock.MagicMock(return_value=False)
 
-        with (
-            mock.patch.object(checkpointing, "save_checkpoint", save_checkpoint_mock),
-            mock.patch.object(train_state_nnx, "to_checkpoint_dict") as to_checkpoint_dict_mock,
-        ):
+        with mock.patch.object(checkpointing, "save_checkpoint", save_checkpoint_mock):
           checkpointing.maybe_save_checkpoint(mgr, state, config, data_iterator=None, step=5)
 
         mgr.latest_step.assert_called_once_with()
         mgr.reached_preemption.assert_called_once_with(5)
         mgr.wait_until_finished.assert_not_called()
-        to_checkpoint_dict_mock.assert_called_once_with(state)
         save_checkpoint_mock.assert_called_once()
 
   def test_maybe_save_checkpoint_allows_mtc_period_with_continuous_policy(
@@ -587,16 +601,12 @@ def test_maybe_save_checkpoint_allows_mtc_period_with_continuous_policy(
     mgr.reached_preemption.return_value = False
     save_checkpoint_mock = mock.MagicMock(return_value=False)
 
-    with (
-        mock.patch.object(checkpointing, "save_checkpoint", save_checkpoint_mock),
-        mock.patch.object(train_state_nnx, "to_checkpoint_dict") as to_checkpoint_dict_mock,
-    ):
+    with mock.patch.object(checkpointing, "save_checkpoint", save_checkpoint_mock):
       checkpointing.maybe_save_checkpoint(mgr, state, config, data_iterator=None, step=5)
 
     mgr.should_save.assert_called_once_with(5)
     mgr.latest_step.assert_called_once_with()
     mgr.reached_preemption.assert_called_once_with(5)
-    to_checkpoint_dict_mock.assert_called_once_with(state)
     save_checkpoint_mock.assert_called_once()
 
   def test_maybe_save_checkpoint_checks_scale_up_after_unsaved_dispatch(self):
@@ -610,12 +620,10 @@ def test_maybe_save_checkpoint_checks_scale_up_after_unsaved_dispatch(self):
 
     with (
         mock.patch.object(checkpointing, "save_checkpoint", save_checkpoint_mock),
-        mock.patch.object(train_state_nnx, "to_checkpoint_dict", return_value={}) as to_checkpoint_dict_mock,
         mock.patch.object(checkpointing.elastic_utils, "maybe_elastic_scale_up") as mock_maybe_scale_up,
     ):
       checkpointing.maybe_save_checkpoint(mgr, state, config, data_iterator=None, step=5)
 
-    to_checkpoint_dict_mock.assert_called_once_with(state)
     save_checkpoint_mock.assert_called_once()
     mock_maybe_scale_up.assert_called_once_with(config, mgr)