From 7da8b8a58a1bbee7422a3e3bde13667bbcfc3b45 Mon Sep 17 00:00:00 2001 From: Lance Wang Date: Wed, 29 Jul 2026 15:26:59 +0000 Subject: [PATCH] [NNX] Delete Linen (pre-train 3/3): collapse dispatch in quantization and model creation - quantizations.maybe_quantize_model always runs the qwix forward pass with the dummy tokens/positions/segment ids (and the MTP decoder targets when mtp_num_layers > 0), then pops the transient nnx.Intermediate variables the traced forward sows. - model_creation_utils.from_pretrained always builds the sharded model through maxtext_utils_nnx.create_nnx_sharded_model. Tests follow: quantizations_test and nnx_quant_guard_test drop their flag arguments and Linen expectations, correctness_tests_nnx_dispatch_test keeps only the NNX case, and forward_pass_logit_checker always loads via from_pretrained. --- src/maxtext/layers/quantizations.py | 48 +++--- src/maxtext/utils/model_creation_utils.py | 7 +- tests/unit/nnx_quant_guard_test.py | 31 +--- tests/unit/quantizations_test.py | 183 ++++++---------------- tests/utils/forward_pass_logit_checker.py | 46 ++---- 5 files changed, 92 insertions(+), 223 deletions(-) diff --git a/src/maxtext/layers/quantizations.py b/src/maxtext/layers/quantizations.py index 0ef3aeca55..f2b964f5d0 100644 --- a/src/maxtext/layers/quantizations.py +++ b/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 diff --git a/src/maxtext/utils/model_creation_utils.py b/src/maxtext/utils/model_creation_utils.py index a2a1403125..4131dec541 100644 --- a/src/maxtext/utils/model_creation_utils.py +++ b/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) diff --git a/tests/unit/nnx_quant_guard_test.py b/tests/unit/nnx_quant_guard_test.py index 50cac5d349..8451dd1264 100644 --- a/tests/unit/nnx_quant_guard_test.py +++ b/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): diff --git a/tests/unit/quantizations_test.py b/tests/unit/quantizations_test.py index d014a66ee0..8d6dc73580 100644 --- a/tests/unit/quantizations_test.py +++ b/tests/unit/quantizations_test.py @@ -395,149 +395,71 @@ 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): @@ -545,7 +467,7 @@ def test_int8_quantization(self): @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): @@ -553,7 +475,7 @@ def test_fp8_quantization(self): @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): @@ -561,7 +483,7 @@ def test_fp8_full_quantization(self): @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 ) diff --git a/tests/utils/forward_pass_logit_checker.py b/tests/utils/forward_pass_logit_checker.py index a51b23980f..66d6fa4343 100644 --- a/tests/utils/forward_pass_logit_checker.py +++ b/tests/utils/forward_pass_logit_checker.py @@ -71,7 +71,6 @@ """ import argparse -import functools import os from pathlib import Path import sys @@ -83,8 +82,6 @@ from maxtext.utils.globals import MAXTEXT_TEST_ASSETS_ROOT, HF_IDS from maxtext.checkpoint_conversion.utils.hf_utils import convert_jax_weight_to_torch from maxtext.common.common_types import DECODING_ACTIVE_SEQUENCE_INDICATOR, MODEL_MODE_TRAIN -from maxtext.layers import quantizations -from maxtext.models import models from maxtext.utils import max_logging from maxtext.utils import maxtext_utils from maxtext.utils import model_creation_utils @@ -356,7 +353,7 @@ def get_data(golden_data_point, config): def main(config, test_args): # pylint: disable=W0621 """Test the Whole Model of model_name""" init_rng = jax.random.PRNGKey(config.init_weights_seed) - init_rng, rng1 = jax.random.split(init_rng) + init_rng, _ = jax.random.split(init_rng) devices_array = maxtext_utils.create_device_mesh(config) mesh = jax.sharding.Mesh(devices_array, config.mesh_axes) @@ -393,19 +390,13 @@ def main(config, test_args): # pylint: disable=W0621 if not test_args.run_hf_model: """Comparing maxtext/huggingface model with pre-loaded golden logitis""" max_logging.log("Initializing MaxText model") - quant = quantizations.configure_quantization(config) - if config.pure_nnx_decoder and config.enable_nnx: - model = model_creation_utils.from_pretrained(config, mesh=mesh, model_mode=MODEL_MODE_TRAIN) - - if config.lora.enable_lora: - model = lora_utils.apply_lora_to_model(model, mesh, config) - if config.lora.lora_restore_path: - lora_utils.restore_lora_from_path(model, config) - state = None - else: - model = models.transformer_as_linen(config, mesh=mesh, quant=quant, model_mode=MODEL_MODE_TRAIN) - init_state_fn = functools.partial(maxtext_utils.init_initial_state, model, None, config, False, rng1) - state, _ = maxtext_utils.setup_decode_state(config, mesh, None, init_state_fn) + model = model_creation_utils.from_pretrained(config, mesh=mesh, model_mode=MODEL_MODE_TRAIN) + + if config.lora.enable_lora: + model = lora_utils.apply_lora_to_model(model, mesh, config) + if config.lora.lora_restore_path: + lora_utils.restore_lora_from_path(model, config) + state = None if test_args.golden_logits_path == "": input_golden_data_path = os.path.join( @@ -641,22 +632,13 @@ def main(config, test_args): # pylint: disable=W0621 raise ImportError("peft library is required to load HF LoRA adapter. Run `pip install peft`.") from exc hf_model = PeftModel.from_pretrained(hf_model, hf_lora_path) - quant = quantizations.configure_quantization(config) - if config.pure_nnx_decoder and config.enable_nnx: - maxtext_model = model_creation_utils.from_pretrained(config, mesh=mesh, model_mode=MODEL_MODE_TRAIN) + maxtext_model = model_creation_utils.from_pretrained(config, mesh=mesh, model_mode=MODEL_MODE_TRAIN) - if config.lora.enable_lora: - maxtext_model = lora_utils.apply_lora_to_model(maxtext_model, mesh, config) - if config.lora.lora_restore_path: - lora_utils.restore_lora_from_path(maxtext_model, config) - maxtext_state = None - else: - maxtext_model = models.transformer_as_linen(config, mesh, quant=quant, model_mode=MODEL_MODE_TRAIN) - init_state_fn = functools.partial(maxtext_utils.init_initial_state, maxtext_model, None, config, False, rng1) - if test_args.ckpt_type == "linen": - maxtext_state, _ = maxtext_utils.setup_decode_state(config, mesh, None, init_state_fn) - else: - maxtext_state, _ = model_creation_utils.setup_decode_state_from_nnx(maxtext_model, config, rng1, mesh) + if config.lora.enable_lora: + maxtext_model = lora_utils.apply_lora_to_model(maxtext_model, mesh, config) + if config.lora.lora_restore_path: + lora_utils.restore_lora_from_path(maxtext_model, config) + maxtext_state = None # The long prompt is required to catch position-dependent regressions (e.g. RoPE); # the short prompts above cannot detect them. See build_long_prompt().