NNX: fix DiLoCo train loop + checkpoint under pure_nnx

Under enable_diloco the state becomes a DiLoCoTrainState, but the pure_nnx path
still merged it against the plain-model graphdef (nnx.merge leaf mismatch +
segfault), and several downstream sites assumed a plain TrainStateNNX. Guard the
merge and surface the graphdef as model; fix get_first_step, jit_model,
params_shardings, setup_params, and the rng args in train_loop; match the diloco
sharding's params to_pure_dict; and handle the DiLoCoTrainState in
maybe_save_checkpoint by saving the synchronized global model. Train + checkpoint
save/restore validated end-to-end on CPU.

Author: ecnal-cienet

src/maxtext/common/checkpointing.py

@@ -1040,7 +1040,8 @@ def maybe_save_checkpoint(checkpoint_manager, state, config, data_iterator, step
     actual_step = int(step)
   else:
     if config.pure_nnx:
-      actual_step = int(state.optimizer.step) - 1
+      # Under DiLoCo the step lives on the DiLoCoTrainState; otherwise on the optimizer.
+      actual_step = int(state.step if config.enable_diloco else state.optimizer.step) - 1
     else:
       # Linen TrainState has .step attribute
       actual_step = int(state.step) - 1
@@ -1051,7 +1052,13 @@ def maybe_save_checkpoint(checkpoint_manager, state, config, data_iterator, step
 
   if config.pure_nnx:
     # Save in the Linen on-disk layout so pure_nnx and Linen checkpoints are interchangeable.
-    state = train_state_nnx.to_linen_checkpoint_dict(state.to_pure_dict())
+    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:
+      state = train_state_nnx.to_linen_checkpoint_dict(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/trainers/pre_train/train.py

@@ -77,6 +77,8 @@
 def get_first_step(model, state):
   if isinstance(model, nn.Module):
     return int(state.step)
+  if hasattr(state, "inner_state"):  # DiLoCoTrainState (NNX DiLoCo): step is the optimizer step var
+    return int(state.step.get_value())
   return int(state.optimizer.step.get_value())
 
 
@@ -598,10 +600,19 @@ def train_loop(config, recorder, state=None):
 
   if isinstance(model, nn.Module):
     jit_model = model
+  elif config.enable_diloco:
+    # state is the DiLoCoTrainState; `model` is already the TrainStateNNX graphdef the inner step needs.
+    jit_model = model
   else:
     jit_model, state = nnx.split(state)
 
-  params_shardings, state_mesh_shardings = sharding.maybe_update_params_sharding_with_opt(config, state_mesh_shardings)
+  if config.pure_nnx and config.enable_diloco:
+    # state_mesh_shardings is a DiLoCoTrainState; its .params field already holds the
+    # plain param shardings the inner train step needs (the Zero-1 opt update, which
+    # reads .model/.optimizer, doesn't apply to the wrapped diloco shardings here).
+    params_shardings = state_mesh_shardings.params
+  else:
+    params_shardings, state_mesh_shardings = sharding.maybe_update_params_sharding_with_opt(config, state_mesh_shardings)
 
   p_train_step, p_eval_step = train_utils.jit_train_and_eval_step(
       config,
@@ -622,7 +633,8 @@ def train_loop(config, recorder, state=None):
     elif config.shard_optimizer_over_data:
       # NNX: reshard state so params match the data-sharded in_shardings (Zero-1 layout)
       state = jax.device_put(state, state_mesh_shardings)
-    if isinstance(model, nn.Module):
+    if isinstance(model, nn.Module) or config.enable_diloco:
+      # The DiLoCo train step takes (state, batch, rng), like the Linen step.
       lower_args = (state, shaped_batch, init_rng)
     else:
       lower_args = (state, shaped_batch)
@@ -638,6 +650,8 @@ def train_loop(config, recorder, state=None):
   # Write train config params, num model params, and XLA flags to tensorboard
   if isinstance(model, nn.Module):
     setup_params = state.params
+  elif config.enable_diloco:
+    setup_params = state.params  # DiLoCoTrainState.params: the outer (global) params
   else:
     _, setup_params, _ = nnx.split(state.model, nnx.Param, ...)
   metric_logger_instance.write_setup_info_to_tensorboard(setup_params)
@@ -652,7 +666,7 @@ def train_loop(config, recorder, state=None):
 
       with jax.profiler.StepTraceAnnotation("train", step_num=step):
         example_batch = data_loader.load_next_batch(rampup_manager=rampup_manager)
-        if isinstance(model, nn.Module):
+        if isinstance(model, nn.Module) or config.enable_diloco:
           # pylint: disable=not-callable
           step_rng_args = (jax.jit(jax.random.fold_in)(init_rng, step),)
         else:

src/maxtext/utils/train_utils.py

@@ -295,7 +295,9 @@ def create_train_state_fn():
         inner_state_shardings = diloco.add_diloco_to_sharding(state_mesh_shardings)
         state_mesh_shardings = diloco.DiLoCoTrainState(
             inner_state_shardings,
-            state_mesh_shardings_params,
+            # Match the outer params' pure-dict structure (build_diloco_state stores
+            # outer_params via to_pure_dict), so the sharding tree matches the state tree.
+            state_mesh_shardings_params.to_pure_dict() if config.pure_nnx else state_mesh_shardings_params,
             outer_opt_state_sharding,
             jax.sharding.NamedSharding(mesh=step_mesh, spec=jax.sharding.PartitionSpec()),
         )
@@ -319,8 +321,16 @@ def create_train_state_fn():
       maxtext_utils.print_shardings_params(state_params, state_mesh_shardings_params, mesh, logical_annotations_params)
 
   if config.pure_nnx:
-    train_state = nnx.merge(state_graphdef, state)
-    model = train_state.model
+    if config.enable_diloco:
+      # `state` is a DiLoCoTrainState whose inner_state holds the broadcast
+      # TrainStateNNX. Don't merge it into the plain-model graphdef. The inner
+      # train step needs that graphdef, so surface it as `model`: train_loop
+      # uses it as jit_model and feeds the DiLoCoTrainState through as state.
+      train_state = state
+      model = state_graphdef
+    else:
+      train_state = nnx.merge(state_graphdef, state)
+      model = train_state.model
   else:
     train_state = state