perf(modelopt): skip redundant weight fake-quant during frozen-weight logprobs

docs/guides/quantization-aware-rl.md

@@ -39,6 +39,34 @@ for most models, but it is not guaranteed for every architecture or recipe. If
 you encounter errors with the standard Megatron layer specs, leave it unset or
 set it to `false` to exercise ModelOpt's Megatron layer-spec path.
 
+## Frozen-Weight Logprob Optimization
+
+QARL can avoid repeatedly fake-quantizing the same frozen weights during the
+no-gradient policy and reference logprob passes. Set
+`policy.quant_fold_frozen_weight_snap: true` to enable it; it defaults to `false`.
+
+When enabled, NeMo RL folds each enabled fake-quant weight quantizer once at the
+start of the pass, using ModelOpt's fold formula (`QuantModule.fold_weight`): the
+fake-quantized value is written into the existing parameter storage and the weight
+quantizer is disabled. Forwards during the pass then read an already-quantized
+weight instead of re-quantizing it on every microbatch. The original weights are
+restored and the quantizers re-enabled before training resumes.
+
+The fold is applied per discovered weight/quantizer pair rather than through
+`mtq.fold_weight`, which crashes on Megatron models with tied word embeddings
+(the tied `output_layer` carries `weight = None`) and needlessly processes
+disabled quantizers. Disabled quantizers (for example `lm_head` and embeddings in
+the standard recipes) are identity at forward time, so they are skipped entirely.
+
+This applies to any quantization format. Only the *weight* quantizer is disabled,
+so recipes that also quantize activations (such as W4A4) keep their input and
+output quantizers running and produce unchanged logprobs. Weight quantizers built
+as a `SequentialQuantizer` (W4A4 double-quant) are not folded, so those modules
+simply do not benefit.
+
+The option costs one temporary copy of each folded weight shard, held only for the
+duration of the pass.
+
 ## Quantization-Aware GRPO (QA-GRPO)
 
 ### Configuration
@@ -51,6 +79,7 @@ defaults: "../configs/grpo_math_8B_megatron.yaml"
 
 policy:
   quant_cfg: "examples/modelopt/quant_configs/nvfp4_a16.yaml"
+  quant_fold_frozen_weight_snap: true
   quant_calib_data: "cnn_dailymail"
   quant_calib_size: 512
   quant_batch_size: 1
@@ -289,6 +318,7 @@ defaults: "../configs/distillation_math_megatron.yaml"
 
 policy:
     quant_cfg: "NVFP4_DEFAULT_CFG"
+    quant_fold_frozen_weight_snap: true
     quant_calib_data: "cnn_dailymail"
     quant_calib_size: 512
     quant_batch_size: 1
@@ -318,6 +348,7 @@ These parameters are added under the `policy` section:
 | `quant_calib_size` | Number of samples for the calibration pass |
 | `quant_batch_size` | Batch size during calibration |
 | `quant_sequence_length` | Sequence length for calibration data |
