[NNX] Delete Linen (12.1a): state setup and train loop

src/maxtext/utils/maxtext_utils.py

@@ -98,10 +98,7 @@ def get_functional_train_with_signature(
   """Get the shardings (both state and data) for `train_step`."""
   functional_train = functools.partial(train_step, model, config, state_mesh_shardings, params_shardings)
   functional_train.__name__ = "train_step"  # pyrefly: ignore[missing-attribute]
-  if config.pure_nnx:
-    in_shardings = (state_mesh_shardings, data_sharding)  # State, batch
-  else:
-    in_shardings = (state_mesh_shardings, data_sharding, None)  # State, batch, rng
+  in_shardings = (state_mesh_shardings, data_sharding)  # State, batch
   out_shardings = (state_mesh_shardings, None)  # State, metrics
   static_argnums = ()  # We partial out the static argnums of model and config
   donate_argnums = 0  # This is the index of the state - we allow the compiler to make use of this memory.
@@ -112,10 +109,7 @@ def get_functional_eval_with_signature(eval_step, data_sharding, state_mesh_shar
   """Get the shardings (both state and data) for `eval_step`."""
   functional_eval = functools.partial(eval_step, model, config)
   functional_eval.__name__ = "eval_step"  # pyrefly: ignore[missing-attribute]
-  if config.pure_nnx:
-    in_shardings = (state_mesh_shardings, data_sharding)  # State, batch (NNX: no rng)
-  else:
-    in_shardings = (state_mesh_shardings, data_sharding, None)  # State, batch, rng
+  in_shardings = (state_mesh_shardings, data_sharding)  # State, batch
   out_shardings = None  # metrics
   static_argnums = ()  # We partial out the static argnums of model, config
   donate_argnums = ()  # state will be kept instead of being donated in eval_step
@@ -279,11 +273,7 @@ def get_train_input_output_trees(func, input_args, input_kwargs):
 
   serialized_compiled = load_serialized_compiled(config.compiled_trainstep_file)
   shaped_batch = get_shaped_batch(config)
-  if config.pure_nnx:
-    shaped_input_args = (state, shaped_batch)
-  else:
-    example_rng = jax.random.PRNGKey(0)
-    shaped_input_args = (state, shaped_batch, example_rng)
+  shaped_input_args = (state, shaped_batch)
   shaped_input_kwargs = {}
   in_tree, out_tree = get_train_input_output_trees(partial_train, shaped_input_args, shaped_input_kwargs)
   p_train_step = deserialize_and_load(serialized_compiled, in_tree, out_tree, execution_devices=execution_devices)
@@ -1879,51 +1869,8 @@ def get_logical_annotations(config, mesh, init_state_fn):
 
 
 def get_abstract_state(config, mesh, init_state_fn, is_training=True):
-  """Get a shaped abstraction of the state (including optimizer)"""
-  if config.pure_nnx:
-    return get_abstract_state_nnx(config, mesh, init_state_fn, is_training)
-
-  init_state_partial = init_state_fn
-
-  with nn_partitioning.axis_rules(config.logical_axis_rules):
-    abstract_state = jax.eval_shape(init_state_partial)
-
-  state_logical_annotations = nn.get_partition_spec(abstract_state)
-
-  state_mesh_shardings = nn.logical_to_mesh_sharding(state_logical_annotations, mesh, config.logical_axis_rules)
-  if is_training and config.shard_optimizer_over_data:
-    # Add data to sharding for optimizer state
-    state_mesh_shardings = state_mesh_shardings.replace(
-        opt_state=jax.tree.map_with_path(
-            functools.partial(sharding.add_data_to_sharding, mesh),
-            max_utils.unbox_logicallypartioned(abstract_state).opt_state,
-            state_mesh_shardings.opt_state,
-        )
-    )
-  if is_training and config.optimizer_memory_host_offload:
-    opt_state = jax.tree_util.tree_map(lambda x: x.with_memory_kind(kind="pinned_host"), state_mesh_shardings.opt_state)
-    state_mesh_shardings = state_mesh_shardings.replace(opt_state=opt_state)
-  if is_training and config.parameter_memory_host_offload:
-    assert config.param_scan_axis == 0, "You must set the scan axis 0 to enable parameter offloading."
-
-    def move(path, x):
-      max_logging.log(f"max_utils.py: Moving {path} to host")
-      return x.with_memory_kind(kind="pinned_host")
-
-    params = jax.tree_util.tree_map_with_path(move, state_mesh_shardings.params)
-    state_mesh_shardings = state_mesh_shardings.replace(params=params)
-
-  abstract_sharded_state = jax.jit(init_state_partial, in_shardings=None, out_shardings=state_mesh_shardings).eval_shape()
-
-  unboxed_abstract_sharded_state = max_utils.unbox_logicallypartioned(abstract_sharded_state)
-  # Initialization
-  with jax.set_mesh(mesh), nn_partitioning.axis_rules(config.logical_axis_rules):
-    state_mesh_annotations = nn.logical_to_mesh(state_logical_annotations)
-  return (
-      unboxed_abstract_sharded_state,
-      state_mesh_annotations,
-      state_mesh_shardings,
-  )
+  """Get a shaped abstraction of the state (including optimizer)."""
+  return get_abstract_state_nnx(config, mesh, init_state_fn, is_training)
 
 
 def get_abstract_state_nnx(config, mesh, nnx_init_trainstate_fn, is_training=True):

src/maxtext/utils/train_utils.py

@@ -20,7 +20,6 @@
 import optax
 import functools
 import orbax.checkpoint.pathways as ocp_pathways
-from functools import partial
 
 from flax import nnx
 from flax.linen import partitioning as nn_partitioning
@@ -252,32 +251,25 @@ def setup_train_loop(config, recorder, devices=None):
   from maxtext.input_pipeline.input_pipeline_interface import create_data_iterator
 
   with maybe_record_goodput(recorder, GoodputEvent.TPU_INIT):
-    is_training = True
     init_rng = jax.random.PRNGKey(config.init_weights_seed)
     mesh = maxtext_utils.get_mesh_from_config(config, devices)
     context_parallel_size = mesh.shape.get(config.context_sharding, 1)
-    if config.pure_nnx:
-      # Create abstract NNX model.
-      _create_model_partial, model = model_creation_utils.create_nnx_abstract_model(config, mesh, devices)
-    else:
-      model = model_creation_utils.from_config(config, devices)
+    # Create abstract NNX model.
+    _create_model_partial, model = model_creation_utils.create_nnx_abstract_model(config, mesh, devices)
     learning_rate_schedule, tx = create_training_optimizer(config, model)
 
-    if config.pure_nnx:
-      # For NNX, the train state is wrapped in the TrainStateNNX module.
-      def create_train_state_fn():
-        model = _create_model_partial()
-        wrt = (
-            getattr(nnx, "LoRAParam", nnx.Param)
-            if getattr(getattr(config, "lora", None), "enable_lora", False)
-            else nnx.Param
-        )
-        optimizer = nnx.Optimizer(model, tx, wrt=wrt)
-        return train_state_nnx.TrainStateNNX(model, optimizer)
-
-      init_state_fn = create_train_state_fn
-    else:
-      init_state_fn = partial(maxtext_utils.init_initial_state, model, tx, config, is_training, init_rng)
+    # The train state is wrapped in the TrainStateNNX module.
+    def create_train_state_fn():
+      model = _create_model_partial()
+      wrt = (
+          getattr(nnx, "LoRAParam", nnx.Param)
+          if getattr(getattr(config, "lora", None), "enable_lora", False)
+          else nnx.Param
+      )
+      optimizer = nnx.Optimizer(model, tx, wrt=wrt)
+      return train_state_nnx.TrainStateNNX(model, optimizer)
+
+    init_state_fn = create_train_state_fn
     checkpoint_manager = create_checkpoint_manager(config, mesh, init_state_fn)
     if checkpoint_manager is not None:
       checkpoint_step = checkpointing.latest_step(checkpoint_manager)
@@ -340,38 +332,34 @@ def create_train_state_fn():
     state, _, state_mesh_shardings, data_iterator, _ = maxtext_utils.setup_training_state(
         data_iterator, config, mesh, checkpoint_manager, init_state_fn
     )
-    if config.pure_nnx:
-      if getattr(getattr(config, "lora", None), "enable_lora", False) and getattr(config.lora, "lora_restore_path", None):
-        # Restore standalone LoRA adapter weights onto the base model state after initialization.
-        target_model_state = (
-            state["model"]
-            if (isinstance(state, (nnx.State, dict)) and "model" in state)
-            else getattr(state, "model", state)
-        )
-        # pyrefly: ignore[bad-argument-type]
-        lora_utils.restore_lora_from_path(target_model_state, config)
-        _, _, state_mesh_shardings = maxtext_utils.get_abstract_state_nnx(config, mesh, init_state_fn, True)
-      with nn_partitioning.axis_rules(config.logical_axis_rules):
-        # We only need the graphdef here; it's merged with state below. Avoid
-        # nnx.get_abstract_model: it eagerly builds a NamedSharding for every variable
-        # under jax.set_mesh(mesh) and rejects any logical name missing from
-        # logical_axis_rules (e.g. concat_embed on the MTP kernel). Tracing shapes
-        # without a mesh skips sharding resolution, so it avoids the crash.
-        state_graphdef = nnx.graphdef(nnx.eval_shape(init_state_fn))
+    if getattr(getattr(config, "lora", None), "enable_lora", False) and getattr(config.lora, "lora_restore_path", None):
+      # Restore standalone LoRA adapter weights onto the base model state after initialization.
+      target_model_state = (
+          state["model"]
+          if (isinstance(state, (nnx.State, dict)) and "model" in state)
+          else getattr(state, "model", state)
+      )
+      # pyrefly: ignore[bad-argument-type]
+      lora_utils.restore_lora_from_path(target_model_state, config)
+      _, _, state_mesh_shardings = maxtext_utils.get_abstract_state_nnx(config, mesh, init_state_fn, True)
+    with nn_partitioning.axis_rules(config.logical_axis_rules):
+      # We only need the graphdef here; it's merged with state below. Avoid
+      # nnx.get_abstract_model: it eagerly builds a NamedSharding for every variable
+      # under jax.set_mesh(mesh) and rejects any logical name missing from
+      # logical_axis_rules (e.g. concat_embed on the MTP kernel). Tracing shapes
+      # without a mesh skips sharding resolution, so it avoids the crash.
+      state_graphdef = nnx.graphdef(nnx.eval_shape(init_state_fn))
 
     if isinstance(state, diloco.DiLoCoTrainState):
       state_params = state.params
       if hasattr(state_mesh_shardings, "model"):
         _, state_mesh_shardings_params, _ = nnx.split(state_mesh_shardings.model, nnx.Param, ...)
       else:
         state_mesh_shardings_params = state_mesh_shardings.params
-    elif config.pure_nnx:
+    else:
       with nn_partitioning.axis_rules(config.logical_axis_rules):
         _, state_params, _ = nnx.split(state.model, nnx.Param, ...)
         _, state_mesh_shardings_params, _ = nnx.split(state_mesh_shardings.model, nnx.Param, ...)
-    else:
-      state_params = state.params
-      state_mesh_shardings_params = state_mesh_shardings.params
 
     if config.enable_diloco:
       with jax.set_mesh(mesh), nn_partitioning.axis_rules(config.logical_axis_rules):
@@ -409,28 +397,21 @@ def create_train_state_fn():
 
     # print weights sharding info under debug sharding mode
     if config.debug_sharding:
-      if config.pure_nnx:
-        # TODO: Study how to get logical annotations of NNX module. Because of eager sharding, we
-        # probably already lost the logical partition info at this moment.
-        logical_annotations_params = None
-      else:
-        logical_annotations = maxtext_utils.get_logical_annotations(config, mesh, init_state_fn)
-        logical_annotations_params = logical_annotations.params
+      # TODO: Study how to get logical annotations of NNX module. Because of eager sharding, we
+      # probably already lost the logical partition info at this moment.
+      logical_annotations_params = None
 
       max_utils.print_non_trivial_mesh_axis(model.mesh)  # pyrefly: ignore[missing-attribute]
       maxtext_utils.print_shardings_params(state_params, state_mesh_shardings_params, mesh, logical_annotations_params)
 
-  if config.pure_nnx:
-    if config.enable_diloco:
-      # Don't merge the DiLoCoTrainState into the plain-model graphdef. The inner
-      # train step needs that graphdef as jit_model; the wrapper passes through as state.
-      train_state = state
-      model = state_graphdef  # pyrefly: ignore[unbound-name]
-    else:
-      train_state = nnx.merge(state_graphdef, state)  # pyrefly: ignore[unbound-name]
-      model = train_state.model
-  else:
+  if config.enable_diloco:
+    # Don't merge the DiLoCoTrainState into the plain-model graphdef. The inner
+    # train step needs that graphdef as jit_model; the wrapper passes through as state.
     train_state = state
+    model = state_graphdef  # pyrefly: ignore[unbound-name]
+  else:
+    train_state = nnx.merge(state_graphdef, state)  # pyrefly: ignore[unbound-name]
+    model = train_state.model
 
   return (
       init_rng,

tests/integration/setup_train_loop_nnx_test.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""Integration test for setup_train_loop with pure_nnx=True.
+"""Integration test for setup_train_loop on the NNX path.
 
 setup_train_loop wires together create_nnx_abstract_model, the training
 optimizer,
@@ -43,7 +43,6 @@ def _tiny_nnx_pyconfig(**overrides):
       "enable_checkpointing": False,
       "dataset_type": "synthetic",
       "model_name": "default",
-      "pure_nnx": True,
       "per_device_batch_size": 1.0,
       "base_emb_dim": 8,
       "base_num_query_heads": 4,
@@ -68,7 +67,7 @@ def _tiny_nnx_pyconfig(**overrides):
 class SetupTrainLoopNNXIntegrationTest(unittest.TestCase):
   """End-to-end check that setup_train_loop returns a usable TrainStateNNX."""
 
-  def test_pure_nnx_setup_returns_train_state_nnx(self):
+  def test_setup_returns_train_state_nnx(self):
     config = _tiny_nnx_pyconfig()
 
     (
@@ -126,7 +125,7 @@ def test_load_balanced_cp_keeps_checkpoint_iterator_unwrapped(self):
     self.assertNotIsInstance(data_iterator, train_utils._ReorderedDataIterator)
     self.assertIsInstance(eval_data_iterator, train_utils._ReorderedDataIterator)
 
-  def test_pure_nnx_setup_param_only_split_matches_model(self):
+  def test_setup_param_only_split_matches_model(self):
     """nnx.split(state.model, nnx.Param, ...) must yield a non-empty Param tree
 
     whose structure matches state_mesh_shardings.model after the same split.

tests/unit/correctness_tests_nnx_dispatch_test.py

@@ -18,7 +18,7 @@
 exercises the changed dispatch code that otherwise has no CPU coverage:
   - `mt.from_config` is exported (the GRPO trainer calls it)
   - the SFT correctness test's `setup_maxtext_model` / `get_maxtext_logits` run on
-    both paths (pure_nnx=True -> NNX, pure_nnx=False -> Linen) and stay finite
+    the NNX path and stay finite
 
 The GRPO NNX building blocks the other dispatch helpers call
 (`compute_log_probs_nnx`, `grpo_loss_fn_nnx`) are already covered by grpo_nnx_test.
@@ -54,15 +54,12 @@
 }
 
 
-def _sft_config(pure_nnx):
+def _sft_config():
   return pyconfig.initialize(
       [sys.argv[0], os.path.join(MAXTEXT_PKG_DIR, "configs/post_train", "sft.yml")],
-      run_name=f"unit-sft-{pure_nnx}",
+      run_name="unit-sft-nnx",
       model_name="default",
       enable_checkpointing=False,
-      pure_nnx=pure_nnx,
-      enable_nnx=pure_nnx,
-      pure_nnx_decoder=pure_nnx,
       **_SMALL,
   )
 
@@ -84,12 +81,7 @@ def test_from_config_is_exported(self):
     self.assertTrue(hasattr(mt, "from_config"))
 
   def test_sft_logits_nnx_path(self):
-    config = _sft_config(pure_nnx=True)
-    logits = sft.get_maxtext_logits(config, _fake_data(config))
-    self.assertTrue(bool(jnp.isfinite(logits).all()))
-
-  def test_sft_logits_linen_path(self):
-    config = _sft_config(pure_nnx=False)
+    config = _sft_config()
     logits = sft.get_maxtext_logits(config, _fake_data(config))
     self.assertTrue(bool(jnp.isfinite(logits).all()))