[NNX] Delete Linen 3/5: collapse dispatch in inference (maxengine, kvcache, vLLM, LoRA)

src/maxtext/inference/vllm_decode.py

@@ -114,8 +114,6 @@ def decode_with_vllm(config: Config) -> None:
               "debug_sharding": config.debug_sharding,
               "prefuse_moe_weights": config.prefuse_moe_weights,
               "scan_layers": config.scan_layers,
-              "enable_nnx": config.enable_nnx,
-              "pure_nnx_decoder": config.pure_nnx_decoder,
           },
           "sharding": {
               "sharding_strategy": {
@@ -271,8 +269,6 @@ def decode_with_tunix(
           "debug_sharding": config.debug_sharding,
           "prefuse_moe_weights": config.prefuse_moe_weights,
           "scan_layers": config.scan_layers,
-          "enable_nnx": config.enable_nnx,
-          "pure_nnx_decoder": config.pure_nnx_decoder,
       }
   }
 

src/maxtext/utils/lora_utils.py

@@ -15,7 +15,6 @@
 """Common LoRA utils needed to support LoRA adapters."""
 
 from collections.abc import Mapping
-from functools import partial
 import json
 import os
 import re
@@ -171,10 +170,7 @@ def load_adapter(config, base_abstract_state_params, adapter_config_path, adapte
     if not gcs_utils.gcs_path_exists(f"{adapter_weights_path}/commit_success.txt"):
       raise FileNotFoundError(f"Failed to read lora_weights from {adapter_weights_path}.")
 
-    if config.pure_nnx:
-      lora_state, _ = get_lora_abstract_state_nnx(base_abstract_state_params, lora_config)
-    else:
-      lora_state, _ = get_lora_abstract_state(base_abstract_state_params, lora_config)
+    lora_state, _ = get_lora_abstract_state_nnx(base_abstract_state_params, lora_config)
 
     with nn_partitioning.axis_rules(config.logical_axis_rules):
       lora_params = checkpointing.load_params_from_path(
@@ -218,37 +214,31 @@ def setup_initial_lora_state(model, data_iterator, tx, config, rng, mesh, checkp
 
   if lora_adapter_path:
     max_logging.log(f"Setting initial state of LoRA with lora_adapter_path = {lora_adapter_path}")
-    if config.pure_nnx:
-      # pylint: disable=import-outside-toplevel
-      from maxtext.common import train_state_nnx
-      from maxtext.utils import model_creation_utils
-
-      _create_model_partial, _ = model_creation_utils.create_nnx_abstract_model(config, mesh)
-
-      def create_train_state_fn():
-        nnx_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(nnx_model, tx, wrt=wrt)
-        return train_state_nnx.TrainStateNNX(nnx_model, optimizer)
+    # pylint: disable=import-outside-toplevel
+    from maxtext.common import train_state_nnx
+    from maxtext.utils import model_creation_utils
+
+    _create_model_partial, _ = model_creation_utils.create_nnx_abstract_model(config, mesh)
+
+    def create_train_state_fn():
+      nnx_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(nnx_model, tx, wrt=wrt)
+      return train_state_nnx.TrainStateNNX(nnx_model, optimizer)
 
-      init_state_fn = create_train_state_fn
-    else:
-      init_state_fn = partial(maxtext_utils.init_initial_state, model, tx, config, True, rng)
+    init_state_fn = create_train_state_fn
     unboxed_abstract_state, _, _ = maxtext_utils.get_abstract_state(config, mesh, init_state_fn, True)
 
     lora_config_path = lora_adapter_path + "adapter_config.json"
 
     lora_config = gcs_utils.read_json_from_gcs(lora_config_path)
 
-    if config.pure_nnx:
-      base_abstract_params = _nnx_param_subtree(unboxed_abstract_state)
-      lora_state, lora_state_annotations = get_lora_abstract_state_nnx(base_abstract_params, lora_config)
-    else:
-      lora_state, lora_state_annotations = get_lora_abstract_state(unboxed_abstract_state.params, lora_config)
+    base_abstract_params = _nnx_param_subtree(unboxed_abstract_state)
+    lora_state, lora_state_annotations = get_lora_abstract_state_nnx(base_abstract_params, lora_config)
 
     lora_weights_path = f"{lora_adapter_path}/0/items"
 

tests/integration/maxengine_test.py

@@ -96,8 +96,8 @@ def test_stack_and_unstack_prefill_cache_nnx(self):
   # default. test_basic_prefill_nnx / test_basic_decode_nnx below cover the NNX path.
 
   def _init_nnx_pyconfig(self, **kwargs):
-    """Same as init_pyconfig but with the NNX flags turned on."""
-    return self.init_pyconfig(pure_nnx=True, enable_nnx=True, pure_nnx_decoder=True, **kwargs)
+    """Same as init_pyconfig (NNX is the only path now)."""
+    return self.init_pyconfig(**kwargs)
 
   def _build_nnx_params(self, cfg, mesh):
     """Materialize an NNX Transformer and return its nnx.Param state."""
@@ -173,7 +173,7 @@ def test_basic_decode_nnx(self):
       )
 
   def test_quantize_passes_gate_for_nnx(self):
-    """pure_nnx + quantization (convert-on-load) reaches the actual machinery in train mode."""
+    """NNX + quantization (convert-on-load) reaches the actual machinery in train mode."""
     # checkpoint_is_quantized defaults to False — full-precision on disk, AQT
     # quantizes per-forward against the loaded kernel (train mode).
     cfg = self._init_nnx_pyconfig(quantization="int8")
@@ -187,7 +187,7 @@ def test_quantize_passes_gate_for_nnx(self):
       pass  # any other failure (e.g. checkpoint not found) is fine for this test
 
   def test_load_pre_quantized_nnx_passes_quant_gate(self):
-    """pure_nnx + quantization + checkpoint_is_quantized=True clears the load gate."""
+    """NNX + quantization + checkpoint_is_quantized=True clears the load gate."""
     cfg = self._init_nnx_pyconfig(quantization="int8", checkpoint_is_quantized=True)
     engine = maxengine.MaxEngine(cfg, jax.devices())
     self.assertEqual(engine._nnx_quant_mode_str, "serve")  # pylint: disable=protected-access
@@ -219,7 +219,7 @@ def test_quantized_prefill_nnx_train_mode(self):
     self.assertTrue(jnp.all(jnp.isfinite(prefill_result["logits"])))
 
   def test_lora_load_single_adapter_reaches_loader_on_nnx(self):
-    """pure_nnx + LoRA: load_single_adapter dispatches to the NNX loader.
+    """NNX + LoRA: load_single_adapter dispatches to the NNX loader.
 
     A nonexistent adapter path should raise FileNotFoundError from the
     loader itself. A NotImplementedError here would mean the dispatch

tests/post_training/unit/lora_utils_test.py

@@ -55,8 +55,6 @@
     "base_mlp_dim": 256,
     "max_prefill_predict_length": 4,
     "model_name": "llama2-7b",
-    "enable_nnx": True,
-    "pure_nnx_decoder": True,
     "override_model_config": True,
     "weight_dtype": "bfloat16",
 }

tests/unit/aqt_serve_roundtrip_nnx_test.py

@@ -74,9 +74,6 @@ def _init_cfg(self, ckpt_path, *, checkpoint_is_quantized):
         sys.argv[0],
         base_yml,
         "model_name=gpt3-52k",
-        "pure_nnx=true",
-        "enable_nnx=true",
-        "pure_nnx_decoder=true",
         "max_target_length=64",
         "max_prefill_predict_length=16",
         "per_device_batch_size=1",

tests/unit/maxengine_nnx_test.py

@@ -41,9 +41,6 @@ def _nnx_config(self, **kwargs):
         "max_target_length": 8,
         "per_device_batch_size": 1,
         "enable_checkpointing": False,
-        "pure_nnx": True,
-        "enable_nnx": True,
-        "pure_nnx_decoder": True,
     } | kwargs
     return pyconfig.initialize([sys.argv[0], get_test_config_path()], **init_kwargs)