diff --git a/docs/guides/quantization-aware-rl.md b/docs/guides/quantization-aware-rl.md index ea846c3d44..9ddaec51bd 100644 --- a/docs/guides/quantization-aware-rl.md +++ b/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. diff --git a/examples/configs/recipes/llm/grpo-nemotron3-super-120BA12B-16n4g-megatron-qa-nvfp4-w4a16-real-300step.yaml b/examples/configs/recipes/llm/grpo-nemotron3-super-120BA12B-16n4g-megatron-qa-nvfp4-w4a16-real-300step.yaml index 67fd6a66d2..f3fae1f2fc 100644 --- a/examples/configs/recipes/llm/grpo-nemotron3-super-120BA12B-16n4g-megatron-qa-nvfp4-w4a16-real-300step.yaml +++ b/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 diff --git a/examples/configs/recipes/llm/grpo-nemotron3-super-120BA12B-16n4g-megatron-qa-nvfp4-w4a4-real-300step.yaml b/examples/configs/recipes/llm/grpo-nemotron3-super-120BA12B-16n4g-megatron-qa-nvfp4-w4a4-real-300step.yaml index 9dec458278..dbe37e0c69 100644 --- a/examples/configs/recipes/llm/grpo-nemotron3-super-120BA12B-16n4g-megatron-qa-nvfp4-w4a4-real-300step.yaml +++ b/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 diff --git a/examples/configs/recipes/llm/grpo-qwen3-8b-base-dapo-2n8g-long-megatron-qa-nvfp4-w4a16.yaml b/examples/configs/recipes/llm/grpo-qwen3-8b-base-dapo-2n8g-long-megatron-qa-nvfp4-w4a16.yaml index 1eaa89ff47..19d8d28930 100644 --- a/examples/configs/recipes/llm/grpo-qwen3-8b-base-dapo-2n8g-long-megatron-qa-nvfp4-w4a16.yaml +++ b/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 diff --git a/examples/modelopt/qa_distillation_math_megatron.yaml b/examples/modelopt/qa_distillation_math_megatron.yaml index 38a3413c71..6ba9e939a8 100644 --- a/examples/modelopt/qa_distillation_math_megatron.yaml +++ b/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 diff --git a/examples/modelopt/qa_grpo_llama8b_megatron.v2.yaml b/examples/modelopt/qa_grpo_llama8b_megatron.v2.yaml index 756deac7dc..e958d375f9 100644 --- a/examples/modelopt/qa_grpo_llama8b_megatron.v2.yaml +++ b/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 diff --git a/examples/modelopt/qa_grpo_math_megatron.yaml b/examples/modelopt/qa_grpo_math_megatron.yaml index 5f5c982730..4128500bae 100644 --- a/examples/modelopt/qa_grpo_math_megatron.yaml +++ b/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 diff --git a/examples/modelopt/qa_grpo_nano3_megatron.yaml b/examples/modelopt/qa_grpo_nano3_megatron.yaml index ee2fe54fd8..32f519b64d 100644 --- a/examples/modelopt/qa_grpo_nano3_megatron.yaml +++ b/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 diff --git a/examples/modelopt/qa_grpo_qwen3_30ba3b_megatron.yaml b/examples/modelopt/qa_grpo_qwen3_30ba3b_megatron.yaml index 8c9a3fac7e..4e2adee790 100644 --- a/examples/modelopt/qa_grpo_qwen3_30ba3b_megatron.yaml +++ b/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 diff --git a/nemo_rl/modelopt/models/policy/workers/megatron_quant_policy_worker.py b/nemo_rl/modelopt/models/policy/workers/megatron_quant_policy_worker.py index 29f7dd505a..4f2cf6b864 100644 --- a/nemo_rl/modelopt/models/policy/workers/megatron_quant_policy_worker.py +++ b/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. diff --git a/nemo_rl/modelopt/models/policy/workers/weight_folding.py b/nemo_rl/modelopt/models/policy/workers/weight_folding.py new file mode 100644 index 0000000000..501a32b8b1 --- /dev/null +++ b/nemo_rl/modelopt/models/policy/workers/weight_folding.py @@ -0,0 +1,297 @@ +# 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. + +"""Reuse fake-quantized weights across 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. +That result only changes when the weight changes, i.e. at an optimizer step, so any +stretch of forwards between weight updates recomputes the identical tensor: + +* the ``get_logprobs`` re-scoring stage (``torch.no_grad()``, no optimizer step at all) + — served by :func:`temporarily_fold_weights`; +* the gradient-accumulation microbatches of one training global batch (the optimizer + steps once, *after* all of them) — served by + :func:`temporarily_cache_weight_quantization`, which must keep the quantizer in the + autograd graph and therefore caches instead of folding. + +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 Callable, Iterator +from typing import Any, cast + +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)." + ) + + +def _cache_eligible(quantizer: TensorQuantizer) -> bool: + """Whether ``quantizer``'s steady-state forward/backward can be replicated exactly. + + The cached forward replays a precomputed output and replicates the backward of + ``TensorQuantizer._fake_quantize`` (pass-through STE, or the amax clip mask when + ``pass_through_bwd`` is disabled). Any quantizer feature that adds other terms to + the forward chain — smoothquant ``pre_quant_scale``, input rotation, the + static-block reshape, bias, calibration collection — would make that replica + wrong, so such quantizers are left on their original forward (correct, just not + accelerated). + """ + return ( + type(quantizer) is TensorQuantizer + and quantizer._if_quant + and not quantizer._if_calib + and getattr(quantizer, "pre_quant_scale", None) is None + and not getattr(quantizer, "rotate_is_enabled", False) + and not getattr(quantizer, "is_static_block_quant", True) + and getattr(quantizer, "bias_calibrator", None) is None + and hasattr(quantizer, "_get_amax") + ) + + +def _make_cached_forward( + quantizer: TensorQuantizer, + weight: torch.Tensor, + cached: torch.Tensor, + amax: torch.Tensor | None, + original_forward: Callable[[torch.Tensor], torch.Tensor], + stats: dict[str, int], +) -> Callable[[torch.Tensor], torch.Tensor]: + """Build a replacement ``forward`` that replays ``cached`` for ``weight``. + + Exactness contract, mirroring ``TensorQuantizer._fake_quantize``: + + * forward: the module's quantized-weight tensor is ``cached``, computed once via + the quantizer's own forward — bit-identical by construction. + * backward: upstream saves ``(inputs, amax)`` only when ``pass_through_bwd`` is + off and amax exists, and then applies ``where(|inputs| <= amax, grad, 0)`` + (``_fake_quant_backward_function``); otherwise the gradient passes through + unchanged. ``amax`` here is pre-resolved to ``None`` for the pass-through case. + + Any call that is not exactly "the same weight storage, quantization enabled" falls + back to the original forward — refit paths call weight quantizers on ``.float()`` + copies, and ``disable_quantization`` can flip state mid-window. + """ + w_ptr = weight.data_ptr() + w_meta = (weight.shape, weight.stride(), weight.dtype, weight.device) + + class _CachedWeightFakeQuant(torch.autograd.Function): + @staticmethod + def forward( # pyrefly: ignore[bad-override] Always ignore torch.autograd.Function.forward's type since it's always more specific than the base class + ctx: Any, + inputs: torch.Tensor, + ) -> torch.Tensor: + if amax is not None: + ctx.save_for_backward(inputs) + return cached.view_as(cached) + + @staticmethod + def backward(ctx: Any, *grad_outputs: torch.Tensor) -> torch.Tensor: + grad = grad_outputs[0] + if not ctx.saved_tensors: + return grad + (inputs,) = ctx.saved_tensors + # Exact upstream clip-mask STE (tensor_quant.py::_fake_tensor_quant_backward). + zero = grad.new_zeros(1) + return torch.where(inputs.abs() <= amax, grad, zero) + + def cached_forward(inputs): + if ( + isinstance(inputs, torch.Tensor) + and inputs.data_ptr() == w_ptr + and (inputs.shape, inputs.stride(), inputs.dtype, inputs.device) == w_meta + and quantizer._if_quant + and not quantizer._if_calib + and quantizer.is_enabled + ): + stats["hits"] += 1 + return _CachedWeightFakeQuant.apply(inputs) + stats["misses"] += 1 + return original_forward(inputs) + + return cached_forward + + +@contextlib.contextmanager +def temporarily_cache_weight_quantization( + model: torch.nn.Module, + *, + verbose: bool = False, + rank: int = 0, +) -> Iterator[None]: + """Serve ``weight_quantizer(weight)`` from a per-stage cache, exact in both passes. + + For each enabled fake-quant weight quantizer, computes the quantized weight once + (via the quantizer's own forward) and patches the quantizer to replay it, with a + backward that replicates ModelOpt's exactly: pass-through STE by default, or the + ``where(|w| <= amax, grad, 0)`` clip mask when the quantizer sets + ``pass_through_bwd=False``. Both directions are bit-identical to the unpatched + quantizer, so this is safe around *training* forward-backward passes — unlike + :func:`temporarily_fold_weights`, which disables the quantizer and thereby drops + the clip mask from the gradient. + + Only valid while the weights are frozen: the cache must be rebuilt after every + optimizer step, so wrap exactly one gradient-accumulation window (one + ``megatron_forward_backward`` call), never a loop that steps the optimizer. + + Never mutates parameters or quantizer state; the patch is an instance-level + ``forward`` override removed on exit (exception-safe). Calls with any other + tensor, or after the quantizer is disabled mid-window, fall back to the original + forward. Costs one cached copy of each quantized weight shard for the window. + """ + stats = {"hits": 0, "misses": 0} + patched: list[TensorQuantizer] = [] + try: + with torch.no_grad(): + for weight, quantizer in _foldable_weight_quantizers(model): + if not _cache_eligible(quantizer) or "forward" in vars(quantizer): + continue + original_forward = quantizer.forward + try: + cached = original_forward(weight).detach() + amax: torch.Tensor | None = ( + None + if quantizer.is_mx_format + or getattr(quantizer, "_pass_through_bwd", True) + else cast("torch.Tensor | None", quantizer._get_amax(weight)) + ) + except Exception: + continue # unexpected quantizer flavor: leave it unpatched + object.__setattr__( + quantizer, + "forward", + _make_cached_forward( + quantizer, weight, cached, amax, original_forward, stats + ), + ) + patched.append(quantizer) + yield + finally: + for quantizer in patched: + object.__delattr__(quantizer, "forward") + if verbose and rank == 0: + print( + f"[weight_folding] cached {len(patched)} weight quantizer(s) for the " + f"frozen-weight window: {stats['hits']} cache hits, " + f"{stats['misses']} fallback calls." + ) diff --git a/nemo_rl/models/policy/__init__.py b/nemo_rl/models/policy/__init__.py index 7994f37d7b..34e5b25ad5 100644 --- a/nemo_rl/models/policy/__init__.py +++ b/nemo_rl/models/policy/__init__.py @@ -562,5 +562,22 @@ 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] + + # Cache each weight's fake-quantized value across the gradient-accumulation + # microbatches of one training global batch (weights only change at the + # optimizer step, which runs after all of them). The cached forward keeps the + # quantizer in the autograd graph and replicates ModelOpt's backward exactly, + # so forward and gradients are bit-identical to the uncached path. Costs one + # cached copy of each quantized weight shard per window. Absent/False = + # disabled. Megatron QAT only. + quant_cache_train_weight_snap: NotRequired[bool] is_vlm: NotRequired[bool] diff --git a/pyrefly.toml b/pyrefly.toml index 5f55643e09..aa78896522 100644 --- a/pyrefly.toml +++ b/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", diff --git a/tests/unit/models/policy/test_weight_folding.py b/tests/unit/models/policy/test_weight_folding.py new file mode 100644 index 0000000000..5142089f5e --- /dev/null +++ b/tests/unit/models/policy/test_weight_folding.py @@ -0,0 +1,484 @@ +# 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. + +import contextlib +import copy +from collections.abc import Iterator +from contextlib import AbstractContextManager + +import pytest +import torch +import torch.nn as nn + +pytestmark = pytest.mark.mcore + +mtq = pytest.importorskip( + "modelopt.torch.quantization", + reason="Requires the nvidia-modelopt package", +) + +from nemo_rl.modelopt.models.policy.workers.weight_folding import ( # noqa: E402 + temporarily_cache_weight_quantization, + temporarily_fold_weights, +) + + +def _quantized_model() -> tuple[nn.Module, torch.Tensor]: + """Build a small INT8 fake-quantized model and the input it was calibrated on.""" + torch.manual_seed(0) + model = nn.Sequential( + nn.Linear(16, 16, bias=False), + nn.Linear(16, 16, bias=False), + ) + inputs = torch.randn(4, 16) + mtq.quantize(model, mtq.INT8_DEFAULT_CFG, lambda m: m(inputs)) + return model, inputs + + +def test_folding_does_not_change_forward_output() -> None: + """The whole point: folded forwards must produce identical logits.""" + model, inputs = _quantized_model() + with torch.no_grad(): + expected = model(inputs) + + with temporarily_fold_weights(model): + folded = model(inputs) + + restored = model(inputs) + + assert torch.equal(expected, folded), "folding changed the forward output" + assert torch.equal(expected, restored), "restore changed the forward output" + + +def test_folds_inside_and_restores_weights_and_quantizers() -> None: + model, _ = _quantized_model() + linear = model[0] + original_weight = linear.weight.detach().clone() + original_storage = linear.weight.data_ptr() + original_amax = linear.weight_quantizer.amax.detach().clone() + + with temporarily_fold_weights(model): + # Inside: weight carries the snapped value and the quantizer steps aside. + assert not torch.equal(original_weight, linear.weight) + assert not linear.weight_quantizer.is_enabled + + assert torch.equal(original_weight, linear.weight) + assert linear.weight.data_ptr() == original_storage, ( + "restore must write through existing parameter storage" + ) + assert linear.weight_quantizer.is_enabled + # keep_attrs=True must have preserved calibration state. + assert linear.weight_quantizer.amax is not None + assert torch.equal(original_amax, linear.weight_quantizer.amax) + + +def test_restores_after_exception_in_body() -> None: + model, _ = _quantized_model() + linear = model[0] + original_weight = linear.weight.detach().clone() + + with pytest.raises(RuntimeError, match="boom"): + with temporarily_fold_weights(model): + raise RuntimeError("boom") + + assert torch.equal(original_weight, linear.weight) + assert linear.weight_quantizer.is_enabled + + +def test_restores_fused_weight_quantizers() -> None: + """Regression: MoE/fused modules expose e.g. ``w13_weight_quantizer``. + + ``fold_weight`` folds any ``*_weight_quantizer``, so restoring only a plain + ``module.weight_quantizer`` would leave these folded and disabled for the rest of + training. + """ + model, _ = _quantized_model() + linear = model[0] + linear.register_parameter("w13_weight", nn.Parameter(torch.randn(16, 16))) + linear.add_module("w13_weight_quantizer", copy.deepcopy(linear.weight_quantizer)) + + original_fused = linear.w13_weight.detach().clone() + + with temporarily_fold_weights(model): + assert not torch.equal(original_fused, linear.w13_weight) + assert not linear.w13_weight_quantizer.is_enabled + + assert torch.equal(original_fused, linear.w13_weight), ( + "fused weight was folded but never restored" + ) + assert linear.w13_weight_quantizer.is_enabled + + +def test_quantizer_disabled_beforehand_is_skipped_and_stays_disabled() -> None: + """Disabled quantizers are identity at forward time, so folding them is a no-op. + + They must be skipped entirely: no wasted restore clone, weight untouched, and the + disabled state preserved through the block. + """ + model, _ = _quantized_model() + model[0].weight_quantizer.disable() + disabled_weight = model[0].weight.detach().clone() + + with temporarily_fold_weights(model): + # Skipped: the disabled quantizer's weight is never folded. + assert torch.equal(disabled_weight, model[0].weight) + assert not model[0].weight_quantizer.is_enabled + # The enabled sibling still folds. + assert not model[1].weight_quantizer.is_enabled + + assert not model[0].weight_quantizer.is_enabled + assert model[1].weight_quantizer.is_enabled + + +def test_none_weight_with_quantizer_is_skipped() -> None: + """Regression: Megatron tied-embedding ``output_layer`` has ``weight = None``. + + ModelOpt attaches a ``weight_quantizer`` to it anyway, and upstream + ``mtq.fold_weight`` crashes with ``AttributeError: 'NoneType' object has no + attribute 'data'`` on such modules (observed on Qwen3-0.6B, which ties word + embeddings). The fold must skip the pair and still fold everything else. + """ + model, inputs = _quantized_model() + + class TiedOutputLayer(nn.Module): + def __init__(self, template: nn.Module) -> None: + super().__init__() + self.weight = None # borrowed from the embedding at forward time + self.add_module( + "weight_quantizer", copy.deepcopy(template.weight_quantizer) + ) + + model.tied_head = TiedOutputLayer(model[0]) + assert model.tied_head.weight_quantizer.is_enabled + + with torch.no_grad(): + expected = model[0](inputs) + + with temporarily_fold_weights(model): # must not raise + assert not model[0].weight_quantizer.is_enabled + with torch.no_grad(): + folded = model[0](inputs) + + assert torch.equal(expected, folded) + assert model.tied_head.weight_quantizer.is_enabled + + +def test_nothing_to_fold_is_a_noop() -> None: + """An unquantized model has no weight quantizers; the context must still work.""" + model = nn.Sequential(nn.Linear(8, 8)) + original = model[0].weight.detach().clone() + + with temporarily_fold_weights(model): + pass + + assert torch.equal(original, model[0].weight) + + +# --------------------------------------------------------------------------- +# temporarily_cache_weight_quantization (training-stage cache) +# --------------------------------------------------------------------------- + + +def _grads(model: nn.Module, inputs: torch.Tensor) -> list[torch.Tensor]: + """One fwd/bwd pass; returns detached copies of all parameter gradients.""" + model.zero_grad() + model(inputs).square().sum().backward() + return [p.grad.detach().clone() for p in model.parameters()] + + +def test_cache_forward_is_bit_identical_and_actually_served_from_cache() -> None: + model, inputs = _quantized_model() + with torch.no_grad(): + expected = model(inputs) + + original_amax = model[0].weight_quantizer.amax.detach().clone() + with temporarily_cache_weight_quantization(model): + assert "forward" in vars(model[0].weight_quantizer) + with torch.no_grad(): + cached_out = model(inputs) + # Corrupt amax in-place: the *real* quantizer would now produce a + # different value, so an unchanged output proves the cache is served. + model[0].weight_quantizer.amax.copy_(original_amax * 100) + still_cached = model(inputs) + model[0].weight_quantizer.amax.copy_(original_amax) + + assert "forward" not in vars(model[0].weight_quantizer) + with torch.no_grad(): + restored = model(inputs) + + assert torch.equal(expected, cached_out), "caching changed the forward output" + assert torch.equal(expected, still_cached), "forward bypassed the cache" + assert torch.equal(expected, restored), "exit did not restore the real quantizer" + + +def test_cache_gradients_bit_identical_pass_through() -> None: + """Default ModelOpt configs use pass-through STE; grads must match bitwise.""" + model, inputs = _quantized_model() + expected = _grads(model, inputs) + + with temporarily_cache_weight_quantization(model): + got = _grads(model, inputs) + after = _grads(model, inputs) + + for e, g, a in zip(expected, got, after): + assert torch.equal(e, g), "cached backward diverged from the real quantizer" + assert torch.equal(e, a), "gradients changed after cache exit" + + +def test_cache_gradients_bit_identical_with_active_clip_mask() -> None: + """With ``pass_through_bwd=False`` ModelOpt clips grads at amax; replicate it.""" + model, inputs = _quantized_model() + for layer in model: + wq = layer.weight_quantizer + wq._pass_through_bwd = False + with torch.no_grad(): + wq.amax.copy_(wq.amax * 0.5) # force a non-trivial clip mask + + clipped = (model[0].weight.abs() > model[0].weight_quantizer.amax).sum() + assert clipped > 0, "test setup failed to activate the clip mask" + + expected = _grads(model, inputs) + assert (expected[0] == 0).sum() >= clipped, "upstream clip mask not in effect" + + with temporarily_cache_weight_quantization(model): + got = _grads(model, inputs) + + for e, g in zip(expected, got): + assert torch.equal(e, g), "cached clip-mask backward diverged" + + +def test_cache_falls_back_for_non_weight_tensors() -> None: + """Refit paths call weight quantizers on ``.float()`` copies; those must not + be served the cached (bf16-shaped) value.""" + model, inputs = _quantized_model() + wq = model[0].weight_quantizer + float_weight = model[0].weight.detach().float() * 0.25 + with torch.no_grad(): + expected = wq(float_weight) + + with temporarily_cache_weight_quantization(model): + with torch.no_grad(): + got = wq(float_weight) + + assert torch.equal(expected, got), "cache served a stale value for a foreign tensor" + + +def test_cache_restores_forward_on_exception() -> None: + model, inputs = _quantized_model() + + with pytest.raises(RuntimeError, match="boom"): + with temporarily_cache_weight_quantization(model): + raise RuntimeError("boom") + + assert "forward" not in vars(model[0].weight_quantizer) + assert "forward" not in vars(model[1].weight_quantizer) + + +def test_cache_rebuilds_from_fresh_weights_per_window() -> None: + """Simulates the per-global-batch usage: weight update between windows.""" + model, inputs = _quantized_model() + + with temporarily_cache_weight_quantization(model): + pass + + with torch.no_grad(): # "optimizer step" + model[0].weight.mul_(1.5) + expected = model(inputs) + + with temporarily_cache_weight_quantization(model): + with torch.no_grad(): + got = model(inputs) + + assert torch.equal(expected, got), "cache went stale across a weight update" + + +def test_cache_skips_ineligible_and_disabled_quantizers() -> None: + model, inputs = _quantized_model() + # pre_quant_scale adds a term to the forward chain the replica cannot + # reproduce; such quantizers must keep their original forward. (Scale of + # ones keeps the reference output unchanged.) + model[0].weight_quantizer._enable_pre_quant_scale = True + model[0].weight_quantizer.pre_quant_scale = torch.ones_like(model[0].weight[0]) + model[1].weight_quantizer.disable() + with torch.no_grad(): + expected = model(inputs) + + with temporarily_cache_weight_quantization(model): + assert "forward" not in vars(model[0].weight_quantizer) + assert "forward" not in vars(model[1].weight_quantizer) + with torch.no_grad(): + got = model(inputs) + + assert torch.equal(expected, got) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="NVFP4 needs CUDA") +@pytest.mark.parametrize("pass_through_bwd", [True, False]) +def test_cache_gradients_bit_identical_nvfp4(pass_through_bwd: bool) -> None: + """The shipped recipes use NVFP4 dynamic block quantization on GPU.""" + torch.manual_seed(0) + model = nn.Sequential( + nn.Linear(32, 32, bias=False), + nn.Linear(32, 32, bias=False), + ).cuda() + inputs = torch.randn(4, 32, device="cuda") + mtq.quantize(model, mtq.NVFP4_DEFAULT_CFG, lambda m: m(inputs)) + for layer in model: + layer.weight_quantizer._pass_through_bwd = pass_through_bwd + + with torch.no_grad(): + expected_out = model(inputs) + expected = _grads(model, inputs) + + with temporarily_cache_weight_quantization(model): + assert "forward" in vars(model[0].weight_quantizer) + with torch.no_grad(): + got_out = model(inputs) + got = _grads(model, inputs) + + assert torch.equal(expected_out, got_out) + for e, g in zip(expected, got): + assert torch.equal(e, g), "NVFP4 cached backward diverged" + + +@pytest.mark.parametrize( + "config, expects_cache", + [ + ({}, False), + ({"quant_cache_train_weight_snap": False}, False), + ({"quant_cache_train_weight_snap": True}, True), + ], +) +def test_quant_worker_routes_train_through_cache( + monkeypatch: pytest.MonkeyPatch, + config: dict[str, bool], + expects_cache: bool, +) -> None: + worker_module = pytest.importorskip( + "nemo_rl.modelopt.models.policy.workers.megatron_quant_policy_worker", + reason="Requires Megatron and Ray", + ) + base_module = worker_module.megatron_policy_worker + + events: list[str] = [] + + def recording_context( + *args: object, **kwargs: object + ) -> AbstractContextManager[None]: + @contextlib.contextmanager + def manager() -> Iterator[None]: + events.append("cache_enter") + try: + yield + finally: + events.append("cache_exit") + + return manager() + + def fake_forward_backward(*args: object, **kwargs: object) -> str: + events.append("fwd_bwd") + return "losses" + + def base_train(self: object, *args: object, **kwargs: object) -> str: + # The real train() resolves megatron_forward_backward from the base + # module's globals once per global batch. + assert base_module.megatron_forward_backward(...) == "losses" + assert base_module.megatron_forward_backward(...) == "losses" + events.append("train_done") + return "result" + + monkeypatch.setattr( + worker_module, "temporarily_cache_weight_quantization", recording_context + ) + monkeypatch.setattr(base_module, "megatron_forward_backward", fake_forward_backward) + monkeypatch.setattr(worker_module.MegatronPolicyWorkerImpl, "train", base_train) + + worker_class = ( + worker_module.MegatronQuantPolicyWorker.__ray_metadata__.modified_class + ) + worker = object.__new__(worker_class) + worker.cfg = config + worker.model = object() + worker.rank = 0 + + assert worker.train() == "result" + if expects_cache: + # One cache window per forward-backward call, none spanning both. + assert events == [ + "cache_enter", + "fwd_bwd", + "cache_exit", + "cache_enter", + "fwd_bwd", + "cache_exit", + "train_done", + ] + else: + assert events == ["fwd_bwd", "fwd_bwd", "train_done"] + # The scoped patch must be unwound after train() returns. + assert base_module.megatron_forward_backward is fake_forward_backward + + +@pytest.mark.parametrize( + "config, expects_fold", + [ + ({}, False), + ({"quant_fold_frozen_weight_snap": False}, False), + ({"quant_fold_frozen_weight_snap": True}, True), + ], +) +def test_quant_worker_routes_logprobs_through_fold( + monkeypatch: pytest.MonkeyPatch, + config: dict[str, bool], + expects_fold: bool, +) -> None: + worker_module = pytest.importorskip( + "nemo_rl.modelopt.models.policy.workers.megatron_quant_policy_worker", + reason="Requires Megatron and Ray", + ) + + events: list[str] = [] + + def recording_context( + *args: object, **kwargs: object + ) -> AbstractContextManager[None]: + @contextlib.contextmanager + def manager() -> Iterator[None]: + events.append("fold_enter") + try: + yield + finally: + events.append("fold_exit") + + return manager() + + def base_get_logprobs(self: object, *args: object, **kwargs: object) -> str: + events.append("base") + return "result" + + monkeypatch.setattr(worker_module, "temporarily_fold_weights", recording_context) + monkeypatch.setattr( + worker_module.MegatronPolicyWorkerImpl, "get_logprobs", base_get_logprobs + ) + + worker_class = ( + worker_module.MegatronQuantPolicyWorker.__ray_metadata__.modified_class + ) + worker = object.__new__(worker_class) + worker.cfg = config + worker.model = object() + worker.rank = 0 + + assert worker.get_logprobs() == "result" + assert events == (["fold_enter", "base", "fold_exit"] if expects_fold else ["base"])