+| `quant_fold_frozen_weight_snap` | Optional boolean, default `false`. During frozen-weight logprob passes, fold each enabled fake-quantized weight into its parameter once (ModelOpt's fold formula) instead of re-quantizing every microbatch. Weights and quantizers are restored afterwards. Safe for any format, including activation-quantized recipes such as W4A4. |
 
 The `policy.generation.quant_cfg` should match `policy.quant_cfg` to ensure consistent quantization between training and generation.
 

examples/configs/recipes/llm/grpo-nemotron3-super-120BA12B-16n4g-megatron-qa-nvfp4-w4a16-real-300step.yaml

@@ -6,6 +6,7 @@ checkpointing:
 policy:
   disable_modelopt_layer_spec: true
   quant_cfg: examples/modelopt/quant_configs/nvfp4_experts_weightonly.yaml
+  quant_fold_frozen_weight_snap: true
   generation:
     quant_cfg: examples/modelopt/quant_configs/nvfp4_experts_weightonly.yaml
     real_quant: true

examples/configs/recipes/llm/grpo-nemotron3-super-120BA12B-16n4g-megatron-qa-nvfp4-w4a4-real-300step.yaml

@@ -6,6 +6,7 @@ checkpointing:
 policy:
   disable_modelopt_layer_spec: true
   quant_cfg: examples/modelopt/quant_configs/nvfp4_experts.yaml
+  quant_fold_frozen_weight_snap: true
   quant_calib_data: cnn_dailymail
   quant_calib_size: 16
   quant_batch_size: 1

examples/configs/recipes/llm/grpo-qwen3-30ba3b-4n4g-megatron-qa-nvfp4-w4a4-real.yaml

@@ -12,6 +12,7 @@ checkpointing:
   checkpoint_must_save_by: 00:03:45:00
 policy:
   quant_cfg: examples/modelopt/quant_configs/nvfp4_experts.yaml
+  quant_fold_frozen_weight_snap: true
   generation:
     quant_cfg: examples/modelopt/quant_configs/nvfp4_experts.yaml
     real_quant: true

examples/configs/recipes/llm/grpo-qwen3-8b-base-dapo-2n8g-long-megatron-qa-nvfp4-w4a16.yaml

@@ -29,6 +29,7 @@ policy:
   max_total_sequence_length: 30720
   refit_buffer_size_gb: 4
   quant_cfg: examples/modelopt/quant_configs/nvfp4_a16_mlp_only.yaml
+  quant_fold_frozen_weight_snap: true
   dtensor_cfg:
     enabled: false
   optimizer: null

examples/modelopt/qa_distillation_math_megatron.yaml

@@ -10,6 +10,7 @@ defaults: "../configs/distillation_math_megatron.yaml"
 policy:
     # Quantization config applied to the student's Megatron training worker.
     quant_cfg: "NVFP4_DEFAULT_CFG"
+    quant_fold_frozen_weight_snap: true
     quant_calib_data: "cnn_dailymail"
     quant_calib_size: 512
     quant_batch_size: 1

examples/modelopt/qa_grpo_llama8b_megatron.v2.yaml

@@ -87,6 +87,7 @@ policy:
 
   # NVFP4 weight-only (W4A16) custom recipe applied to the Megatron training worker.
   quant_cfg: "examples/modelopt/quant_configs/nvfp4_a16.yaml"
+  quant_fold_frozen_weight_snap: true
   quant_calib_data: "cnn_dailymail"
   quant_calib_size: 512
   quant_batch_size: 1

examples/modelopt/qa_grpo_math_megatron.yaml

@@ -10,6 +10,7 @@ defaults: "../configs/grpo_math_1B_megatron.yaml"
 policy:
   # Quantization config applied to the Megatron training worker.
   quant_cfg: "NVFP4_DEFAULT_CFG"
+  quant_fold_frozen_weight_snap: true
   quant_calib_data: "cnn_dailymail"
   quant_calib_size: 512
   quant_batch_size: 1

examples/modelopt/qa_grpo_nano3_megatron.yaml

@@ -25,6 +25,7 @@ policy:
   # Nano3 is a hybrid MoE/Mamba model. This recipe keeps attention and the
   # known Nano3-sensitive layers in BF16, while applying NVFP4 to weights.
   quant_cfg: "examples/modelopt/quant_configs/nano3_nvfp4_weightonly.yaml"
+  quant_fold_frozen_weight_snap: true
   quant_calib_data: "cnn_dailymail"
   quant_calib_size: 512
   quant_batch_size: 1

examples/modelopt/qa_grpo_qwen3_30ba3b_megatron.yaml

@@ -23,6 +23,7 @@ policy:
   # Built-in NVFP4 weight-only recipe keeps activations in native dtype while
   # still exercising ModelOpt quantization on Qwen3 MoE/MLP weights.
   quant_cfg: NVFP4_MLP_WEIGHT_ONLY_CFG
+  quant_fold_frozen_weight_snap: true
   quant_calib_data: cnn_dailymail
   quant_calib_size: 16
   quant_batch_size: 1

nemo_rl/modelopt/models/policy/workers/megatron_quant_policy_worker.py

@@ -40,6 +40,9 @@
     quantize_model,
     symlink_pre_quantized_model,
 )
+from nemo_rl.modelopt.models.policy.workers.weight_folding import (
+    temporarily_fold_weights,
+)
 from nemo_rl.modelopt.utils import (
     MODELOPT_REAL_QUANT_ZMQ_TIMEOUT_MS,
     resolve_nvfp4_real_quant_mode,
@@ -431,6 +434,22 @@ def get_quantizer_stats(self) -> dict:
             "positive_amax": positive_amax,
         }
 
