Return was_restored flag from setup_training_state to fix standalone checkpointer restoration

Author: xibinliu

src/maxtext/checkpoint_conversion/standalone_scripts/convert_gpt3_ckpt_from_paxml.py

@@ -117,7 +117,7 @@ def init_state_fn():
       cfg.checkpoint_period,
   )
 
-  state, _, _, _ = maxtext_utils.setup_training_state(None, cfg, mesh, checkpoint_manager, init_state_fn)
+  state, _, _, _, _ = maxtext_utils.setup_training_state(None, cfg, mesh, checkpoint_manager, init_state_fn)
   max_logging.log("start")
   max_utils.print_mem_stats("After params initialized")
 

src/maxtext/experimental/rl/grpo_trainer.py

@@ -864,7 +864,7 @@ def init_state_fn():
 
   with maybe_record_goodput(recorder, GoodputEvent.TRAINING_PREPARATION):
     data_iterator = grpo_input_pipeline.create_data_iterator(config_inference, inference_mesh)
-    state, _, state_mesh_shardings, data_iterator = maxtext_utils.setup_training_state(
+    state, _, state_mesh_shardings, data_iterator, _ = maxtext_utils.setup_training_state(
         data_iterator, config, mesh, checkpoint_manager, init_state_fn
     )
 

src/maxtext/utils/generate_param_only_checkpoint.py

@@ -167,7 +167,7 @@ def init_state_fn():
     tx = optimizers.get_optimizer(config, learning_rate_schedule)
     init_state_fn = functools.partial(maxtext_utils.init_initial_state, model, tx, config, True, rng)
 
-  state, state_mesh_notations, state_mesh_shardings, _ = maxtext_utils.setup_training_state(
+  state, state_mesh_notations, state_mesh_shardings, _, _ = maxtext_utils.setup_training_state(
       None, config, mesh, checkpoint_manager, init_state_fn
   )
   if config.pure_nnx:
@@ -400,13 +400,9 @@ def generate_decode_checkpoint(config):
   training_state, training_state_annotations = _read_train_checkpoint(config, checkpoint_manager, mesh)
   if config.pure_nnx:
     # NNX state is a flat nnx.State; opt_state lives under the optimizer sub-state.
-    assert (
-        training_state.optimizer.opt_state
-    ), "missing opt_state in training checkpoint"
+    assert training_state.optimizer.opt_state, "missing opt_state in training checkpoint"
   else:
-    assert (
-        training_state.opt_state != {}
-    ), "missing opt_state in training checkpoint"
+    assert training_state.opt_state != {}, "missing opt_state in training checkpoint"
 
   _possibly_unroll_params(config, training_state, training_state_annotations, mesh)
 

src/maxtext/utils/maxtext_utils.py

@@ -1392,7 +1392,9 @@ def get_abstract_param(model, config):
       {"params": key, "dropout": key, "aqt": key},
       np.ones(input_shape, dtype=jnp.int32),
       np.ones(input_shape, dtype=jnp.int32),
-      encoder_images=np.ones(image_shape, dtype=jnp.int32) if config.use_multimodal else None,  # pyrefly: ignore[no-matching-overload]
+      encoder_images=np.ones(image_shape, dtype=jnp.int32)
+      if config.use_multimodal
+      else None,  # pyrefly: ignore[no-matching-overload]
       encoder_audios=np.ones(audio_shape, dtype=jnp.float32) if config.use_audio else None,
   )
   return abstract_vars
@@ -1413,7 +1415,7 @@ def setup_decode_state(config, mesh, checkpoint_manager, init_state_fn):
   if not config.load_parameters_path:
     # generate random params
     max_logging.log("No decode checkpoint specified - generating random weights.")
-    state, state_mesh_annotations, _, _ = setup_initial_state(
+    state, state_mesh_annotations, _, _, _ = setup_initial_state(
         None, config, mesh, checkpoint_manager, init_state_fn, False
     )
   else:
@@ -1468,6 +1470,9 @@ def setup_initial_state(
   Returns:
     train_state: the initialized train state. For NNX, this is a TrainStateNNX instance
     state_mesh_annotations: the mesh annotations for the train state
+    state_mesh_shardings: the mesh shardings for the train state
+    data_iterator: the updated data iterator
+    was_restored: True if state or params were restored from checkpoint, False if freshly initialized
   """
 
   unboxed_abstract_state, state_mesh_annotations, state_mesh_shardings = get_abstract_state(
@@ -1493,6 +1498,8 @@ def setup_initial_state(
         expansion_factor_real_data=config.expansion_factor_real_data,
         maxtext_config=config,
     )
+    # Partial or fully restored
+    was_restored = bool(restored is not None or raw_params is not None)
 
     if restored:
       if isinstance(
@@ -1548,7 +1555,7 @@ def _merge_params(p_raw, p_init):
             state = state.replace(params=raw_params)
   if not config.pure_nnx:
     state = max_utils.unbox_logicallypartioned(state)
-  return state, state_mesh_annotations, state_mesh_shardings, data_iterator
+  return state, state_mesh_annotations, state_mesh_shardings, data_iterator, was_restored
 
 
 def get_logical_annotations(config, mesh, init_state_fn):

src/maxtext/utils/standalone_checkpointer.py

@@ -25,7 +25,6 @@
 
 from absl import app
 from flax import nnx
-from flax.linen import partitioning as nn_partitioning
 import jax
 from jax import numpy as jnp
 from maxtext.configs import pyconfig
@@ -69,44 +68,34 @@ def init_state_fn():
     mesh = model.mesh
     _, tx = train_utils.create_training_optimizer(config, model)
     init_state_fn = partial(maxtext_utils.init_initial_state, model, tx, config, True, init_rng)
+
   checkpoint_manager = train_utils.create_checkpoint_manager(config, mesh, init_state_fn)
 
-  unboxed_abstract_state, _, _ = maxtext_utils.get_abstract_state(config, mesh, init_state_fn, is_training=True)
   # A barrier to sync all hosts before starting to restore checkpoint
   jax.experimental.multihost_utils.sync_global_devices("Barrier before load")
-  checkpoint_load_start = datetime.datetime.now()
-  with nn_partitioning.axis_rules(config.logical_axis_rules):
-    state, _ = checkpointing.load_state_if_possible(
-        checkpoint_manager,
-        None,
-        config.load_parameters_path,
-        config.load_full_state_path,
-        config.checkpoint_storage_concurrent_gb,
-        unboxed_abstract_state,
-        use_ocdbt=config.checkpoint_storage_use_ocdbt,
-        use_zarr3=config.checkpoint_storage_use_zarr3,
-    )
-    if state:
-      state = state["items"]
 
+  checkpoint_load_start = datetime.datetime.now()
+  # Delegate checkpoint restoration or state initialization to setup_training_state
+  state, _, _, _, was_restored = maxtext_utils.setup_training_state(None, config, mesh, checkpoint_manager, init_state_fn)
   jax.block_until_ready(state)
   checkpoint_load_end = datetime.datetime.now()
-  if state is not None:  # Checkpoint was available for restore
+
+  if was_restored:
     if jax.process_index() == 0:
       max_logging.log(
           "STANDALONE CHECKPOINTER : Checkpoint restored in :" f" {checkpoint_load_end - checkpoint_load_start}"
       )
-  else:  # Checkpoint was unavailable, state needs to be initialized
-    state, _, _, _ = maxtext_utils.setup_training_state(None, config, mesh, checkpoint_manager, init_state_fn)
-  state = add_entropy_to_checkpoint(state)
+  else:  # Checkpoint was unavailable, fresh state needs to be perturbed with entropy
+    state = add_entropy_to_checkpoint(state)
 
   start_step = get_first_step(model, state)  # this is the start_step for training
   for step in np.arange(start_step, config.steps):
     if checkpoint_manager is not None:
       start_time = datetime.datetime.now()
       # A barrier to sync all hosts before starting to save checkpoint
       jax.experimental.multihost_utils.sync_global_devices("Barrier before save")
-      if checkpointing.save_checkpoint(checkpoint_manager, int(step), state):
+      state_to_save = train_state_nnx.to_linen_checkpoint_dict(state.to_pure_dict()) if config.pure_nnx else state
+      if checkpointing.save_checkpoint(checkpoint_manager, int(step), state_to_save):
         checkpoint_manager.wait_until_finished()
         end_time = datetime.datetime.now()
         if jax.process_index() == 0:

src/maxtext/utils/train_utils.py

@@ -299,7 +299,7 @@ def create_train_state_fn():
     # Create data_loader AFTER reordering wrapper is applied
     data_loader = create_dataloader(config, mesh, data_iterator, recorder, rampup_manager)
 
-    state, _, state_mesh_shardings, data_iterator = maxtext_utils.setup_training_state(
+    state, _, state_mesh_shardings, data_iterator, _ = maxtext_utils.setup_training_state(
         data_iterator, config, mesh, checkpoint_manager, init_state_fn
     )
     if config.pure_nnx:
@@ -321,11 +321,7 @@ def create_train_state_fn():
         state, outer_opt_state_sharding = diloco.build_diloco_state(config, lambda: state, mesh=mesh)
 
         # create state_mesh_shardings for the DilocoState
-        step_mesh = (
-            state_mesh_shardings.optimizer.step.mesh
-            if config.pure_nnx
-            else state_mesh_shardings.step.mesh
-        )
+        step_mesh = state_mesh_shardings.optimizer.step.mesh if config.pure_nnx else state_mesh_shardings.step.mesh
         inner_state_shardings = diloco.add_diloco_to_sharding(state_mesh_shardings)
         state_mesh_shardings = diloco.DiLoCoTrainState(
             inner_state_shardings,

tests/unit/maxtext_utils_test.py

@@ -387,7 +387,8 @@ def create_train_state_fn():
       init_state_fn = create_train_state_fn
     else:
       init_state_fn = functools.partial(maxtext_utils.init_initial_state, self.model, tx, self.config, True, rng)
-    state, _, _, _ = maxtext_utils.setup_initial_state(None, self.config, self.mesh, None, init_state_fn)
+    state, _, _, _, was_restored = maxtext_utils.setup_initial_state(None, self.config, self.mesh, None, init_state_fn)
+    self.assertFalse(was_restored)
     if self.config.pure_nnx:
       self.assertIsNotNone(state.optimizer)
     else:
@@ -1413,7 +1414,8 @@ def create_train_state_fn():
           True,
           rng,
       )
-    state, _, _, _ = maxtext_utils.setup_training_state(None, self.config, self.mesh, None, init_state_fn)
+    state, _, _, _, was_restored = maxtext_utils.setup_training_state(None, self.config, self.mesh, None, init_state_fn)
+    self.assertFalse(was_restored)
     if self.config.pure_nnx:
       self.assertIsNotNone(state.optimizer)
     else: