perf(modelopt): cache fake-quantized weights across training microbatches

docs/guides/quantization-aware-rl.md

@@ -39,6 +39,59 @@ 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.
+
+## Frozen-Weight Training Optimization
+
+The same redundancy exists inside training: within one global batch, the
+gradient-accumulation microbatches all run forwards against identical weights (the
+optimizer steps once, after all of them), yet each forward re-quantizes every
+weight. Set `policy.quant_cache_train_weight_snap: true` to fake-quantize each
+weight once per global batch instead; it defaults to `false`.
+
+Folding cannot be reused here: training needs the weight quantizer in the autograd
+graph, because ModelOpt's backward is straight-through estimation that can carry an
+amax clip mask (`pass_through_bwd: false`). Instead, each enabled weight quantizer
+is patched for the duration of one `megatron_forward_backward` call to replay a
+precomputed quantized weight, with a backward that replicates ModelOpt's exactly —
+pass-through by default, or `where(|w| <= amax, grad, 0)` when the config disables
+pass-through. Forward outputs and gradients are bit-identical to the uncached
+path. The cache is rebuilt from fresh weights for every global batch, so it can
+never span an optimizer step.
+
+Parameters and quantizer state are never mutated. Quantizers whose forward chain
+this replica cannot reproduce exactly (smoothquant `pre_quant_scale`, rotation,
+static block quantization, bias quantization, calibration mode) and any call with
+a tensor other than the module's weight fall back to the original quantizer —
+correct, just not accelerated. The option costs one cached copy of each quantized
+weight shard, held for the duration of one global batch.
+
 ## Quantization-Aware GRPO (QA-GRPO)
 
 ### Configuration
@@ -51,6 +104,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 +343,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 +373,8 @@ 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. |
+| `quant_cache_train_weight_snap` | Optional boolean, default `false`. During training, fake-quantize each weight once per global batch and replay the cached value across the gradient-accumulation microbatches, with a backward replicating ModelOpt's exactly. Forward and gradients are bit-identical to the uncached path; the cache is rebuilt after every optimizer step. |
 
 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-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,10 @@
     quantize_model,
     symlink_pre_quantized_model,
 )
+from nemo_rl.modelopt.models.policy.workers.weight_folding import (
+    temporarily_cache_weight_quantization,
+    temporarily_fold_weights,
+)
 from nemo_rl.modelopt.utils import (
     MODELOPT_REAL_QUANT_ZMQ_TIMEOUT_MS,
     resolve_nvfp4_real_quant_mode,
@@ -431,6 +435,55 @@ 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 train(self, *args, **kwargs):
+        """Train, optionally caching the fake-quantized weights per global batch.
+
+        With ``policy.quant_cache_train_weight_snap`` set, each weight is fake-quantized
+        once per gradient-accumulation window instead of on every microbatch forward:
+        every ``megatron_forward_backward`` call inside the base ``train`` — one per
+        global batch, strictly between ``zero_grad`` and ``optimizer.step()`` — is
+        wrapped in :func:`temporarily_cache_weight_quantization`, so the cache is
+        rebuilt from fresh weights after every optimizer step and can never go stale.
+
+        Unlike the ``get_logprobs`` fold, the quantizer stays in the autograd graph:
+        the cached forward replicates ModelOpt's backward exactly (pass-through STE, or
+        the amax clip mask when ``pass_through_bwd`` is off), so gradients are
+        bit-identical to the uncached path. Off by default; when unset this is exactly
+        the base implementation.
+        """
+        if not self.cfg.get("quant_cache_train_weight_snap"):
+            return super().train(*args, **kwargs)
+
+        original_forward_backward = megatron_policy_worker.megatron_forward_backward
+
+        def caching_forward_backward(*fb_args, **fb_kwargs):
+            with temporarily_cache_weight_quantization(
+                self.model, verbose=True, rank=self.rank
+            ):
+                return original_forward_backward(*fb_args, **fb_kwargs)
+
+        megatron_policy_worker.megatron_forward_backward = caching_forward_backward
+        try:
+            return super().train(*args, **kwargs)
+        finally:
+            megatron_policy_worker.megatron_forward_backward = original_forward_backward
+
     def generate(self, **kwargs):
         """Quantized Megatron generation is not supported.