test(quant): Add NNX versions of Qwix quantization unit tests

Author: hsuan-lun-chiang

src/maxtext/layers/nnx_wrappers.py

@@ -499,7 +499,16 @@ def maybe_unbox(x):
       for path, _ in unknown_state_flat.items():
         paths_str += f"\n  - {'/'.join(map(str, path))}"
 
-      warnings.warn(f"Found unknown module paths in incoming state:{paths_str}")
+        # Dynamically reconstruct the unknown variables
+        curr = module
+        for p in path[:-1]:
+          if not hasattr(curr, p):
+            setattr(curr, p, nnx.Module())
+          curr = getattr(curr, p)
+
+      warnings.warn(
+          f"Found unknown module paths in incoming state:{paths_str}. Intermediate modules have been reconstructed."
+      )
 
     nnx.update(module, new_state)
     _refresh_variable_trace_state(module)

src/maxtext/layers/quantizations.py

@@ -38,9 +38,13 @@
 from flax.linen import fp8_ops
 from flax.linen import initializers as flax_initializers
 import flax.linen as nn
+from flax import nnx
+
+from qwix._src import flax_util
 
 from maxtext.common.common_types import DType, Config
 from maxtext.inference.kvcache import KVQuant
+from maxtext.layers import nnx_wrappers
 
 # Params used to define mixed precision quantization configs
 DEFAULT = "__default__"  # default config
@@ -708,18 +712,16 @@ def configure_kv_quant(config):
 
 
 def _apply_linen_module_in_nnx(linen_module_cls, op_id, *args, **kwargs):
+  """Applies a Linen module within an NNX context."""
   try:
-    from qwix._src import flax_util
     parent = flax_util.get_current_module()
-    from flax import nnx
     is_nnx = isinstance(parent, nnx.Module)
-  except Exception:
+  except Exception:  # pylint: disable=broad-exception-caught
     is_nnx = False
 
   if is_nnx:
     attr_name = f"_qwix_fp8_gpu_{op_id}"
     if not hasattr(parent, attr_name):
-      from maxtext.layers import nnx_wrappers
       rngs = getattr(parent, "qwix_rngs", None)
       if rngs is None:
         rngs = nnx.Rngs(0)
@@ -838,11 +840,17 @@ def maybe_quantize_model(model, config):
     if quantization_provider:
       if config.pure_nnx:
         input_shape = (config.micro_batch_size_to_train_on, config.max_target_length)
-        import jax.numpy as jnp
         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)
-        model = qwix.quantize_model(model, quantization_provider, dummy_tokens, dummy_positions, dummy_segment_ids, enable_dropout=False)
+        model = qwix.quantize_model(
+            model,
+            quantization_provider,
+            dummy_tokens,
+            dummy_positions,
+            dummy_segment_ids,
+            enable_dropout=False,
+        )
       else:
         model = qwix.quantize_model(model, quantization_provider)
   return model

src/maxtext/trainers/pre_train/train.py

@@ -349,7 +349,9 @@ def train_step(model, config, state_mesh_shardings, params_shardings, state, dat
           is_train=True,
       )
     else:
-      OverwriteWithGradient = nnx.variablelib.variable_type_from_name(maxtext_utils.OVERWRITE_WITH_GRADIENT, allow_register=True)
+      OverwriteWithGradient = nnx.variablelib.variable_type_from_name(
+          maxtext_utils.OVERWRITE_WITH_GRADIENT, allow_register=True
+      )
       model_graphdef, curr_params, overwrite_vars, rest = nnx.split(state.model, nnx.Param, OverwriteWithGradient, ...)
       if config.parameter_memory_host_offload:
         # Params are kept on host (pinned_host) in in_shardings. Move only Param
@@ -379,7 +381,9 @@ def diff_wrapper(param, overwrite_vars, rest, config, data):
         return loss, (aux, new_overwrite_vars, new_rest)
 
       grad_func = jax.value_and_grad(diff_wrapper, argnums=(0, 1), has_aux=True)
-      (loss, (aux, new_overwrite_vars, new_rest)), (raw_grads, overwrite_grads) = grad_func(curr_params, overwrite_vars, rest, config, data)
+      (loss, (aux, _, new_rest)), (raw_grads, overwrite_grads) = grad_func(
+          curr_params, overwrite_vars, rest, config, data
+      )
       nnx.update(state.model, new_rest)
       nnx.update(state.model, overwrite_grads)
 

tests/unit/quantizations_test.py

@@ -22,6 +22,7 @@
 from aqt.jax.v2 import aqt_tensor
 from aqt.jax.v2.flax import aqt_flax
 from flax import nnx
+from flax.nnx import traversals
 import jax
 from jax import lax
 from jax import numpy as jnp
@@ -48,7 +49,7 @@ def __init__(
       self,
       quantization: quantizations.AqtQuantization,
       data_type: Any,
-      rngs: nnx.Rngs,
+      rngs: nnx.Rngs,  # pylint: disable=unused-argument
   ):
     self.quantization = quantization
     self.identity = jnp.identity(2, dtype=data_type)
@@ -441,10 +442,9 @@ def loss_quant(model):
 
       # nnx.grad returns a State object which is a mapping of paths to gradients.
       # Flatten them to check for tolerance.
-      from flax.nnx import traversals
       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.
@@ -542,29 +542,46 @@ def loss_quant_linen(all_vars, inputs):
   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)
+
   @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)
+
   @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)
+
   @pytest.mark.gpu_only
   @pytest.mark.external_serving
   def test_fp8_gpu_quantization(self):
     self.quantization_config("fp8_gpu", grad_tolerance=1.5)
 
-  # @pytest.mark.gpu_only
+  @pytest.mark.gpu_only
   @pytest.mark.external_serving
-  def test_fp8_gpu_quantization(self):
+  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)
 
   @pytest.mark.gpu_only
   @pytest.mark.external_serving
   def test_fp8_nanoo_quantization(self):
     self.quantization_config("fp8_nanoo", grad_tolerance=1.5)
 
+  @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)
+
   @pytest.mark.skip(reason="No runner with GPU arch >= 89 is available")
   @pytest.mark.gpu_only
   def test_fp8_te_fp8_delayedscaling_quantization(self):