+    def get_logprobs(self, *args, **kwargs):
+        """Compute logprobs, optionally folding the frozen weights for the stage.
+
+        With ``policy.quant_fold_frozen_weight_snap`` set, each weight is fake-quantized
+        once for the whole stage instead of on every microbatch forward. Off by default;
+        when unset this is exactly the base implementation.
+
+        Safe only because the re-scoring pass runs under ``no_grad`` with no optimizer
+        step -- the fold is written into the parameter and reverted on exit.
+        """
+        if not self.cfg.get("quant_fold_frozen_weight_snap"):
+            return super().get_logprobs(*args, **kwargs)
+
+        with temporarily_fold_weights(self.model, verbose=True, rank=self.rank):
+            return super().get_logprobs(*args, **kwargs)
+
     def generate(self, **kwargs):
         """Quantized Megatron generation is not supported.
 

nemo_rl/modelopt/models/policy/workers/weight_folding.py

@@ -0,0 +1,140 @@
+# Copyright (c) 2026, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Temporarily fold fake-quantized weights for frozen-weight QAT stages.
+
+In ModelOpt QAT the weight quantizer sits inside the linear, so every forward recomputes
+``weight_quantizer(weight)`` — an elementwise pass over the weight shard, uncached.
+During the ``get_logprobs`` re-scoring stage the weights are frozen (``torch.no_grad()``,
+no optimizer step between microbatches), so that result is identical across every
+microbatch and is pure wasted work.
+
+Folding writes the fake-quantized value into the weight and disables the weight
+quantizer — exactly the frozen-weight steady state. ModelOpt ships this as
+:func:`modelopt.torch.quantization.fold_weight`, and this module applies the same
+per-weight formula (``quant_module.py::QuantModule.fold_weight``), but folds each
+discovered pair directly instead of delegating to the utility, for two reasons:
+
+* ``fold_weight`` selects on ``fake_quant`` alone and dereferences ``weight.data``
+  unconditionally, so it crashes with ``AttributeError`` on Megatron models with tied
+  embeddings, where the ``output_layer`` exposes a ``weight_quantizer`` but its
+  ``weight`` is ``None`` (the embedding weight is borrowed at forward time).
+* Folding a *disabled* quantizer is an identity no-op (a disabled quantizer returns
+  its input unchanged), so cloning those weights for restore would only waste memory —
+  on standard QARL recipes the disabled ``lm_head``/embedding quantizers account for
+  ~40% of the quantized-weight bytes.
+
+The fold is reversible: original weights are cloned before folding and written back
+through their existing parameter storage on exit, so it can wrap a single stage of an
+otherwise-continuing QAT run.
+"""
+
+import contextlib
+from collections.abc import Iterator
+
+import torch
+
+from modelopt.torch.quantization.nn.modules.tensor_quantizer import TensorQuantizer
+
+_QUANTIZER_SUFFIX = "weight_quantizer"
+
+
+def _foldable_weight_quantizers(
+    model: torch.nn.Module,
+) -> Iterator[tuple[torch.Tensor, TensorQuantizer]]:
+    """Yield the enabled ``(weight, quantizer)`` pairs whose fold changes the forward.
+
+    Mirrors ``QuantModule.fold_weight``'s discovery — any attribute whose name ends in
+    ``weight_quantizer`` holding a ``fake_quant`` ``TensorQuantizer``, paired with the
+    weight named by dropping the ``_quantizer`` suffix. Matching the suffix scan matters
+    for fused and MoE modules, which expose names like ``w13_weight_quantizer`` and
+    ``gate_up_proj_weight_quantizer``: a plain ``module.weight_quantizer`` lookup would
+    miss those, leaving them folded and disabled after the stage.
+
+    Beyond the upstream scan, pairs are skipped when:
+
+    * the quantizer is disabled — its forward is the identity, so folding it is a no-op
+      that would only cost a wasted restore clone (recipes routinely disable
+      ``lm_head``/embedding quantizers, which hold ~40% of the quantized-weight bytes);
+    * the weight is not a tensor — Megatron tied-embedding ``output_layer`` modules
+      carry ``weight = None`` and borrow the embedding weight at forward time
+      (upstream ``fold_weight`` crashes on these);
+    * the quantizer is a ``SequentialQuantizer`` (W4A4 double-quant) — it subclasses
+      ``nn.Sequential``, not ``TensorQuantizer``, and upstream skips it too.
+    """
+    for module in model.modules():
+        for name in dir(module):
+            if not name.endswith(_QUANTIZER_SUFFIX):
+                continue
+            quantizer = getattr(module, name, None)
+            if (
+                not isinstance(quantizer, TensorQuantizer)
+                or not quantizer.fake_quant
+                or not quantizer.is_enabled
+            ):
+                continue
+            weight = getattr(module, name[: -len("_quantizer")], None)
+            if isinstance(weight, torch.Tensor):
+                yield weight, quantizer
+
+
+@contextlib.contextmanager
+def temporarily_fold_weights(
+    model: torch.nn.Module,
+    *,
+    verbose: bool = False,
+    rank: int = 0,
+) -> Iterator[None]:
+    """Fold fake-quantized weights into the parameters for the duration of the block.
+
+    Applies ModelOpt's fold formula (``quantizer(weight.float()).to(weight.dtype)``,
+    from ``QuantModule.fold_weight``) to each enabled weight quantizer, snapping the
+    quantized value into the existing parameter storage and disabling the quantizer,
+    then restores the original weights and re-enables the quantizers on exit. Forwards
+    inside the block read an already-snapped weight and short-circuit out of the
+    disabled quantizer, so the snap happens once per stage instead of once per
+    microbatch. Calibration state (``_amax`` / ``_pre_quant_scale``) is never touched.
+
+    Only valid while the weights are frozen — ``no_grad`` with no optimizer step —
+    which is exactly the ``get_logprobs`` re-scoring stage. Folding across a live
+    training step would corrupt the weights, since the fold is written into the
+    parameter itself.
+
+    Activation quantization is unaffected: only *weight* quantizers are disabled, so a
+    recipe's ``input_quantizer`` / ``output_quantizer`` keep running and W4A4 logprobs
+    are unchanged.
+
+    Costs one temporary copy of each folded weight shard, held only for the block.
+    """
+    folded = list(_foldable_weight_quantizers(model))
+    original_weights = [(weight, weight.detach().clone()) for weight, _ in folded]
+
+    try:
+        with torch.no_grad():
+            for weight, quantizer in folded:
+                # Exact upstream fold formula (quant_module.py::fold_weight).
+                weight.data.copy_(quantizer(weight.float()).to(weight.dtype))
+                quantizer.disable()
+        yield
+    finally:
+        with torch.no_grad():
+            for weight, original in original_weights:
+                weight.data.copy_(original)
+        for _, quantizer in folded:
+            quantizer.enable()
+    if verbose and rank == 0:
+        print(
+            f"[weight_folding] frozen-weight stage: folded and restored "
+            f"{len(folded)} weight quantizer(s)."
+        )

