Merge pull request AI-Hypercomputer#3526 from AI-Hypercomputer:feat/nnx-set-defaults-true

PiperOrigin-RevId: 940139961

Author: Google-ML-Automation

src/maxtext/checkpoint_conversion/to_maxtext.py

@@ -320,22 +320,23 @@ def get_maxtext_model_info(config):
   # Get abstract model structure (name, shape) without materializing the weights to save memory
   abstract_params_tree = maxtext_utils.get_abstract_param(maxtext_model_flax, config)["params"]
 
-  abstract_params_flat, _ = jax.tree_util.tree_flatten_with_path(abstract_params_tree)
-  # Standardize abstract tree for later unflattening
-  abstract_params_tree = jax.tree.map(
-      lambda _: 0,
-      abstract_params_tree,
-      is_leaf=lambda x: isinstance(x, nn.LogicallyPartitioned),
+  abstract_params_flat, abstract_params_treedef = (
+      jax.tree_util.tree_flatten_with_path(
+          abstract_params_tree,
+          is_leaf=lambda x: isinstance(x, nn.LogicallyPartitioned),
+      )
   )
-  abstract_params_treedef = jax.tree_util.tree_structure(abstract_params_tree)
 
   max_logging.log("MaxText abstract model and state initialized.")
 
   # preprocess state
   maxtext_abstract_dict = {}
   for mt_target_idx, (path_tuple, abstract_leaf_value) in enumerate(abstract_params_flat):
     mt_param_key = "params-" + "-".join(param_key_parts_from_path(path_tuple))
-    mt_target_shape = abstract_leaf_value.shape
+    if isinstance(abstract_leaf_value, nn.LogicallyPartitioned):
+      mt_target_shape = abstract_leaf_value.value.shape
+    else:
+      mt_target_shape = abstract_leaf_value.shape
     maxtext_abstract_dict[mt_param_key] = (mt_target_idx, mt_target_shape)
 
   return maxtext_abstract_dict, abstract_params_treedef

src/maxtext/checkpoint_conversion/utils/utils.py

@@ -968,7 +968,9 @@ def extract_nnx_weights(weights_dict: dict) -> dict[str, np.ndarray]:
   for path_tuple, leaf_value in leaves_with_paths:
     path_keys = param_key_parts_from_path(path_tuple)
     # Skip NNX RNG state variables (not model weights)
-    if "to_nnx__rngs" in path_keys or any(k.endswith("_rngs") for k in path_keys):
+    if "to_nnx__rngs" in path_keys or any(
+        k == "rngs" or k.endswith("_rngs") for k in path_keys
+    ):
       continue
     maxtext_param_key = "params-" + "-".join(path_keys)
     if not isinstance(leaf_value, (jax.Array, np.ndarray)):

src/maxtext/common/checkpointing.py

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+
 """Create an Orbax CheckpointManager with specified (Async or not) Checkpointer."""
 
 import time
@@ -357,11 +358,19 @@ def combine_sharding(sds, shardings):
         use_ocdbt=use_ocdbt,
         use_zarr3=use_zarr3,
     )
+    # NNX checkpoints are saved as pure dicts (see maybe_save_checkpoint); the
+    # restore target must match — a boxed nnx.State wouldn't.
+    restore_target = abstract_unboxed_pre_state
+    if isinstance(abstract_unboxed_pre_state, nnx.State):
+      restore_target = abstract_unboxed_pre_state.to_pure_dict()
     # Provide sharding info to ensure restoration returns JAX arrays (not NumPy arrays).
     restore_args = jax.tree_util.tree_map(
-        lambda x: ocp.type_handlers.ArrayRestoreArgs(sharding=x.sharding), abstract_unboxed_pre_state
+        lambda x: ocp.type_handlers.ArrayRestoreArgs(sharding=x.sharding),
+        restore_target,
+    )
+    return ocp.Checkpointer(handler).restore(
+        p, restore_target, restore_args=restore_args
     )
