[NNX] Delete Linen (12.1c): collapse dispatch in quantization and model creation

src/maxtext/layers/quantizations.py

@@ -938,32 +938,28 @@ def maybe_quantize_model(model, config):
   if config.quantization and config.use_qwix_quantization and not config.use_batch_split_schedule:
     quantization_provider = get_qt_provider(config)
     if quantization_provider:
-      if config.pure_nnx:
-        input_shape = (config.micro_batch_size_to_train_on, config.max_target_length)
-        dummy_tokens = jnp.ones(input_shape, dtype=jnp.int32)
-        dummy_positions = jnp.ones(input_shape, dtype=jnp.int32)
-        dummy_segment_ids = jnp.ones(input_shape, dtype=jnp.int32)
-        # The MTP block reads the decoder targets, so the qwix forward pass needs them.
-        # The Linen path supplies them from the is_initializing() guard in Transformer.
-        dummy_targets = {}
-        if config.mtp_num_layers > 0:
-          dummy_targets["decoder_target_tokens"] = jnp.ones(input_shape, dtype=jnp.int32)
-          dummy_targets["decoder_target_mask"] = jnp.ones(input_shape, dtype=jnp.int32)
-        model = qwix.quantize_model(
-            model,
-            quantization_provider,
-            dummy_tokens,
-            dummy_positions,
-            dummy_segment_ids,
-            enable_dropout=False,
-            **dummy_targets,
-        )
-        # Qwix quantization runs a forward pass during tracing, which sows transient nnx.Intermediate variables
-        # (e.g. max_logits from QK-Clip, MTP losses) into the model. Popping them here prevents structural mismatches
-        # between the initial setup GraphDef/state_mesh_shardings and the stripped states during train steps.
-        nnx.pop(model, nnx.Intermediate)
-      else:
-        model = qwix.quantize_model(model, quantization_provider)
+      input_shape = (config.micro_batch_size_to_train_on, config.max_target_length)
+      dummy_tokens = jnp.ones(input_shape, dtype=jnp.int32)
+      dummy_positions = jnp.ones(input_shape, dtype=jnp.int32)
+      dummy_segment_ids = jnp.ones(input_shape, dtype=jnp.int32)
+      # The MTP block reads the decoder targets, so the qwix forward pass needs them.
+      dummy_targets = {}
+      if config.mtp_num_layers > 0:
+        dummy_targets["decoder_target_tokens"] = jnp.ones(input_shape, dtype=jnp.int32)
+        dummy_targets["decoder_target_mask"] = jnp.ones(input_shape, dtype=jnp.int32)
+      model = qwix.quantize_model(
+          model,
+          quantization_provider,
+          dummy_tokens,
+          dummy_positions,
+          dummy_segment_ids,
+          enable_dropout=False,
+          **dummy_targets,
+      )
+      # Qwix quantization runs a forward pass during tracing, which sows transient nnx.Intermediate variables
+      # (e.g. max_logits from QK-Clip, MTP losses) into the model. Popping them here prevents structural mismatches
+      # between the initial setup GraphDef/state_mesh_shardings and the stripped states during train steps.
+      nnx.pop(model, nnx.Intermediate)
       for _, val in nnx.graph.iter_graph(model):
         if hasattr(val, "__dict__") and "qwix_rngs" in val.__dict__:
           del val.qwix_rngs

src/maxtext/utils/model_creation_utils.py

