[NNX] Delete Linen (12.1b): collapse dispatch in sharding and muon

src/maxtext/utils/maxtext_utils.py

@@ -96,10 +96,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.
@@ -110,10 +107,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
@@ -254,11 +248,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)
@@ -1794,51 +1784,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/muon_utils.py

@@ -33,8 +33,6 @@
 import jax
 from maxtext.configs import pyconfig
 from maxtext.utils.globals import MAXTEXT_PKG_DIR
-from maxtext.layers import quantizations
-from maxtext.models import models
 from maxtext.utils import maxtext_utils, model_creation_utils
 from optax.contrib._muon import MuonDimensionNumbers as mdn
 
@@ -134,26 +132,17 @@ def get_transform_tree(tree, path=()):
 def get_muon_weight_dimension_numbers(model, config, verbose=False):
   """Extract muon dimension number from model structure."""
 
-  if isinstance(model, nnx.Module):
-    _, abstract_param, _ = nnx.split(model, nnx.Param, ...)
+  _, abstract_param, _ = nnx.split(model, nnx.Param, ...)
 
-    def apply_transform_nnx(path: Tuple[jax.tree_util.KeyEntry, ...], leaf):
-      # Convert jax.tree_util.KeyEntry path to Tuple[str, ...]
-      path_strings = tuple(p.key for p in path if isinstance(p, jax.tree_util.DictKey))
-      return transform_logic(path_strings)
+  def apply_transform_nnx(path: Tuple[jax.tree_util.KeyEntry, ...], leaf):
+    # Convert jax.tree_util.KeyEntry path to Tuple[str, ...]
+    path_strings = tuple(p.key for p in path if isinstance(p, jax.tree_util.DictKey))
+    return transform_logic(path_strings)
 
-    # NNX abstract_param is an nnx.State (not Linen's dict of LogicallyPartitioned leaves);
-    # tree_map_with_path round-trips that structure so each Param.value holds the mdn result.
-    muon_weight_dimension_numbers = jax.tree_util.tree_map_with_path(
-        apply_transform_nnx, nnx.to_pure_dict(abstract_param)
-    )
-    muon_weight_dimension_numbers = nnx.State(muon_weight_dimension_numbers)
-
-  else:  # Linen
-    # quickly get param structure without materialization
-    abstract_param = maxtext_utils.get_abstract_param(model, config)
-    # get muon dimension number from param
-    muon_weight_dimension_numbers = get_transform_tree(abstract_param)
+  # tree_map_with_path handles NNX's PyTree structure; result is an nnx.State with the
+  # same structure, where each Param's value holds the mdn result.
+  muon_weight_dimension_numbers = jax.tree_util.tree_map_with_path(apply_transform_nnx, nnx.to_pure_dict(abstract_param))
+  muon_weight_dimension_numbers = nnx.State(muon_weight_dimension_numbers)
 
   if verbose:
     _print_structure_debug(abstract_param, muon_weight_dimension_numbers)
@@ -185,7 +174,7 @@ def get_leaf_info(leaf):
   print("\nIs this reasonable?")
 
 
-def get_model_mdn(model_name, scan_layers=True, verbose=False, pure_nnx=False):
+def get_model_mdn(model_name, scan_layers=True, verbose=False):
   """Initializes a model and retrieves its Muon dimension numbers.
 
   This function sets up the configuration for a given model, initializes the
@@ -209,30 +198,16 @@ def get_model_mdn(model_name, scan_layers=True, verbose=False, pure_nnx=False):
       f"model_name={model_name}",
       f"scan_layers={scan_layers}",
       "attention=dot_product",
-      f"pure_nnx={pure_nnx}",
       "skip_jax_distributed_system=True",
   ]
-  if not pure_nnx:
-    argv.extend(
-        [
-            "enable_nnx=False",
-            "pure_nnx_decoder=False",
-        ]
-    )
   config = pyconfig.initialize(argv)
   # Setup model
   devices_array = maxtext_utils.create_device_mesh(config)
   mesh = jax.sharding.Mesh(devices_array, config.mesh_axes)
-  quant = quantizations.configure_quantization(config)
-  if pure_nnx:
-    _, model = model_creation_utils.create_nnx_abstract_model(config, mesh)
-  else:
-    model = models.transformer_as_linen(config, mesh=mesh, quant=quant)
+  _, model = model_creation_utils.create_nnx_abstract_model(config, mesh)
   # Get dimension number
   muon_weight_dimension_numbers = get_muon_weight_dimension_numbers(model, config, verbose=verbose)
-  if pure_nnx:
-    muon_weight_dimension_numbers = {"params": nnx.to_pure_dict(muon_weight_dimension_numbers)}
-  return muon_weight_dimension_numbers
+  return {"params": nnx.to_pure_dict(muon_weight_dimension_numbers)}
 
 
 if __name__ == "__main__":
@@ -241,4 +216,4 @@ def get_model_mdn(model_name, scan_layers=True, verbose=False, pure_nnx=False):
     sys.exit(1)
   model_name_arg = sys.argv[1]
   scan_layers_arg = sys.argv[2].lower() == "true"
-  get_model_mdn(model_name_arg, scan_layers_arg, verbose=True, pure_nnx=False)
+  get_model_mdn(model_name_arg, scan_layers_arg, verbose=True)

src/maxtext/utils/sharding.py

@@ -557,26 +557,7 @@ def maybe_update_params_sharding_with_opt(config, state_mesh_shardings):
       - updated_state_mesh_shardings: State mesh shardings with updated params field
         (unchanged if shard_optimizer_over_data is False)
   """
-  if config.pure_nnx:
-    return maybe_update_params_sharding_with_opt_nnx(config, state_mesh_shardings)
-  prev_params_shardings = state_mesh_shardings.params
-  if config.shard_optimizer_over_data:
-    if isinstance(state_mesh_shardings.opt_state, optax.ScaleByAdamState):
-      sharded_fp32_params = state_mesh_shardings.opt_state.mu
-    elif isinstance(state_mesh_shardings.opt_state, tuple) and isinstance(
-        state_mesh_shardings.opt_state[0], optax.ScaleByAdamState
-    ):
-      sharded_fp32_params = state_mesh_shardings.opt_state[0].mu
-    else:
-      raise NotImplementedError(f"Could not find optimizer state shardings from {type(state_mesh_shardings.opt_state)}")
-    if "params" not in sharded_fp32_params.keys():
-      # When quantization=fp8 is enabled the sharded_fp32_params
-      # are not wrapped in `params`. Here we wrap them back.
-      sharded_fp32_params = {"params": sharded_fp32_params}
-    state_mesh_shardings = state_mesh_shardings.replace(
-        params=dict(prev_params_shardings, **sharded_fp32_params)
-    )  # pyrefly: ignore[bad-unpacking]
-  return prev_params_shardings, state_mesh_shardings
+  return maybe_update_params_sharding_with_opt_nnx(config, state_mesh_shardings)
 
 
 def maybe_update_params_sharding_with_opt_nnx(
@@ -708,8 +689,6 @@ def build_zero1_input_state_mesh_shardings(config, state_mesh_shardings, params_
   """
   if not config.shard_optimizer_over_data:
     return state_mesh_shardings
-  if not config.pure_nnx:
-    return state_mesh_shardings.replace(params=params_shardings)
   # nnx.State has no .replace: shallow-copy via tree_map (preserves nested container
   # types) and overlay params_shardings under input_state.model.
   input_state = jax.tree_util.tree_map(

src/maxtext/utils/train_utils.py

@@ -19,7 +19,6 @@
 import jax
 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
@@ -228,27 +227,20 @@ 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()
-        optimizer = nnx.Optimizer(model, tx, wrt=nnx.Param)
-        return train_state_nnx.TrainStateNNX(model, optimizer)
+    # The train state is wrapped in the TrainStateNNX module.
+    def create_train_state_fn():
+      model = _create_model_partial()
+      optimizer = nnx.Optimizer(model, tx, wrt=nnx.Param)
+      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)
+    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)
@@ -308,38 +300,30 @@ 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:
-      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))
-        _, 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
+    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))
+      _, state_params, _ = nnx.split(state.model, nnx.Param, ...)
+      _, state_mesh_shardings_params, _ = nnx.split(state_mesh_shardings.model, nnx.Param, ...)
 
     if config.enable_diloco:
       with jax.set_mesh(mesh), nn_partitioning.axis_rules(config.logical_axis_rules):
         state, outer_opt_state_sharding = diloco.build_diloco_state(config, lambda: state, mesh=mesh)
 
         # create state_mesh_shardings for the DilocoState
-        step_mesh = state_mesh_shardings.optimizer.step.mesh if config.pure_nnx else state_mesh_shardings.step.mesh
+        step_mesh = state_mesh_shardings.optimizer.step.mesh
         inner_state_shardings = diloco.add_diloco_to_sharding(state_mesh_shardings)
         state_mesh_shardings = diloco.DiLoCoTrainState(
             inner_state_shardings,
             # 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()  # pyrefly: ignore[missing-attribute]
-            if config.pure_nnx
-            else state_mesh_shardings_params,
+            state_mesh_shardings_params.to_pure_dict(),
             outer_opt_state_sharding,
-            jax.sharding.NamedSharding(  # pyrefly: ignore[bad-argument-type]
-                mesh=step_mesh, spec=jax.sharding.PartitionSpec()
-            ),
+            jax.sharding.NamedSharding(mesh=step_mesh, spec=jax.sharding.PartitionSpec()),
         )
 
     # TODO(aireenmei, hengtaoguo): support sharding in vit for multimodal
@@ -349,28 +333,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
+  else:
+    train_state = nnx.merge(state_graphdef, state)
+    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()
 
     (
@@ -109,7 +108,7 @@ def test_pure_nnx_setup_returns_train_state_nnx(self):
     # flag them as unused — they're part of the public return contract.
     del checkpoint_manager, rampup_manager, eval_data_iterator
 
-  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.