-    return ocp.Checkpointer(handler).restore(p, abstract_unboxed_pre_state, restore_args=restore_args)
 
 
 def create_orbax_checkpoint_manager(
@@ -1028,7 +1037,10 @@ 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
@@ -1039,7 +1051,19 @@ 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/configs/base.yml

@@ -494,6 +494,7 @@ shard_mode: "auto" # can be either auto or explicit
 custom_mesh_and_rule: "" # replace default mesh and logical rule by specifying yml name under config/mesh_and_rule/.
 mesh_axes: ['diloco', 'data', 'stage', 'fsdp', 'fsdp_transpose', 'context', 'context_autoregressive', 'tensor', 'tensor_transpose', 'tensor_sequence', 'expert', 'autoregressive']
 logical_axis_rules: [
+                      ['circular_repeats', []],
                       # ==========================================
                       # Vocabulary Embedding
                       # ==========================================
@@ -1235,9 +1236,9 @@ position_id_per_seconds: 25
 subslice_shape: ""
 
 # NNX
-enable_nnx: false
-pure_nnx_decoder: false
-pure_nnx: false
+enable_nnx: true
+pure_nnx_decoder: true
+pure_nnx: true
 
 ################################## Qwen3-Next Specific Configs ##################################
 # Kernel size for the 1D convolution in the Gated Delta Net

src/maxtext/layers/nnx_decoders.py

@@ -18,33 +18,34 @@
 
 import functools
 import inspect
-import warnings
 from typing import Any
+import warnings
 
-import jax
-import jax.numpy as jnp
 from flax import linen as nn
 from flax import nnx
-from maxtext.layers import nnx_wrappers
+import jax
 from jax.ad_checkpoint import checkpoint_name
+import jax.numpy as jnp
 from jax.sharding import Mesh
-
 from maxtext.common.common_types import (
+    Config,
+    DecoderBlockType,
     MODEL_MODE_AUTOREGRESSIVE,
     MODEL_MODE_PREFILL,
     MODEL_MODE_TRAIN,
-    Config,
-    DecoderBlockType,
     MultimodalInput,
     ShardMode,
 )
 from maxtext.layers import initializers, linears, mhc, normalizations, quantizations
+from maxtext.layers import nnx_wrappers
 from maxtext.layers.attentions import Attention
 from maxtext.layers.embeddings import Embed, PositionalEmbedding, attend_on_embedding
 from maxtext.layers.normalizations import RMSNorm
+from maxtext.layers.pipeline import create_nnx_pipeline
 from maxtext.layers.quantizations import AqtQuantization as Quant
 from maxtext.models import (
     deepseek,
+    deepseek4,
     deepseek_batchsplit,
     deepseek_batchsplit_fp8,
     gemma,
@@ -68,7 +69,6 @@
 from maxtext.multimodal import utils as mm_utils
 from maxtext.utils import max_logging, max_utils, maxtext_utils, maxtext_utils_nnx, sharding
 from maxtext.utils.sharding import create_sharding
-from maxtext.layers.pipeline import create_nnx_pipeline
 
 # ------------------------------------------------------------------------------
 # The network: Decoder Definitions
@@ -1111,21 +1111,38 @@ def get_deepseek():
         DecoderBlockType.GEMMA: [gemma.GemmaDecoderLayer],
         DecoderBlockType.GEMMA2: [gemma2.Gemma2DecoderLayer],
         DecoderBlockType.GEMMA3: [gemma3.Gemma3DecoderLayer],
-        DecoderBlockType.GEMMA4: get_scannable(gemma4.Gemma4DecoderLayer, gemma4.Gemma4ScannableBlock),
+        DecoderBlockType.GEMMA4: get_scannable(
+            gemma4.Gemma4DecoderLayer, gemma4.Gemma4ScannableBlock
+        ),
         DecoderBlockType.GEMMA4_SMALL: [gemma4_small.Gemma4SmallDecoderLayer],
         DecoderBlockType.GPT3: [gpt3.Gpt3DecoderLayer],
         DecoderBlockType.QWEN2: [qwen2.Qwen2DecoderLayer],
         DecoderBlockType.QWEN3: [qwen3.Qwen3DecoderLayer],
         DecoderBlockType.QWEN3_MOE: [qwen3.Qwen3MoeDecoderLayer],
-        DecoderBlockType.QWEN3_CUSTOM_MOE: [qwen3_custom.Qwen3CustomMoeDecoderLayer],
+        DecoderBlockType.QWEN3_CUSTOM_MOE: [
+            qwen3_custom.Qwen3CustomMoeDecoderLayer
+        ],
         DecoderBlockType.SIMPLE: [simple_layer.SimpleDecoderLayer],
         DecoderBlockType.SIMPLE_MLP: [simple_layer.SimpleMlpDecoderLayer],
         DecoderBlockType.DEEPSEEK: get_deepseek(),
-        DecoderBlockType.GPT_OSS: get_scannable(gpt_oss.GptOssDecoderLayer, gpt_oss.GptOssScannableBlock),
-        DecoderBlockType.QWEN3_NEXT: get_scannable(qwen3.Qwen3NextDecoderLayer, qwen3.Qwen3NextScannableBlock),
-        DecoderBlockType.QWEN3_5: get_scannable(qwen3_5.Qwen3_5DecoderLayer, qwen3_5.Qwen3_5ScannableBlock),
-        DecoderBlockType.LLAMA4: get_scannable(llama4.Llama4DecoderLayer, llama4.Llama4ScannableBlock),
-        DecoderBlockType.OLMO3: get_scannable(olmo3.Olmo3DecoderLayer, olmo3.Olmo3ScannableBlock),
+        DecoderBlockType.DEEPSEEK4: get_scannable(
+            deepseek4.DeepSeek4DecoderLayer, deepseek4.DeepSeek4ScannableBlock
+        ),
+        DecoderBlockType.GPT_OSS: get_scannable(
+            gpt_oss.GptOssDecoderLayer, gpt_oss.GptOssScannableBlock
+        ),
+        DecoderBlockType.QWEN3_NEXT: get_scannable(
+            qwen3.Qwen3NextDecoderLayer, qwen3.Qwen3NextScannableBlock
+        ),
+        DecoderBlockType.QWEN3_5: get_scannable(
+            qwen3_5.Qwen3_5DecoderLayer, qwen3_5.Qwen3_5ScannableBlock
+        ),
+        DecoderBlockType.LLAMA4: get_scannable(
+            llama4.Llama4DecoderLayer, llama4.Llama4ScannableBlock
+        ),
+        DecoderBlockType.OLMO3: get_scannable(
+            olmo3.Olmo3DecoderLayer, olmo3.Olmo3ScannableBlock
+        ),
     }
 
     if cfg.decoder_block not in layer_map:

src/maxtext/layers/nnx_wrappers.py

@@ -26,7 +26,6 @@
 from flax.core import FrozenDict
 from flax.core import meta
 from flax.nnx import graph
-from flax.nnx import tracers as nnx_tracers
 from flax.nnx import variablelib
 from flax.nnx.bridge import module as bdg_module
 from flax.nnx.module import Module
@@ -183,19 +182,6 @@ def is_linen_initializing() -> bool:
   return False
 
 
-def _refresh_variable_trace_state(module: Module) -> None:
-  """Resets stale ``_trace_state`` on Variables to unblock downstream ``nnx.split``.
-
-  ``nnx.update`` called with JAX tracer values uses ``_unsafe_bypass_check=True``,
-  which leaves Variables with a stale ``_trace_state`` from the outer Python
-  context and breaks ``nnx.split`` with "Cannot extract graph node from different
-  trace level". Resets ``_trace_state`` on any Variable whose ``_can_update`` is False.
-  """
-  for _, v in nnx.graph.iter_graph(module):
-    if isinstance(v, variablelib.Variable) and not v._can_update:  # pylint: disable=protected-access
-      object.__setattr__(v, "_trace_state", nnx_tracers.TraceState())
-
-
 class ToNNX(Module):
   """A wrapper to turn any Linen module into an NNX module.
 
@@ -542,9 +528,21 @@ def maybe_unbox(x):
           f"Found unknown module paths in incoming state:{paths_str}. Intermediate modules have been reconstructed."
       )
 
+      # Filter out unknown paths so we don't try to assign them to static attributes
+      filtered_state_flat = {
+          k: v for k, v in new_state_flat.items() if k not in unknown_state_flat
+      }
+      new_state = nnx.State(
+          nnx.traversals.unflatten_mapping(filtered_state_flat)
+      )
+
+    # Rebind the module to the current trace via split / update / merge.
+    # nnx.update directly on the live module can leave stale tracers.
+    graphdef, full_state = nnx.split(module)
+    nnx.update(full_state, new_state)
+    module = nnx.merge(graphdef, full_state)
+
     _fix_for_qwix_quantization(module)
-    nnx.update(module, new_state)
-    _refresh_variable_trace_state(module)
     method_fn = _get_module_method(module, nnx_method)
     out = method_fn(module, *args, **kwargs)
     self._update_variables(module)

src/maxtext/layers/quantizations.py

@@ -1031,8 +1031,7 @@ def generate_quantizer_set(
           fp8_recipe: recipe.Recipe | None = default_recipe,
           n_groups: int | None = None,
       ):
-        """Generates a set of quantizers for TransformerEngine."""
-
+        """Route quantizer-set generation through the OVERWRITE_WITH_GRADIENT collection."""
         OVERWRITE_WITH_GRADIENT = "_overwrite_with_gradient"
         return super().generate_quantizer_set(  # pytype: disable=wrong-keyword-args
             postfix=postfix,

src/maxtext/trainers/diloco/diloco.py

@@ -28,13 +28,12 @@
 import drjax
 from flax import nnx
 from flax import struct
-from flax.training import train_state
 import jax
 import jax.numpy as jnp
 from jaxtyping import Array, Int32, Key, PyTree, UInt32
-import optax
-
+from maxtext.common.train_state_nnx import TrainStateNNX
 from maxtext.configs import pyconfig
+import optax
 
 Batch = Any
 Params = PyTree
@@ -157,8 +156,12 @@ def add_diloco_dim(x):
   # For NNX, model params (Param variables only) live under abstract_state.model;
   # for Linen under abstract_state.params.
   if config.pure_nnx:
-    model_params = abstract_state.model.filter(nnx.Param)
-    model_params_sharding = state_mesh_shardings.model.filter(nnx.Param)
+    _, model_params, _ = nnx.split(abstract_state.model, nnx.Param, ...)
+    model_params = model_params.to_pure_dict()
+    _, model_params_sharding, _ = nnx.split(
+        state_mesh_shardings.model, nnx.Param, ...
+    )
+    model_params_sharding = model_params_sharding.to_pure_dict()
   else:
     model_params = abstract_state.params
     model_params_sharding = state_mesh_shardings.params
@@ -216,7 +219,11 @@ def init_diloco_state() -> tuple[DiLoCoTrainState, PyTree]:
     # Outer state retains a single copy of the model parameters and optimizer state.
     # For NNX, model params (Param variables only) live under state.model;
     # for Linen under state.params.
-    outer_params = state.model.filter(nnx.Param) if config.pure_nnx else state.params
+    if config.pure_nnx:
+      _, outer_params, _ = nnx.split(state.model, nnx.Param, ...)
+      outer_params = outer_params.to_pure_dict()
+    else:
+      outer_params = state.params
     outer_opt_state = outer_optimizer.init(outer_params)
     outer_opt_state_sharding = jax.tree_util.tree_map(lambda x: x.sharding, outer_opt_state)
     # For NNX, the step counter lives at state.optimizer.step; for Linen at state.step.
@@ -258,9 +265,13 @@ def synchronize(state):
     # state (since last synchronization).
     broadcast_outer_params = drjax.broadcast(state.params, mesh=mesh)
     # For NNX, model Param vars live under inner_state.model; for Linen under inner_state.params.
-    inner_model_params = (
-        nnx.filter_state(state.inner_state.model, nnx.Param) if config.pure_nnx else state.inner_state.params
-    )
+    if config.pure_nnx:
+      _, inner_model_params, _ = nnx.split(
+          state.inner_state.model, nnx.Param, ...
+      )
+      inner_model_params = inner_model_params.to_pure_dict()
+    else:
+      inner_model_params = state.inner_state.params
     model_delta = jax.tree.map(lambda x, y: y - x, inner_model_params, broadcast_outer_params)
     # Treat the average delta as the outer optimizer's gradient and apply to
     # the global (outer) model params.
@@ -273,15 +284,40 @@ def synchronize(state):
     if config.pure_nnx:
       # For NNX: merge new Param vars back with the non-Param model vars (e.g. RNG state).
       def replace_nnx_model_params(s, new_params):
-        non_param_model = nnx.filter_state(s.model, nnx.Not(nnx.Param))
-        new_model = nnx.merge_state(non_param_model, new_params)
-        # Assign via __setitem__ so nested States are stored as plain dicts (matching
-        # nnx.state()'s pytree structure). The dict-literal constructor keeps them as
-        # State objects, which makes jax.lax.cond see mismatched pytree structures.
-        result = type(s)({})
-        result["model"] = new_model
-        result["optimizer"] = s["optimizer"]
-        return result
+        s_model = s["model"] if hasattr(s, "keys") else s.model
+        s_opt = s["optimizer"] if hasattr(s, "keys") else s.optimizer
+
+        graphdef, _, non_param_state = nnx.split(s_model, nnx.Param, ...)
+        new_model = nnx.merge(graphdef, new_params, non_param_state)
+
+        if type(s_model).__name__ == "State":
+          new_model = nnx.state(new_model)
+        elif isinstance(s_model, dict):
+          new_model = nnx.to_pure_dict(new_model)
+
+        if hasattr(s, "keys"):
+          # Replace "model" leaves by path, keeping s's treedef. Picking by position
+          # (leaves[N:]) breaks if a key sorts before "model"; reconstructing via
+          # type(s)({...}) breaks the lax.cond match — nnx.State recursive-wraps.
+          leaves_with_paths, treedef = jax.tree_util.tree_flatten_with_path(s)
+          new_model_iter = iter(jax.tree_util.tree_leaves(new_model))
+
+          def _is_model_leaf(path):
+            if not path:
+              return False
+            k = path[0]
+            return (
+                getattr(k, "key", None) == "model"
+                or getattr(k, "name", None) == "model"
+            )
+
+          new_leaves = [
+              next(new_model_iter) if _is_model_leaf(p) else leaf
+              for p, leaf in leaves_with_paths
+          ]
+          return jax.tree_util.tree_unflatten(treedef, new_leaves)
+        else:
+          return TrainStateNNX(new_model, s_opt)
 
       new_inner_state = drjax.map_fn(
           lambda s: replace_nnx_model_params(s, new_outer_params),