Share one RNG stream across the bridged NVFP4 quantizers

src/maxtext/layers/linears.py

@@ -212,9 +212,12 @@ def __init__(
       block_size = getattr(quant, "get_block_size", lambda: 1)()  # needed for TE MXFP8
       dummy_inputs = jnp.zeros((block_size, *self.in_features_shape), dtype=self.dtype)
       self(dummy_inputs, _initializing=True)
-      # Backends that never draw at apply time leave dead RNG state in the model.
-      if not quant.needs_apply_rngs:
+      # Left alone, every bridged wrapper carries its own forked Rngs, which the
+      # unrolled decoder pays for once per layer.
+      if quant.apply_rngs is quantizations.ApplyRngs.NONE:
         quant_dot_general.release_rngs()
+      elif quant.apply_rngs is quantizations.ApplyRngs.SHARED:
+        quant_dot_general.share_rngs(rngs)
     else:
       self._quant_dot_general_name = None
 

src/maxtext/layers/moe.py

@@ -411,8 +411,10 @@ def __init__(
       dummy_inputs = jnp.zeros((1, *self.in_features_shape), dtype=self.dtype)
       self(dummy_inputs, _initializing=True)
       # See the matching comment in linears.py.
-      if not quant.needs_apply_rngs:
+      if quant.apply_rngs is quantizations.ApplyRngs.NONE:
         quant_dot_general.release_rngs()
+      elif quant.apply_rngs is quantizations.ApplyRngs.SHARED:
+        quant_dot_general.share_rngs(rngs)
     else:
       self._quant_dot_general_name = None
 

src/maxtext/layers/nnx_wrappers.py

@@ -239,6 +239,20 @@ def release_rngs(self):
     """
     self.to_nnx__rngs = None
 
+  def share_rngs(self, rngs: Rngs | None):
+    """Draws from ``rngs`` instead of from this wrapper's own fork.
+
+    ``__init__`` still forked, so the caller's streams have advanced exactly as they
+    would have and parameter initialization is unchanged. What changes is where
+    apply-time draws come from. NNX stores a shared ``Rngs`` once however many modules
+    reference it, so a decoder with ``scan_layers=False`` stops paying for one per layer.
+
+    Only for wrapped modules whose draws are decorrelated by something other than the
+    stream itself, as TransformerEngine's stochastic rounding is by its per-quantizer
+    hash.
+    """
+    self.to_nnx__rngs = rngs
+
   def __getattr__(self, name: str):
     if hasattr(super(), name):
       return super().__getattribute__(name)

src/maxtext/layers/quantizations.py

@@ -14,6 +14,7 @@
 
 """Quantization library."""
 
+import enum
 import functools
 import json
 import qwix.pallas as qpl
@@ -76,14 +77,30 @@ def _safe_find_param(x, ptq_array_type=None):
 _TILE_SIZE = "tile_size"  # Tile size for subchannel
 
 
+class ApplyRngs(enum.Enum):
+  """How a quantization backend needs the Linen->NNX bridge to handle its RNGs.
+
+  The bridge forks the caller's `Rngs` into every wrapper it creates, and anything a
+  module holds is model state. Unrolled, that is paid once per layer, so a backend that
+  does not need its own fork should say so.
+  """
+
+  # Never draws at apply time, so the bridge can drop the fork entirely.
+  NONE = "none"
+  # Draws, but its values are decorrelated by something other than the stream, so one
+  # shared stream serves every wrapper.
+  SHARED = "shared"
+  # Draws and needs its own stream. The safe default.
+  PRIVATE = "private"
+
+
 @dataclass
 class Quantization:
   """Base class for quantization configurations"""
 
-  # Whether this backend's dot_general draws RNGs at apply time, as AQT and NVFP4
-  # stochastic rounding do. False lets the Linen->NNX bridge drop its forked Rngs after
-  # init; the default makes opting out deliberate.
-  needs_apply_rngs: ClassVar[bool] = True
+  # How the bridge should handle this backend's RNGs. The default makes anything else
+  # deliberate.
+  apply_rngs: ClassVar[ApplyRngs] = ApplyRngs.PRIVATE
 
   def dot_general_cls(self, mesh_axes: Tuple[str, ...] = ()):
     """Placeholder for dot_general implementation in subclasses."""
@@ -151,7 +168,7 @@ class AqtQuantization:
 
   # Declared, not inherited: this class sits outside the Quantization hierarchy. AQT
   # sets rng_type="jax.uniform" and stochastic rounding draws at apply time.
-  needs_apply_rngs: ClassVar[bool] = True
+  apply_rngs: ClassVar[ApplyRngs] = ApplyRngs.PRIVATE
 
   quant_dg: aqt_config.DotGeneral
   quant_mode: aqt_flax.QuantMode = aqt_flax.QuantMode.TRAIN
@@ -236,7 +253,7 @@ class QwixQuantization:
   """Configures Qwix quantization github.com/google/qwix, for training only."""
 
   # Declared, not inherited: this class sits outside the Quantization hierarchy.
-  needs_apply_rngs: ClassVar[bool] = True
+  apply_rngs: ClassVar[ApplyRngs] = ApplyRngs.PRIVATE
 
   quant_mode = "train"  # needed by external call
   act_calibration_method: str = "absmax"
@@ -319,7 +336,7 @@ class Fp8Quantization(Quantization):
   """Configures Fp8 quantization for NVIDIA GPUs"""
 
   # Flax's fp8 ops scale from amax history, never from make_rng at apply time.
-  needs_apply_rngs: ClassVar[bool] = False
+  apply_rngs: ClassVar[ApplyRngs] = ApplyRngs.NONE
 
   quant_mode = "train"
   # The forward dtype, for callers that quantize an operand themselves rather than through
@@ -414,7 +431,7 @@ class NANOOFp8Quantization(Quantization):
   """Configures NANOO Fp8 quantization for AMD MI300/MI325 GPUs"""
 
   # Same as Fp8Quantization: nn.NANOOFp8DotGeneralOp never draws at apply time.
-  needs_apply_rngs: ClassVar[bool] = False
+  apply_rngs: ClassVar[ApplyRngs] = ApplyRngs.NONE
 
   quant_mode = "train"
   quantize_dtype = jnp.float8_e4m3fnuz
@@ -1117,17 +1134,19 @@ def _get_recipe(recipe_name: str):
     return RECIPES[recipe_name]()
 
   @property
-  def needs_apply_rngs(self) -> bool:
-    """Whether this recipe draws RNGs at apply time.
+  def apply_rngs(self) -> ApplyRngs:
+    """How the bridge should handle this recipe's RNGs.
 
-    Only NVFP4 does, and only while stochastic rounding is on: TE draws ``sr_rng`` for
-    the DGRAD quantizer. The other recipes take their scales from the tensors.
+    Only NVFP4 draws, and only while stochastic rounding is on: TE draws ``sr_rng`` for
+    the DGRAD quantizer. It folds a per-quantizer hash into whatever it draws, so the
+    wrappers do not need streams of their own. The other recipes take their scales from
+    the tensors and never draw.
     """
     from transformer_engine.common import recipe  # pylint: disable=import-outside-toplevel # pytype: disable=import-error
 
     if not isinstance(self._recipe, recipe.NVFP4BlockScaling):  # pytype: disable=module-attr
-      return False
-    return not self._recipe.disable_stochastic_rounding
+      return ApplyRngs.NONE
+    return ApplyRngs.SHARED if not self._recipe.disable_stochastic_rounding else ApplyRngs.NONE
 
   def get_block_size(self):
     """Get the block size for quantization for recipes that require blocks.

tests/unit/nnx_quant_bridge_rng_test.py

@@ -55,7 +55,7 @@ def __call__(self, inputs, kernel, dims, precision=None, **kwargs):
 class _OptOutQuant(quantizations.Quantization):
   """A backend that never draws RNGs at apply time (as TransformerEngine does not)."""
 
-  needs_apply_rngs = False
+  apply_rngs = quantizations.ApplyRngs.NONE
   quant_mode = "train"  # read by quantizations.in_serve_mode()
 
   def dot_general_cls(self, mesh_axes=()):
@@ -64,7 +64,7 @@ def dot_general_cls(self, mesh_axes=()):
 
 
 class _DefaultQuant(quantizations.Quantization):
-  """A backend that leaves `needs_apply_rngs` at its safe default of True."""
+  """A backend that leaves `apply_rngs` at its safe default of PRIVATE."""
 
   quant_mode = "train"
 
@@ -97,7 +97,13 @@ def dot_general_cls(self, mesh_axes=()):
 class _MisdeclaredQuant(_DrawingQuant):
   """A drawing backend that wrongly claims it does not draw."""
 
-  needs_apply_rngs = False
+  apply_rngs = quantizations.ApplyRngs.NONE
+
+
+class _SharedQuant(_DrawingQuant):
+  """A drawing backend whose values do not depend on having a stream of its own."""
+
+  apply_rngs = quantizations.ApplyRngs.SHARED
 
 
 def _rng_state_paths(module) -> list[str]:
@@ -126,15 +132,15 @@ def test_opt_out_backend_leaves_no_rng_state(self):
     self.assertEqual(
         _rng_state_paths(dense),
         [],
-        "A backend with needs_apply_rngs=False must not retain the bridge's forked "
+        "A backend with apply_rngs=NONE must not retain the bridge's forked "
         "Rngs: its counters would be incremented on device every step, once per "
         "unrolled layer.",
     )
 
   def test_default_backend_keeps_rng_state(self):
     """Checks that the opt-out is deliberate and the default stays safe."""
     paths = _rng_state_paths(_make_dense(_DefaultQuant()))
-    self.assertNotEqual(paths, [], "needs_apply_rngs defaults to True, so the Rngs must be kept.")
+    self.assertNotEqual(paths, [], "apply_rngs defaults to PRIVATE, so the Rngs must be kept.")
 
   def test_unquantized_dense_has_no_bridge_state(self):
     self.assertEqual(_rng_state_paths(_make_dense(None)), [])
@@ -213,37 +219,37 @@ def test_zero_rate_dropout_is_a_passthrough(self):
 
 
 class QuantizationFlagTest(unittest.TestCase):
-  """`needs_apply_rngs` must stay safe-by-default and correct per backend."""
+  """`apply_rngs` must stay safe-by-default and correct per backend."""
 
-  def test_base_class_defaults_to_keeping_rngs(self):
-    self.assertTrue(quantizations.Quantization.needs_apply_rngs)
+  def test_base_class_defaults_to_a_private_stream(self):
+    self.assertIs(quantizations.Quantization.apply_rngs, quantizations.ApplyRngs.PRIVATE)
 
-  def test_backends_that_may_draw_at_apply_time_keep_rngs(self):
-    """AQT's config enables jax.uniform RNG, so it must never be opted out silently."""
+  def test_backends_that_may_draw_at_apply_time_keep_their_own_rngs(self):
+    """AQT's config enables jax.uniform RNG, so it must never be narrowed silently."""
     for cls in (quantizations.AqtQuantization, quantizations.QwixQuantization):
-      self.assertTrue(cls.needs_apply_rngs, f"{cls.__name__} must keep its RNGs")
+      self.assertIs(cls.apply_rngs, quantizations.ApplyRngs.PRIVATE, f"{cls.__name__} must keep its own RNGs")
 
   def test_fp8_backends_opt_out(self):
     """Flax's fp8 ops scale from amax history, so they never draw at apply time."""
     for cls in (quantizations.Fp8Quantization, quantizations.NANOOFp8Quantization):
-      self.assertFalse(cls.needs_apply_rngs, f"{cls.__name__} must release the bridge's Rngs")
+      self.assertIs(cls.apply_rngs, quantizations.ApplyRngs.NONE, f"{cls.__name__} must release the bridge's Rngs")
 
   def test_every_backend_declares_the_flag(self):
     """Every backend must carry the flag, so the call sites can read it directly.
 
     `AqtQuantization` and `QwixQuantization` sit outside the `Quantization` hierarchy
     and declare it themselves. A backend that is missing it raises AttributeError at
     the call site rather than silently taking a default. `TransformerEngineQuantization`
-    is excluded because its answer depends on the recipe, so it derives the flag per
-    instance; `TransformerEngineRecipeRngTest` covers it.
+    is excluded because its answer depends on the recipe, so it derives it per instance;
+    `TransformerEngineRecipeRngTest` covers it.
     """
     for cls in (
         quantizations.AqtQuantization,
         quantizations.QwixQuantization,
         quantizations.Fp8Quantization,
         quantizations.NANOOFp8Quantization,
     ):
-      self.assertIsInstance(cls.needs_apply_rngs, bool, f"{cls.__name__} must declare needs_apply_rngs")
+      self.assertIsInstance(cls.apply_rngs, quantizations.ApplyRngs, f"{cls.__name__} must declare apply_rngs")
 
 
 class ApplyTimeRngTest(unittest.TestCase):
@@ -297,17 +303,54 @@ def test_rng_state_does_not_grow_with_layer_count(self):
     self.assertEqual(counts, [0, 0, 0], f"fp8 RNG state must not scale with layer count, got {counts}")
 
 
+class SharedRngStateTest(unittest.TestCase):
+  """A backend that shares one stream must not add state per unrolled layer.
+
+  This is what NVFP4 needs: it draws every step, so the fork cannot simply be dropped,
+  but TransformerEngine folds a per-quantizer hash into whatever it draws, so one stream
+  serves every wrapper. NNX stores a shared `Rngs` once however many modules hold it.
+  """
+
+  def _stack(self, quant, num_layers):
+    """Builds layers the way an unrolled decoder does, from one `Rngs`."""
+    rngs = nnx.Rngs(params=0, dropout=1, aqt=2)
+    return [
+        linears.DenseGeneral(in_features_shape=8, out_features_shape=4, quant=quant, rngs=rngs) for _ in range(num_layers)
+    ]
+
+  def _rng_leaves(self, layers) -> int:
+    return len(_rng_state_paths(nnx.List(layers)))
+
+  def test_state_does_not_grow_with_layer_count(self):
+    counts = [self._rng_leaves(self._stack(_SharedQuant(), n)) for n in (1, 2, 8)]
+    self.assertEqual(
+        len(set(counts)),
+        1,
+        f"a shared stream must cost the same at any depth, got {counts} for 1/2/8 layers",
+    )
+
+  def test_a_private_backend_still_grows(self):
+    """The contrast, so the test above cannot pass for the wrong reason."""
+    counts = [self._rng_leaves(self._stack(_DrawingQuant(), n)) for n in (1, 2, 8)]
+    self.assertEqual(len(set(counts)), 3, f"a private fork per wrapper should scale, got {counts}")
+
+  def test_sharing_keeps_the_layer_callable(self):
+    out = self._stack(_SharedQuant(), 1)[0](jnp.ones((2, 8), jnp.float32))
+    self.assertEqual(out.shape, (2, 4))
+    self.assertTrue(jnp.all(jnp.isfinite(out)))
+
+
 class TransformerEngineRecipeRngTest(unittest.TestCase):
-  """Only NVFP4 with stochastic rounding on draws at apply time."""
+  """Only NVFP4 with stochastic rounding on draws at apply time, and it can share."""
 
   # te_nvfp4 and te_nvfp4_no_rht both leave disable_stochastic_rounding at its False
   # default, so both draw; disable_rht does not affect rounding.
   EXPECTED = {
-      "te_fp8_delayedscaling": False,
-      "te_fp8_currentscaling": False,
-      "te_mxfp8": False,
-      "te_nvfp4": True,
-      "te_nvfp4_no_rht": True,
+      "te_fp8_delayedscaling": quantizations.ApplyRngs.NONE,
+      "te_fp8_currentscaling": quantizations.ApplyRngs.NONE,
+      "te_mxfp8": quantizations.ApplyRngs.NONE,
+      "te_nvfp4": quantizations.ApplyRngs.SHARED,
+      "te_nvfp4_no_rht": quantizations.ApplyRngs.SHARED,
   }
 
   def test_flag_follows_the_recipe(self):
@@ -320,7 +363,7 @@ def test_flag_follows_the_recipe(self):
       with self.subTest(recipe=name):
         config = types.SimpleNamespace(quantization=name, te_comm_gemm_overlap=None)
         quant = quantizations.TransformerEngineQuantization(config)
-        self.assertEqual(quant.needs_apply_rngs, expected)
+        self.assertIs(quant.apply_rngs, expected)
 
 
 if __name__ == "__main__":