nemo_rl/models/policy/__init__.py

@@ -562,5 +562,13 @@ class PolicyConfig(TypedDict):
     # If true, use standard Megatron layer specs while keeping ModelOpt
     # quantization enabled. Useful for faster QARL runs and logged in configs.
     disable_modelopt_layer_spec: NotRequired[bool]
+    # Opt-in (QAT only): during the frozen-weight logprob re-scoring stage, fold the
+    # fake-quantized ("snapped") weight into the parameter once (ModelOpt's fold
+    # formula) and disable the weight quantizer, instead of re-snapping on every
+    # microbatch forward. The weights and quantizers are restored on exit. Safe because
+    # that stage runs under no_grad with no optimizer step. Costs one extra copy of the
+    # weight shard for the duration of the stage. Absent/False = disabled.
+    # Megatron QAT only.
+    quant_fold_frozen_weight_snap: NotRequired[bool]
 
     is_vlm: NotRequired[bool]

pyrefly.toml

@@ -156,6 +156,7 @@ project-includes = [
   "nemo_rl/modelopt/models/policy/workers/__init__.py",
   "nemo_rl/modelopt/models/policy/workers/dtensor_quant_policy_worker.py",
   "nemo_rl/modelopt/models/policy/workers/dtensor_quant_policy_worker_v2.py",
+  "nemo_rl/modelopt/models/policy/workers/weight_folding.py",
   "nemo_rl/modelopt/registry.py",
   "nemo_rl/models/__init__.py",
   "nemo_rl/models/automodel/__init__.py",