[NNX] Delete Linen 5/5: remove the pure_nnx / enable_nnx / pure_nnx_decoder config flags

docs/guides/distillation.md

@@ -63,9 +63,8 @@ Student and teacher must share the same vocabulary. The trainer asserts `student
 If `distill_beta > 0`, the model `sow`s the attention `out_projection` activations at every layer so the loss can read them. This requires:
 
 - `scan_layers: True` — activations are stacked along the leading scan axis; the loss does `jnp.take(features, layer_indices, axis=0)` over that axis.
-- `enable_nnx: True` — `sow(nnx.Intermediate, ...)` is an NNX-specific call.
 
-The trainer validates both at config initialization. Logit-only runs (`distill_beta = 0`) have no such constraint.
+The trainer validates this at config initialization. Logit-only runs (`distill_beta = 0`) have no such constraint.
 
 ## Loss anatomy
 
@@ -221,7 +220,6 @@ distill_feature_loss_type: cosine
 distill_layer_indices: [3, 7, 11, 15, 19, 23, 27, 31]   # for 32-layer student
 
 scan_layers: True
-enable_nnx: True
 ```
 
 ### Logit-only baseline (cheapest; no feature extraction overhead)
@@ -258,7 +256,7 @@ The trainer logs the following to TensorBoard (configured by `tensorboard_dir`,
 
 | Symptom                                                                             | Likely cause                                                                                           | Fix                                                                                                                           |
 | ----------------------------------------------------------------------------------- | ------------------------------------------------------------------------------------------------------ | ----------------------------------------------------------------------------------------------------------------------------- |
-| `ValueError: a value of self.distill_beta > 0.0 requires self.scan_layers = True`   | Feature loss enabled without scanned layers.                                                           | Add `scan_layers=True enable_nnx=True` to your CLI / yml.                                                                     |
+| `ValueError: a value of self.distill_beta > 0.0 requires self.scan_layers = True`   | Feature loss enabled without scanned layers.                                                           | Add `scan_layers=True` to your CLI / yml.                                                                                     |
 | `Vocab size mismatch! Student: X, Teacher: Y`                                       | Different tokenizers.                                                                                  | Use teacher and student with the same vocab; the trainer cannot match logits across vocabularies.                             |
 | `Teacher model path is missing`                                                     | `teacher_overrides.load_parameters_path` not set in non-offline mode.                                  | Set it in `teacher_overrides` in the yml or pass via CLI.                                                                     |
 | `Features extracted from student or teacher model are None, but distill_beta > 0.0` | Model architecture doesn't sow `out_projection_activations` (e.g. uses an unsupported attention path). | Verify the attention layer in use sets `self.sow(nnx.Intermediate, "out_projection_activations", out)` (see `attentions.py`). |

docs/tutorials/posttraining/knowledge_distillation.md

@@ -248,7 +248,7 @@ Key knobs (see the [Distillation guide](../../guides/distillation.md) for the fu
 ```yaml
 distill_alpha: 0.5             # weight on KL(teacher||student)
 distill_temperature: 1.0
-distill_beta: 0.0              # >0 enables feature distillation; requires scan_layers=True, enable_nnx=True
+distill_beta: 0.0              # >0 enables feature distillation; requires scan_layers=True
 distill_layer_indices: None
 ```
 
@@ -312,7 +312,7 @@ python3 -m maxtext.trainers.post_train.distillation.train_distill \
   distill_temperature=2.0 \
   distill_beta=1.0 distill_beta_end=0.1 distill_beta_schedule=cosine \
   distill_layer_indices=[2,5,8,11,14,17,20,23] \
-  scan_layers=True enable_nnx=True \
+  scan_layers=True \
   profiler=xplane
 ```
 

docs/tutorials/posttraining/lora.md

@@ -117,8 +117,6 @@ python3 -m maxtext.trainers.post_train.sft.train_sft \
     per_device_batch_size="${PER_DEVICE_BATCH_SIZE?}" \
     max_target_length="${MAX_TARGET_LENGTH?}" \
     learning_rate="${LEARNING_RATE?}" \
-    enable_nnx=True \
-    pure_nnx_decoder=True \
     lora.enable_lora=True \
     lora.lora_rank="${LORA_RANK?}" \
     lora.lora_alpha="${LORA_ALPHA?}"
@@ -174,8 +172,6 @@ python3 -m maxtext.trainers.post_train.sft.train_sft \
     per_device_batch_size="${PER_DEVICE_BATCH_SIZE?}" \
     max_target_length="${MAX_TARGET_LENGTH?}" \
     learning_rate="${LEARNING_RATE?}" \
-    enable_nnx=True \
-    pure_nnx_decoder=True \
     lora.enable_lora=True \
     lora.lora_rank="${LORA_RANK?}" \
     lora.lora_alpha="${LORA_ALPHA?}"

docs/tutorials/posttraining/lora_on_multi_host.md

@@ -186,8 +186,6 @@ python3 -m maxtext.trainers.post_train.sft.train_sft \
   max_target_length=${MAX_TARGET_LENGTH?} \
   learning_rate=${LEARNING_RATE?} \
   chat_template_path=${CHAT_TEMPLATE_PATH?} \
-  enable_nnx=True \
-  pure_nnx_decoder=True \
   lora.enable_lora=True \
   lora.lora_rank=${LORA_RANK?} \
   lora.lora_alpha=${LORA_ALPHA?} \
@@ -267,8 +265,6 @@ python3 -m maxtext.trainers.post_train.sft.train_sft \
   lora.lora_restore_path=${LORA_RESTORE_PATH?} \
   learning_rate=${LEARNING_RATE?} \
   chat_template_path=${CHAT_TEMPLATE_PATH?} \
-  enable_nnx=True \
-  pure_nnx_decoder=True \
   lora.enable_lora=True \
   lora.lora_rank=${LORA_RANK?} \
   lora.lora_alpha=${LORA_ALPHA?} \

src/maxtext/checkpoint_conversion/inspect_checkpoint.py

@@ -48,7 +48,7 @@
 
   [Mode 2: MaxText Architecture]
     python -m maxtext.checkpoint_conversion.inspect_checkpoint maxtext \
-        model_name=<maxtext_model_name> scan_layers=<True | False> enable_nnx=<True | False>
+        model_name=<maxtext_model_name> scan_layers=<True | False>
         (Optional: other maxtext config)
 
   [Mode 3: Orbax]
@@ -220,13 +220,9 @@ def inspect_maxtext(args, remaining_args):
   import jax
   from maxtext.checkpoint_conversion.utils.utils import param_key_parts_from_path
   from maxtext.configs import pyconfig
-  from maxtext.layers import quantizations
-  from maxtext.models import models
   from maxtext.utils import max_utils, maxtext_utils
   from maxtext.utils.globals import MAXTEXT_PKG_DIR
 
-  Transformer = models.transformer_as_linen
-
   # Configure the PyConfig environment.
   # The first argument in argv is typically the script name.
   argv = (
@@ -242,19 +238,11 @@ def inspect_maxtext(args, remaining_args):
   print(argv)
   config = pyconfig.initialize(argv)
 
-  print(
-      f"\n--- Inspecting MaxText Architecture: {config.model_name} "
-      f"(scan_layers: {config.scan_layers}, enable_nnx: {config.enable_nnx}) ---"
-  )
+  print(f"\n--- Inspecting MaxText Architecture: {config.model_name} " f"(scan_layers: {config.scan_layers}) ---")
   devices_array = maxtext_utils.create_device_mesh(config)
   mesh = jax.sharding.Mesh(devices_array, config.mesh_axes)
-  if config.enable_nnx:
-    _, abstract_model = create_nnx_abstract_model(config, mesh=mesh)
-    _, abstract_param, _ = nnx.split(abstract_model, nnx.Param, ...)
-  else:
-    quant = quantizations.configure_quantization(config)
-    model = Transformer(config, mesh=mesh, quant=quant)
-    abstract_param = maxtext_utils.get_abstract_param(model, config)
+  _, abstract_model = create_nnx_abstract_model(config, mesh=mesh)
+  _, abstract_param, _ = nnx.split(abstract_model, nnx.Param, ...)
 
   # Calculate and display the total parameter count based purely on abstract shapes.
   num_params = max_utils.calculate_num_params_from_pytree(abstract_param)
@@ -273,7 +261,7 @@ def inspect_maxtext(args, remaining_args):
     # "params.params.decoder.decoder_norm.scale" (for standard model weights)
     # "params.Tid2EidVar.decoder.layers_0.mlp.MoeBlock_0.tid2eid" (for legacy custom collections)
     key_str = ".".join(key_parts)
-    if config.enable_nnx and not key_str.startswith(("params", "Tid2EidVar")):
+    if not key_str.startswith(("params", "Tid2EidVar")):
       param_key = "params.params." + key_str
     else:
       param_key = "params." + key_str
@@ -392,7 +380,7 @@ def main():
     inspect_hf(args)
   elif args.mode == "maxtext":
     # remaining_args accepts maxtext config, like `model_name=<maxtext_model_name>
-    # scan_layers=<True | False> enable_nnx=<True | False>`
+    # scan_layers=<True | False>`
     inspect_maxtext(args, remaining_args)
   elif args.mode == "orbax":
     inspect_orbax(args)

src/maxtext/checkpoint_conversion/standalone_scripts/convert_gpt3_ckpt_from_paxml.py

@@ -35,24 +35,18 @@
 """
 
 import argparse
-import functools
 import gc
 import os
 import sys
 
 from flax import nnx
-from flax import serialization
 import jax
 from jax import random
 from jax.sharding import Mesh
 from maxtext.configs import pyconfig
 from maxtext.utils.globals import MAXTEXT_PKG_DIR
 from maxtext.common import checkpointing
-from maxtext.common.common_types import MODEL_MODE_TRAIN
-from maxtext.layers import quantizations
 from maxtext.common import train_state_nnx
-from maxtext.models.models import transformer_as_linen
-from maxtext.optimizers import optimizers
 from maxtext.utils import max_logging
 from maxtext.utils import max_utils
 from maxtext.utils import maxtext_utils
@@ -93,23 +87,15 @@ def convert(paxml_ckpt_path, maxtext_model_name, base_output_directory, run_name
   devices_array = maxtext_utils.create_device_mesh(cfg)
   mesh = Mesh(devices_array, cfg.mesh_axes)
 
-  if cfg.pure_nnx:
-    rngs = maxtext_utils_nnx.create_nnx_rngs(cfg, rng_key=init_rng)
-    model = model_creation_utils.from_config(cfg, mesh=mesh, rngs=rngs)
-    _, tx = train_utils.create_training_optimizer(cfg, model)
-    _create_model_partial, _ = model_creation_utils.create_nnx_abstract_model(cfg, mesh)
+  rngs = maxtext_utils_nnx.create_nnx_rngs(cfg, rng_key=init_rng)
+  model = model_creation_utils.from_config(cfg, mesh=mesh, rngs=rngs)
+  _, tx = train_utils.create_training_optimizer(cfg, model)
+  _create_model_partial, _ = model_creation_utils.create_nnx_abstract_model(cfg, mesh)
 
-    def init_state_fn():
-      nnx_model = _create_model_partial()
-      optimizer = nnx.Optimizer(nnx_model, tx, wrt=nnx.Param)
-      return train_state_nnx.TrainStateNNX(nnx_model, optimizer)
-
-  else:
-    quant = quantizations.configure_quantization(cfg)
-    model = transformer_as_linen(cfg, mesh, quant=quant, model_mode=MODEL_MODE_TRAIN)
-    learning_rate_schedule = maxtext_utils.create_learning_rate_schedule(cfg)
-    tx = optimizers.get_optimizer(cfg, learning_rate_schedule)
-    init_state_fn = functools.partial(maxtext_utils.init_initial_state, model, tx, cfg, True, init_rng)
+  def init_state_fn():
+    nnx_model = _create_model_partial()
+    optimizer = nnx.Optimizer(nnx_model, tx, wrt=nnx.Param)
+    return train_state_nnx.TrainStateNNX(nnx_model, optimizer)
 
   checkpoint_manager = checkpointing.create_orbax_checkpoint_manager(
       cfg.checkpoint_dir,
@@ -119,11 +105,8 @@ 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)
+  state = train_state_nnx.to_checkpoint_dict(state)
+  state.pop("nnx_aux", None)
 
   max_logging.log("start")
   max_utils.print_mem_stats("After params initialized")

src/maxtext/common/checkpointing.py

@@ -186,7 +186,7 @@ def _load_linen_checkpoint_into_nnx(
     enable_single_replica_ckpt_restoring: bool = False,
     config=None,
 ):
-  """Restores a Linen-layout checkpoint into an NNX state (pure_nnx resume).
+  """Restores a Linen-layout checkpoint into an NNX state.
 
   Restores a Linen-shape target that includes `nnx_aux`, then reshapes back via
   `_restored_linen_to_nnx`. rngs/dropout/batch stats come from `items/nnx_aux` when
@@ -230,7 +230,7 @@ def _restored_linen_to_nnx(restored_linen, abstract_nnx_state, config=None):
   """Reshapes a restored Linen-layout tree into the NNX state.
 
   Raises if the checkpoint is missing a weight. Every NNX restore path ends here: the load
-  itself is the Linen one, since pure_nnx reads and writes the Linen on-disk layout.
+  itself is the Linen one, since MaxText reads and writes the Linen on-disk layout.
   """
   _raise_on_weight_mismatch(*_expected_and_restored_params(abstract_nnx_state, restored_linen), config=config)
   return _linen_items_to_nnx(restored_linen, abstract_nnx_state)
@@ -309,7 +309,7 @@ def _load_full_state_from_path(
     The loaded state.
   """
   if source_checkpoint_layout == "orbax":
-    # pure_nnx checkpoints are stored in the Linen on-disk layout; reshape to NNX.
+    # Checkpoints are stored in the Linen on-disk layout; reshape to NNX.
     if isinstance(abstract_unboxed_pre_state, nnx.State):
       return _load_linen_checkpoint_into_nnx(
           path,
@@ -350,7 +350,7 @@ def combine_sharding(sds, shardings):
       sharded_abstract_state = jax.tree.map(combine_sharding, simple_abstract_state, shardings)
       pre_transformed_state = ocp.load(path, sharded_abstract_state)
     state = conversion_fn(pre_transformed_state)
-    # The conversion fn returns MaxText's on-disk (Linen) layout, which is what pure_nnx reads,
+    # The conversion fn returns MaxText's on-disk (Linen) layout, which is what NNX reads,
     # so NNX needs the same reshape as every other restore. An NNX state passes through.
     if isinstance(abstract_unboxed_pre_state, nnx.State) and not isinstance(state, nnx.State):
       state = _restored_linen_to_nnx(state, abstract_unboxed_pre_state, config=maxtext_config)
@@ -533,7 +533,7 @@ def load_state_if_possible(
     load_parameters_from_path = _normalize_checkpoint_root(load_parameters_from_path)
   if load_full_state_from_path:
     load_full_state_from_path = _normalize_checkpoint_root(load_full_state_from_path)
-  # pure_nnx saves in the Linen on-disk layout, so every branch below loads the same tree Linen
+  # Checkpoints are saved in the Linen on-disk layout, so every branch below loads the same tree Linen
   # does: the NNX abstract is converted to that layout going in, and what comes back is reshaped
   # into the NNX state on the way out.
   is_nnx = isinstance(abstract_unboxed_pre_state, (nnx.State, train_state_nnx.TrainStateNNX))
@@ -842,12 +842,8 @@ def maybe_save_checkpoint(checkpoint_manager, state, config, data_iterator, step
   if step is not None:
     actual_step = int(step)
   else:
-    if config.pure_nnx:
-      # 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
+    # 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
 
   # Determine if a checkpoint save should be forced, overriding the usual
   # `config.checkpoint_period` logic.
@@ -927,8 +923,8 @@ def save_checkpoint(checkpoint_manager, step, state, config=None, data_iterator=
 
   if config and getattr(config, "enable_diloco", False):
     state = diloco_checkpoint_utils.to_diloco_checkpoint_dict(state, config)
-  elif config and getattr(config, "pure_nnx", False):
-    # Save in the Linen on-disk layout so pure_nnx and Linen checkpoints are interchangeable.
+  elif config:
+    # Save in the Linen on-disk layout so NNX and Linen checkpoints are interchangeable.
     if isinstance(state, nnx.State):
       state = train_state_nnx.to_checkpoint_dict(state)
 

src/maxtext/common/emergency_checkpointing.py

@@ -65,7 +65,7 @@ def create_emergency_checkpoint_manager(
 
   persistent_p = gcs_utils.mkdir_and_check_permissions(persistent_checkpoint_dir)
 
-  # pure_nnx saves via to_checkpoint_dict (Linen params/opt_state/step plus an nnx_aux
+  # Checkpoints are saved via to_checkpoint_dict (Linen params/opt_state/step plus an nnx_aux
   # subtree), but the emergency manager restores against the abstract it is built with.
   # Convert it the same way so it matches what is on disk; restore reshapes back to NNX.
   if isinstance(abstract_state, nnx.State):

src/maxtext/common/train_state_nnx.py

@@ -60,7 +60,7 @@ def apply_gradients(self, grads: Any, **kwargs):
 
 # On-disk checkpoint format.
 #
-# A pure_nnx run saves in the same on-disk layout as a Linen run, so the two are
+# A run saves in the same on-disk layout that Linen used, so old and new checkpoints are
 # interchangeable. The NNX state pure dict differs from Linen's in three ways, all
 # reshaped below at save time:
 #   1. top-level keys: {model, optimizer:{step, opt_state}} -> {params:{params:...}, step, opt_state}
@@ -234,7 +234,7 @@ def to_checkpoint_dict(state: nnx.State | nnx.Module):
   """Reshapes an nnx.State into the on-disk checkpoint layout.
 
   Weights (nnx.Param) map to the Linen `params` collection and the optimizer to
-  opt_state/step, so pure_nnx and Linen checkpoints stay interchangeable. Everything else that
+  opt_state/step, so NNX and Linen checkpoints stay interchangeable. Everything else that
   must persist -- rngs/dropout, batch stats, and any custom variable -- goes under an `nnx_aux`
   subtree. Works on a concrete state (save) or an abstract state (restore target).
   """

src/maxtext/configs/base.yml

@@ -1339,11 +1339,6 @@ position_id_per_seconds: 25
 # Example: "8,8" to use a 8x8 subgrid (64 chips) of a full pod (16x16) of trillium.
 subslice_shape: ""
 
-# NNX
-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
 gdn_conv_kernel_dim: 4

src/maxtext/configs/inference/vllm.yml

@@ -15,8 +15,6 @@
 attention: "vllm_rpa"
 model_call_mode: "inference"
 
-# NNX required for vLLM integration
-enable_nnx: true
 # Avoid re-initializing JAX distributed system when using vLLM
 skip_jax_distributed_system: true
 # Scanned layers are not supported with vLLM integration

src/maxtext/configs/post_train/distillation-qwen3.yml

@@ -3,7 +3,6 @@
 # Inherit MaxText defaults
 base_config: "base.yml"
 override_model_config: True
-enable_nnx: True
 enable_checkpointing: True
 
 # --- Student Specifics ---

src/maxtext/configs/post_train/distillation_gpt_oss_20b.yml

@@ -19,7 +19,6 @@ distill_alpha: 0.5
 distill_temperature: 1.0
 distill_beta: 0
 distill_layer_indices: []
-enable_nnx: True
 load_balance_loss_weight: 0.001
 
 # Megablox grouped-matmul m-tile (batch_seq). The k/n dims already default to

src/maxtext/configs/post_train/distillation_qwen3_30b_base.yml

@@ -21,7 +21,6 @@ distill_alpha: 0.6
 distill_temperature: 1.0
 distill_beta: 1.0
 distill_layer_indices: [0,1,2,3,4,5,6,7]
-enable_nnx: True
 load_balance_loss_weight: 0.001
 
 ici_fsdp_parallelism: -1

src/maxtext/configs/post_train/distillation_qwen3_30b_base_pdbs8.yml

@@ -21,7 +21,6 @@ distill_alpha: 0.6
 distill_temperature: 1.0
 distill_beta: 1.0
 distill_layer_indices: [0,1,2,3,4,5,6,7]
-enable_nnx: True
 load_balance_loss_weight: 0.001
 
 ici_fsdp_parallelism: -1

src/maxtext/configs/pyconfig_deprecated.py

@@ -194,8 +194,7 @@ def validate_expert_shard_attention_option(expert_shard_attention_option: str) -
     )
 
 
-def validate_vocab_tiling(num_vocab_tiling: int, per_device_batch_size: int, max_target_length: int, enable_nnx: bool):
-  del enable_nnx  # NNX vocab tiling supported via vocab_tiling_nnx_loss in vocabulary_tiling.py
+def validate_vocab_tiling(num_vocab_tiling: int, per_device_batch_size: int, max_target_length: int):
   if (per_device_batch_size * max_target_length) % num_vocab_tiling != 0:
     raise ValueError("Per device batch size times sequence length should be divisible by the number of vocab tiles.")
 
@@ -239,9 +238,7 @@ def validate_keys(keys):
   validate_model_call_mode(keys["model_call_mode"])
   validate_prefill_and_target_lengths(keys["max_prefill_predict_length"], keys["max_target_length"])
   validate_rope_type(keys["rope_type"])
-  validate_vocab_tiling(
-      keys["num_vocab_tiling"], keys["per_device_batch_size"], keys["max_target_length"], keys["enable_nnx"]
-  )
+  validate_vocab_tiling(keys["num_vocab_tiling"], keys["per_device_batch_size"], keys["max_target_length"])
   if keys["enable_rampup_batch_size"]:
     validate_rampup_batch_size(
         keys["per_device_batch_size_start"],