From 9be42795c1118ee2e1672c7f906d9995c01f872c Mon Sep 17 00:00:00 2001 From: babyplutokurt Date: Fri, 7 Aug 2026 02:55:37 -0400 Subject: [PATCH 1/2] perf(modelopt): optimize frozen-weight QAT logprobs In ModelOpt QAT the weight quantizer sits inside the linear, so every forward recomputes weight_quantizer(weight). During the get_logprobs re-scoring stage the weights are frozen (no_grad, no optimizer step between microbatches), so that result is identical across every microbatch and is pure wasted work. Add an opt-in policy.quant_fold_frozen_weight_snap that wraps the stage in a context manager which folds each enabled fake-quant weight quantizer using ModelOpt's fold formula (QuantModule.fold_weight): the fake-quantized value is written into the parameter and the weight quantizer is disabled -- exactly the frozen-weight steady state. Forwards during the stage then read an already-snapped weight instead of re-snapping per microbatch. Weights and quantizers are restored on exit, including on exception. Fold per discovered pair rather than delegating to mtq.fold_weight: the upstream utility selects on fake_quant alone and dereferences weight.data unconditionally, so it crashes with AttributeError on Megatron models with tied word embeddings, where the tied output_layer exposes a weight_quantizer but carries weight = None (the embedding weight is borrowed at forward time). Verified end-to-end on Qwen3-0.6B QA-GRPO, which ties embeddings. Skip disabled quantizers during discovery: their forward is the identity, so folding them is a no-op and cloning their weights for restore is pure memory waste (the disabled lm_head/embedding quantizers in the standard recipes hold ~40% of the quantized-weight bytes). Calibration state (_amax/_pre_quant_scale) is never touched. Discover quantizers the same way fold_weight does, by the *_weight_quantizer attribute-name suffix, rather than looking up a plain module.weight_quantizer. Fused and MoE modules expose names like w13_weight_quantizer and gate_up_proj_weight_quantizer; a narrower lookup would leave them folded and disabled for the rest of training. The option applies to any quantization format. Only the weight quantizer is disabled, so activation-quantized recipes (W4A4) keep their input/output quantizers running and produce unchanged logprobs. Default off, and scoped to the no-grad frozen-weight stage. Includes config examples, documentation, and unit coverage that exercises the real ModelOpt library: forward output is bit-identical folded vs unfolded, weights restore through their original storage, amax survives, restoration holds on exception, fused *_weight_quantizer modules are restored, tied-embedding modules with weight=None are skipped, and disabled quantizers are left untouched. Signed-off-by: babyplutokurt --- docs/guides/quantization-aware-rl.md | 31 +++ ...-megatron-qa-nvfp4-w4a16-real-300step.yaml | 1 + ...g-megatron-qa-nvfp4-w4a4-real-300step.yaml | 1 + ...ba3b-4n4g-megatron-qa-nvfp4-w4a4-real.yaml | 1 + ...apo-2n8g-long-megatron-qa-nvfp4-w4a16.yaml | 1 + .../qa_distillation_math_megatron.yaml | 1 + .../modelopt/qa_grpo_llama8b_megatron.v2.yaml | 1 + examples/modelopt/qa_grpo_math_megatron.yaml | 1 + examples/modelopt/qa_grpo_nano3_megatron.yaml | 1 + .../qa_grpo_qwen3_30ba3b_megatron.yaml | 1 + .../workers/megatron_quant_policy_worker.py | 19 ++ .../models/policy/workers/weight_folding.py | 140 +++++++++++ nemo_rl/models/policy/__init__.py | 8 + pyrefly.toml | 1 + .../unit/models/policy/test_weight_folding.py | 238 ++++++++++++++++++ 15 files changed, 446 insertions(+) create mode 100644 nemo_rl/modelopt/models/policy/workers/weight_folding.py create mode 100644 tests/unit/models/policy/test_weight_folding.py diff --git a/docs/guides/quantization-aware-rl.md b/docs/guides/quantization-aware-rl.md index ea846c3d44..0eb134d067 100644 --- a/docs/guides/quantization-aware-rl.md +++ b/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. 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-30ba3b-4n4g-megatron-qa-nvfp4-w4a4-real.yaml b/examples/configs/recipes/llm/grpo-qwen3-30ba3b-4n4g-megatron-qa-nvfp4-w4a4-real.yaml index d324a5a47d..45b2b7280d 100644 --- a/examples/configs/recipes/llm/grpo-qwen3-30ba3b-4n4g-megatron-qa-nvfp4-w4a4-real.yaml +++ b/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 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..5684bc6f59 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,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. 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..79976d948a --- /dev/null +++ b/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)." + ) diff --git a/nemo_rl/models/policy/__init__.py b/nemo_rl/models/policy/__init__.py index 7994f37d7b..5d2cb32f0d 100644 --- a/nemo_rl/models/policy/__init__.py +++ b/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] 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..afcf8bf3c1 --- /dev/null +++ b/tests/unit/models/policy/test_weight_folding.py @@ -0,0 +1,238 @@ +# 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_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) + + +@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"]) From e7a80ba23e56870b4ffccb3eb94bd820d6251abb Mon Sep 17 00:00:00 2001 From: babyplutokurt Date: Sun, 9 Aug 2026 04:03:50 -0400 Subject: [PATCH 2/2] perf(modelopt): cache fake-quantized weights across training microbatches The frozen-weight redundancy addressed for get_logprobs by quant_fold_frozen_weight_snap also 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 every microbatch forward recomputes weight_quantizer(weight). Add an opt-in policy.quant_cache_train_weight_snap that fake-quantizes each weight once per global batch. Folding cannot be reused here because training needs the weight quantizer in the autograd graph: ModelOpt's backward is straight-through estimation that can carry an amax clip mask (pass_through_bwd=false), which a disabled quantizer would silently drop. Instead, each enabled weight quantizer's forward is patched for the duration of one megatron_forward_backward call to replay a precomputed quantized weight, with a backward replicating 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 (unit-verified for INT8 per-channel and NVFP4 dynamic block quantization, in both backward modes). The wrap happens per megatron_forward_backward call -- one per global batch, strictly between zero_grad and optimizer.step() -- so the cache is rebuilt from fresh weights after every optimizer step and can never go stale. Parameters and quantizer state are never mutated; the patch is an instance-level forward override removed on exit, including on exception. Quantizers whose forward chain the replica cannot reproduce exactly (smoothquant pre_quant_scale, rotation, static block quantization, bias quantization, calibration mode) and calls with any tensor other than the module's weight (e.g. refit exports quantizing .float() copies) fall back to the original quantizer: correct, just not accelerated. Default off. Costs one cached copy of each quantized weight shard, held for the duration of one global batch. Includes documentation and unit coverage against the real ModelOpt library. Also minimizes the w4a4-real recipe YAML (drops a quant_fold_frozen_weight_snap already set by its parent config), flagged by the configs-minimize-check hook. Signed-off-by: babyplutokurt --- docs/guides/quantization-aware-rl.md | 26 ++ ...ba3b-4n4g-megatron-qa-nvfp4-w4a4-real.yaml | 1 - .../workers/megatron_quant_policy_worker.py | 34 +++ .../models/policy/workers/weight_folding.py | 169 +++++++++++- nemo_rl/models/policy/__init__.py | 9 + .../unit/models/policy/test_weight_folding.py | 246 ++++++++++++++++++ 6 files changed, 478 insertions(+), 7 deletions(-) diff --git a/docs/guides/quantization-aware-rl.md b/docs/guides/quantization-aware-rl.md index 0eb134d067..9ddaec51bd 100644 --- a/docs/guides/quantization-aware-rl.md +++ b/docs/guides/quantization-aware-rl.md @@ -67,6 +67,31 @@ 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 @@ -349,6 +374,7 @@ These parameters are added under the `policy` section: | `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-qwen3-30ba3b-4n4g-megatron-qa-nvfp4-w4a4-real.yaml b/examples/configs/recipes/llm/grpo-qwen3-30ba3b-4n4g-megatron-qa-nvfp4-w4a4-real.yaml index 45b2b7280d..d324a5a47d 100644 --- a/examples/configs/recipes/llm/grpo-qwen3-30ba3b-4n4g-megatron-qa-nvfp4-w4a4-real.yaml +++ b/examples/configs/recipes/llm/grpo-qwen3-30ba3b-4n4g-megatron-qa-nvfp4-w4a4-real.yaml @@ -12,7 +12,6 @@ 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 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 5684bc6f59..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 @@ -41,6 +41,7 @@ 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 ( @@ -450,6 +451,39 @@ def get_logprobs(self, *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 index 79976d948a..501a32b8b1 100644 --- a/nemo_rl/modelopt/models/policy/workers/weight_folding.py +++ b/nemo_rl/modelopt/models/policy/workers/weight_folding.py @@ -12,13 +12,19 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Temporarily fold fake-quantized weights for frozen-weight QAT stages. +"""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. -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. +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 @@ -41,10 +47,10 @@ """ import contextlib -from collections.abc import Iterator +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" @@ -138,3 +144,154 @@ def temporarily_fold_weights( 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 5d2cb32f0d..34e5b25ad5 100644 --- a/nemo_rl/models/policy/__init__.py +++ b/nemo_rl/models/policy/__init__.py @@ -571,4 +571,13 @@ class PolicyConfig(TypedDict): # 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/tests/unit/models/policy/test_weight_folding.py b/tests/unit/models/policy/test_weight_folding.py index afcf8bf3c1..5142089f5e 100644 --- a/tests/unit/models/policy/test_weight_folding.py +++ b/tests/unit/models/policy/test_weight_folding.py @@ -29,6 +29,7 @@ ) from nemo_rl.modelopt.models.policy.workers.weight_folding import ( # noqa: E402 + temporarily_cache_weight_quantization, temporarily_fold_weights, ) @@ -184,6 +185,251 @@ def test_nothing_to_fold_is_a_noop() -> None: 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", [