[NNX] NNX migration (12/N): delete Linen code paths, classes, and NNX compatibility 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/standalone_scripts/convert_gpt3_ckpt_from_paxml.py

@@ -35,7 +35,6 @@
 """
 
 import argparse
-import functools
 import gc
 import os
 import sys
@@ -47,11 +46,7 @@
 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
@@ -92,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,
@@ -201,24 +188,18 @@ 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),
-    }
+  # 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),
+  }
 
   def get_layer_prefix(keystr_pax):
     # different path format between decoder_layer variable
@@ -231,26 +212,15 @@ 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"['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,
+    )
 
   def verify_fn(key_path, _):
     keystr = jax.tree_util.keystr(key_path)
@@ -302,7 +272,7 @@ 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
+  step_value = int(converted_state.optimizer.step.value)
   if checkpointing.save_checkpoint(checkpoint_manager, step_value, converted_state):
     max_logging.log(f"saved a checkpoint at step {step_value}")
   # Upon preemption, exit when and only when all ongoing saves are complete.

src/maxtext/common/checkpointing.py

@@ -222,7 +222,7 @@ def _load_linen_checkpoint_into_nnx(
     use_ocdbt,
     use_zarr3,
 ):
-  """Restores a Linen-layout checkpoint into an NNX state (pure_nnx resume).
+  """Restores a Linen-layout checkpoint into an NNX state.
 
   Restores against a Linen-shape abstract, reshapes back via
   `from_linen_checkpoint_dict`, then fills NNX-only rngs/dropout with defaults.
@@ -361,7 +361,7 @@ def combine_sharding(sds, shardings):
     else:
       raise ocp_v1.errors.InvalidLayoutError(f"Unknown checkpoint layout: {source_checkpoint_layout}")
   else:
-    # pure_nnx saves in the Linen on-disk layout; reshape it back into the NNX state.
+    # Checkpoints use the Linen on-disk layout; reshape it back into the NNX state.
     if isinstance(abstract_unboxed_pre_state, nnx.State):
       return _load_linen_checkpoint_into_nnx(
           path,
@@ -853,7 +853,7 @@ def map_to_pspec(data):
         )
         ocp.type_handlers.register_type_handler(jax.Array, array_handler, override=True)
 
-      # pure_nnx saves in the Linen on-disk layout; reshape it back into the NNX state.
+      # Checkpoints use the Linen on-disk layout; reshape it back into the NNX state.
       # (Emergency managers use their own restore path below.)
       if isinstance(abstract_unboxed_pre_state, nnx.State) and not isinstance(
           checkpoint_manager,
@@ -1093,26 +1093,21 @@ 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
 
   if checkpoint_manager.latest_step() == actual_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:
-      state = train_state_nnx.to_linen_checkpoint_dict(state.to_pure_dict())
+  # Save in the Linen on-disk layout so checkpoints stay interchangeable across formats.
+  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/common/train_state_nnx.py

@@ -58,10 +58,9 @@ 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
-# interchangeable. The NNX state pure dict differs from Linen's in three ways,
-# all
-# reshaped below at save time:
+# NNX saves in the same on-disk layout as the old Linen format, so checkpoints
+# stay interchangeable. The NNX state pure dict differs from that layout in three
+# ways, all reshaped below at save time:
 #   1. top-level keys: {model, optimizer:{step, opt_state}} -> {params:{params:...}, step, opt_state}
 #   2. weights: model/... -> params/params/... (Linen `params` collection)
 #   3. opt_state: int-keyed dict (empty entries skipped) -> list with None for EmptyState,

src/maxtext/configs/base.yml

@@ -1235,11 +1235,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

@@ -16,8 +16,6 @@ base_config: "base.yml"
 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_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

@@ -193,8 +193,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.")
 
@@ -238,9 +237,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"],