@@ -928,11 +928,8 @@ def from_pretrained(
   _, _abs_state_for_specs = nnx.split(abstract_model)
   specs = nnx.get_partition_spec(_abs_state_for_specs)
 
-  if config.pure_nnx:
-    model = maxtext_utils_nnx.create_nnx_sharded_model(abstract_model, _create_model, mesh=mesh)
-    # TODO: print debug_sharding info
-  else:
-    model = create_nnx_sharded_model_hybrid(config, mesh, devices, model_mode, rng_key)
+  model = maxtext_utils_nnx.create_nnx_sharded_model(abstract_model, _create_model, mesh=mesh)
+  # TODO: print debug_sharding info
 
   sharded_state = nnx.state(model)
 

tests/unit/nnx_quant_guard_test.py

@@ -12,43 +12,18 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""qwix + NNX coverage: the config guard and the ToNNX->Linen bridge.
+"""qwix + NNX coverage for the ToNNX->Linen bridge.
 
-- Config guard: qwix quantization under pure_nnx requires the pure NNX decoder.
-  The bridged Linen decoder (pure_nnx_decoder=False) is invisible to qwix, so
-  quantization/sparsity would silently no-op; validation must reject that combo.
-- Bridge: nnx_attrs_to_linen_vars must skip qwix's non-Variable bookkeeping attrs
-  (qwix_path/qwix_rngs/disable_quant_stats_update) instead of raising.
+nnx_attrs_to_linen_vars must skip qwix's non-Variable bookkeeping attrs
+(qwix_path/qwix_rngs/disable_quant_stats_update) instead of raising.
 """
 
-import sys
 import unittest
 
 import jax.numpy as jnp
 from flax import nnx
 
-from maxtext.configs import pyconfig
 from maxtext.layers import nnx_wrappers
-from tests.utils.test_helpers import get_test_config_path
-
-
-class QwixNnxQuantGuardTest(unittest.TestCase):
-
-  def _init(self, **overrides):
-    overrides.setdefault("enable_checkpointing", False)
-    return pyconfig.initialize([sys.argv[0], get_test_config_path()], **overrides)
-
-  def test_bridged_decoder_with_qwix_quant_raises(self):
-    with self.assertRaisesRegex(Exception, "pure_nnx_decoder"):
-      self._init(pure_nnx=True, pure_nnx_decoder=False, use_qwix_quantization=True, quantization="fp8_full")
-
-  def test_pure_nnx_decoder_with_qwix_quant_ok(self):
-    cfg = self._init(pure_nnx=True, pure_nnx_decoder=True, use_qwix_quantization=True, quantization="fp8_full")
-    self.assertTrue(cfg.pure_nnx_decoder)
-
-  def test_bridged_decoder_without_quant_ok(self):
-    cfg = self._init(pure_nnx=True, pure_nnx_decoder=False, quantization="")
-    self.assertEqual(cfg.quantization, "")
 
 
 class NnxAttrsToLinenVarsBridgeTest(unittest.TestCase):

tests/unit/quantizations_test.py

@@ -395,173 +395,95 @@ def quantization_config(self, quant, logits_tolerance=2e-1, grad_tolerance=5e-1,
     cfg = self.init_pyconfig(quantization=quant, **kwargs)
     ids, decoder_segment_ids, decoder_positions = self.get_data()
 
-    if cfg.pure_nnx:
-      qt_model = model_creation_utils.create_model(cfg, self.mesh, rngs=nnx.Rngs(0))
-      if getattr(self.__class__, "_cached_base_results_nnx", None) is None:
-        base_cfg = self.init_pyconfig(quantization="", **kwargs)
-        base_model = model_creation_utils.create_model(base_cfg, self.mesh, rngs=nnx.Rngs(0))
-
-        def loss_base(model):
-          logits = model(
-              decoder_input_tokens=ids,
-              decoder_positions=decoder_positions,
-              decoder_segment_ids=decoder_segment_ids,
-              enable_dropout=False,
-          )
-          return jnp.mean((logits) ** 2)
-
-        grads_base = nnx.grad(loss_base)(base_model)
-        logits_base = base_model(
-            decoder_input_tokens=ids,
-            decoder_positions=decoder_positions,
-            decoder_segment_ids=decoder_segment_ids,
-            enable_dropout=False,
-        )
-        self.__class__._cached_base_results_nnx = (grads_base, logits_base)
-
-      grads_base, logits = self.__class__._cached_base_results_nnx
+    qt_model = model_creation_utils.create_model(cfg, self.mesh, rngs=nnx.Rngs(0))
+    if getattr(self.__class__, "_cached_base_results_nnx", None) is None:
+      base_cfg = self.init_pyconfig(quantization="", **kwargs)
+      base_model = model_creation_utils.create_model(base_cfg, self.mesh, rngs=nnx.Rngs(0))
 
-      def loss_quant(model):
-        logits_q = model(
+      def loss_base(model):
+        logits = model(
             decoder_input_tokens=ids,
             decoder_positions=decoder_positions,
             decoder_segment_ids=decoder_segment_ids,
             enable_dropout=False,
         )
-        return jnp.mean((logits_q) ** 2)
+        return jnp.mean((logits) ** 2)
 
-      grads_quant = nnx.grad(loss_quant)(qt_model)
-      quant_logits = qt_model(
+      grads_base = nnx.grad(loss_base)(base_model)
+      logits_base = base_model(
           decoder_input_tokens=ids,
           decoder_positions=decoder_positions,
           decoder_segment_ids=decoder_segment_ids,
           enable_dropout=False,
       )
+      self.__class__._cached_base_results_nnx = (grads_base, logits_base)
 
-      print("relative error in logits:" f" {jnp.abs(quant_logits - logits).mean() / jnp.abs(logits).mean()}")
-      assert jnp.abs(quant_logits - logits).mean() / jnp.abs(logits).mean() < logits_tolerance
-
-      # nnx.grad returns a State object which is a mapping of paths to gradients.
-      # Flatten them to check for tolerance.
-      grads_base_flat = traversals.flatten_mapping(grads_base)
-      grads_quant_flat = traversals.flatten_mapping(grads_quant)
-
-      # Filter for param collections to compare only parameters and not stats/buffers if any
-      # Note: NNX grads structure might contain variables like 'kernel', 'bias'.
-      # For simplicity we compare all matching keys.
-      def flatten_and_filter(grads_flat):
-        return {k: v for k, v in grads_flat.items() if hasattr(v, "shape") and "quant_stats" not in str(k)}
-
-      gb_f = flatten_and_filter(grads_base_flat)
-      gq_f = flatten_and_filter(grads_quant_flat)
-
-      for k in gb_f:
-        if k in gq_f:
-          diff = jnp.abs(gb_f[k] - gq_f[k]).mean() / (jnp.abs(gb_f[k]).mean() + 1e-8)
-          if diff > grad_tolerance:
-            print(f"Gradient mismatch for {k}: rel_error = {diff}")
-            assert diff <= grad_tolerance
-    else:
-      qt_model = model_creation_utils.create_model(cfg, self.mesh)
-      if not hasattr(self.__class__, "_cached_base_results"):
-        model = model_creation_utils.create_model(self.cfg, self.mesh)
-        var = model.init(
-            {"params": self.rng, "aqt": self.rng, "dropout": self.rng},
-            ids,
-            decoder_positions,
-            decoder_segment_ids,
-            enable_dropout=False,
-            mutable=True,
-        )
-
-        def loss_base_linen(all_vars, inputs):
-          logits_b, _ = model.apply(
-              all_vars,
-              *inputs,
-              enable_dropout=False,
-              rngs={"params": self.rng},
-              mutable=True,
-          )
-          return jnp.mean((logits_b) ** 2)
-
-        grads_base_linen = jax.grad(loss_base_linen)(var, (ids, decoder_positions, decoder_segment_ids))
-        logits_b, _ = model.apply(
-            var,
-            ids,
-            decoder_positions,
-            decoder_segment_ids,
-            enable_dropout=False,
-            rngs={"params": self.rng},
-            mutable=True,
-        )
-        self.__class__._cached_base_results = (grads_base_linen, logits_b)
+    grads_base, logits = self.__class__._cached_base_results_nnx
 
-      grads_base_linen, logits = self.__class__._cached_base_results
-
-      quantized_vars = qt_model.init(
-          {"params": self.rng, "aqt": self.rng, "dropout": self.rng},
-          ids,
-          decoder_positions,
-          decoder_segment_ids,
+    def loss_quant(model):
+      logits_q = model(
+          decoder_input_tokens=ids,
+          decoder_positions=decoder_positions,
+          decoder_segment_ids=decoder_segment_ids,
           enable_dropout=False,
-          mutable=True,
       )
+      return jnp.mean((logits_q) ** 2)
+
+    grads_quant = nnx.grad(loss_quant)(qt_model)
+    quant_logits = qt_model(
+        decoder_input_tokens=ids,
+        decoder_positions=decoder_positions,
+        decoder_segment_ids=decoder_segment_ids,
+        enable_dropout=False,
+    )
 
-      def loss_quant_linen(all_vars, inputs):
-        logits_q, _ = qt_model.apply(
-            all_vars,
-            *inputs,
-            enable_dropout=False,
-            rngs={"params": self.rng},
-            mutable=True,
-        )
-        return jnp.mean((logits_q) ** 2)
+    print("relative error in logits:" f" {jnp.abs(quant_logits - logits).mean() / jnp.abs(logits).mean()}")
+    assert jnp.abs(quant_logits - logits).mean() / jnp.abs(logits).mean() < logits_tolerance
 
-      grads_quant_linen = jax.grad(loss_quant_linen)(quantized_vars, (ids, decoder_positions, decoder_segment_ids))
+    # nnx.grad returns a State object which is a mapping of paths to gradients.
+    # Flatten them to check for tolerance.
+    grads_base_flat = traversals.flatten_mapping(grads_base)
+    grads_quant_flat = traversals.flatten_mapping(grads_quant)
 
-      quant_logits, _ = qt_model.apply(
-          quantized_vars,
-          ids,
-          decoder_positions,
-          decoder_segment_ids,
-          enable_dropout=False,
-          rngs={"params": self.rng},
-          mutable=True,
-      )
-      print("relative error in logits:" f" {jnp.abs(quant_logits - logits).mean() / jnp.abs(logits).mean()}")
-      assert jnp.abs(quant_logits - logits).mean() / jnp.abs(logits).mean() < logits_tolerance
-      self.print_grad_diff(grads_base_linen["params"], grads_quant_linen["params"])
-      self.assertTrue(
-          self.pytree_allclose(
-              grads_base_linen["params"],
-              grads_quant_linen["params"],
-              tolerance=grad_tolerance,
-          )
-      )
+    # Filter for param collections to compare only parameters and not stats/buffers if any
+    # Note: NNX grads structure might contain variables like 'kernel', 'bias'.
+    # For simplicity we compare all matching keys.
+    def flatten_and_filter(grads_flat):
+      return {k: v for k, v in grads_flat.items() if hasattr(v, "shape") and "quant_stats" not in str(k)}
+
+    gb_f = flatten_and_filter(grads_base_flat)
+    gq_f = flatten_and_filter(grads_quant_flat)
+
+    for k in gb_f:
+      if k in gq_f:
+        diff = jnp.abs(gb_f[k] - gq_f[k]).mean() / (jnp.abs(gb_f[k]).mean() + 1e-8)
+        if diff > grad_tolerance:
+          print(f"Gradient mismatch for {k}: rel_error = {diff}")
+          assert diff <= grad_tolerance
 
   @pytest.mark.tpu_only
   def test_int8_quantization(self):
     self.quantization_config("int8")
 
   @pytest.mark.tpu_only
   def test_int8_quantization_nnx(self):
-    self.quantization_config("int8", enable_nnx=True, pure_nnx_decoder=True, pure_nnx=True)
+    self.quantization_config("int8")
 
   @pytest.mark.tpu_only
   def test_fp8_quantization(self):
     self.quantization_config("fp8")
 
   @pytest.mark.tpu_only
   def test_fp8_quantization_nnx(self):
-    self.quantization_config("fp8", enable_nnx=True, pure_nnx_decoder=True, pure_nnx=True)
+    self.quantization_config("fp8")
 
   @pytest.mark.tpu_only
   def test_fp8_full_quantization(self):
     self.quantization_config("fp8_full")
 
   @pytest.mark.tpu_only
   def test_fp8_full_quantization_nnx(self):
-    self.quantization_config("fp8_full", enable_nnx=True, pure_nnx_decoder=True, pure_nnx=True)
+    self.quantization_config("fp8_full")
 
   @pytest.mark.gpu_only
   @pytest.mark.external_serving
@@ -571,7 +493,7 @@ def test_fp8_gpu_quantization(self):
   @pytest.mark.gpu_only
   @pytest.mark.external_serving
   def test_fp8_gpu_quantization_nnx(self):
-    self.quantization_config("fp8_gpu", grad_tolerance=1.5, enable_nnx=True, pure_nnx_decoder=True, pure_nnx=True)
+    self.quantization_config("fp8_gpu", grad_tolerance=1.5)
 
   @pytest.mark.gpu_only
   @pytest.mark.external_serving
@@ -581,7 +503,7 @@ def test_fp8_nanoo_quantization(self):
   @pytest.mark.gpu_only
   @pytest.mark.external_serving
   def test_fp8_nanoo_quantization_nnx(self):
-    self.quantization_config("fp8_nanoo", grad_tolerance=1.5, enable_nnx=True, pure_nnx_decoder=True, pure_nnx=True)
+    self.quantization_config("fp8_nanoo", grad_tolerance=1.5)
 
   @pytest.mark.skip(reason="No runner with GPU arch >= 89 is available")
   @pytest.mark.gpu_only
@@ -662,8 +584,6 @@ def test_maybe_quantize_model_pops_intermediates(self):
         quantization="int8",
         use_qwix_quantization=True,
         use_batch_split_schedule=False,
-        pure_nnx=True,
-        pure_nnx_decoder=True,
         micro_batch_size_to_train_on=1,
         max_target_length=2,
     )
@@ -700,7 +620,6 @@ def test_nnx_abstract_state_has_no_intermediates(self):
         enable_checkpointing=False,
         model_name="deepseek3-tiny",
         attention="dot_product",
-        pure_nnx=True,
         use_qwix_quantization=True,
         use_qk_clip=True,  # This sows QK clip intermediates during the forward pass
     )