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

[babyplutokurt]

What does this PR do?

Adds an opt-in policy.quant_fold_frozen_weight_snap flag that eliminates redundant weight fake-quantization during the QAT logprob stages.

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.

With the flag set, the stage is wrapped in a context manager that folds each enabled fake-quant weight quantizer once, using ModelOpt's fold formula (QuantModule.fold_weight): the fake-quantized
value is written into the parameter and the weight quantizer is disabled, which is exactly the frozen-weight steady state. Forwards then read an already-snapped weight instead of re-snapping per
microbatch. Weights and quantizers are restored on exit, including on exception. Default off; when unset the path is byte-for-byte the base implementation.

Why fold per pair instead of calling mtq.fold_weight

The upstream utility selects quantizers on fake_quant alone and dereferences weight.data unconditionally (quant_module.py, line 147). On Megatron models with tied word embeddings the
output_layer is built with skip_weight_param_allocation and carries weight = None (the embedding weight is borrowed at forward time), while ModelOpt still attaches a weight_quantizer to it.
mtq.fold_weight crashes there with AttributeError: 'NoneType' object has no attribute 'data'. Reproduced end to end on Qwen3-0.6B QA-GRPO, which ties embeddings; HF never hits this because it
expresses tying as an aliased tensor.

The per-pair fold applies the identical upstream formula but:

• skips pairs whose weight is not a tensor (the tied output_layer case), and
• skips disabled quantizers, whose forward is the identity: folding them is a no-op and cloning their weights for restore wastes memory (the disabled lm_head/embedding quantizers hold roughly 40% of
  the quantized-weight bytes on Qwen3-0.6B).

Discovery uses the same *_weight_quantizer attribute-name suffix scan as fold_weight, so fused and MoE modules (w13_weight_quantizer, gate_up_proj_weight_quantizer) are folded and, critically,
restored.

Correctness and performance verification

• Unit tests (10, running 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 quantizers are restored, tied-embedding modules with weight = None are skipped, disabled quantizers are untouched.
• Mechanism check on GPU (NVFP4 W4A16 and W4A4, bf16): logits and state_dict bit-identical with and without the fold. Only the weight quantizer is disabled, so W4A4 keeps its activation quantizers
  running.
• E2E A/B on Qwen3-0.6B QA-GRPO (2 GPUs, fixed seed): both arms complete 3/3 steps; step-1 metrics identical; 112 quantizers folded and restored per policy logprob pass, 0 on the reference pass (which
  runs unquantized by design).
• Speedup scales as 1 + Wq/G (weight-quant cost over GEMM cost per forward), so the win is largest for short-sequence/small-microbatch scoring: up to ~1.7x on the logprob stage at ~1k tokens per
  forward on the smoke setup, shrinking toward 1.0x for long-sequence recipes such as the 30k-token DAPO config.

Usage

policy:
  quant_cfg: "examples/modelopt/quant_configs/nvfp4_a16.yaml"
  quant_fold_frozen_weight_snap: true

Enabled in the QA recipe configs in this PR; costs one temporary copy of each folded weight shard for the duration of the stage. Megatron QAT only.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[babyplutokurt]

Hi @terrykong

This PR improve the speed of the QAT Training, by stashing the fake quantization cache and remove unncessary per-forward wrapper. 62% speed up compare to current in-house w4a16 QAT in logprobs stage.

[terrykong]

@sharonyu-115 @mxinO could you review?

[sharonyu-115]

/ok to test 60b1e30

[mxinO]

Thanks for the cotribution, this is a good idea. But this is over complicated solution, modelopt has a fold_weight utility, this can be a simple context manager, and we don't need split to two different methods and introduce so many helpers, an example can be,

@contextmanager
def temporarily_fold_weights(model):
    modules = [
        m
        for m in model.modules()
        if hasattr(m, "weight") and hasattr(m, "weight_quantizer")
    ]
    original_weights = [(m.weight, m.weight.detach().clone()) for m in modules]
    active_quantizers = [m.weight_quantizer for m in modules if m.weight_quantizer.is_enabled]

    try:
        mtq.fold_weight(model, keep_attrs=True)
        yield
    finally:
        with torch.no_grad():
            for weight, original in original_weights:
                weight.copy_(original)
        for quantizer in active_quantizers:
            quantizer.enable()

and use it in the get_logprobs,

def get_logprobs(self, *args, **kwargs):
    if not self.cfg.get("quant_fold_frozen_weight_snap", False):
        return super().get_logprobs(*args, **kwargs)

    with temporarily_fold_weights(self.model):
        return super().get_logprobs(*args, **kwargs)

[babyplutokurt]

@mxinO Thanks, this is much better, and I've reimplemented it your way. Pushed as a single commit.

Result: Three context managers and eight helpers collapse to one context manager and one helper; two mutually-exclusive config keys (with precedence rules) collapse to a single
policy.quant_fold_frozen_weight_snap. All the monkeypatching is gone, no more patching TensorQuantizer.forward, DynamicModule.__getattr__, or MRO forward shadowing. get_logprobs is now the
four lines you sketched.

You were also right that the split was unnecessary. I'd assumed W4A4 needed its own mode, but fold_weight disables only the weight quantizer, so input_quantizer/output_quantizer keep running and
activation-quantized recipes are unaffected. Verified empirically: forward output is bit-identical folded vs. unfolded with input_quantizer enabled.

[babyplutokurt]

I've moved off mtq.fold_weight. It crashes on tied-embedding Megatron models.

While validating the fold end-to-end (QA-GRPO, Megatron + vLLM), the first get_logprobs
call with the flag enabled died inside the utility:

modelopt/torch/quantization/nn/modules/quant_module.py:147, in fold_weight
    weight.data.copy_(attr(weight.float()).to(weight.dtype))
AttributeError: 'NoneType' object has no attribute 'data'

Root cause: fold_weight selects pairs on fake_quant alone; it checks neither
is_enabled nor whether the weight is a tensor. On Megatron models with tied word
embeddings the tied output_layer is built with skip_weight_param_allocation, so it
carries weight = None and borrows the embedding weight at forward time. ModelOpt still
attaches an output_layer.weight_quantizer (disabled by every standard recipe via
'*output_layer*': enable: false, but fake_quant stays True), so the filter passes
and line 147 dereferences the None weight. Reproduced on Qwen3-0.6B
(tie_word_embeddings: true). HF models don't hit it because they express tying as a
shared tensor rather than None, which is why unit tests alone missed it.

There's no way to filter from outside: mtq.fold_weight(model) walks all modules, and
even per-module module.fold_weight() can't skip an individual pair inside a fused
module. So the context manager now applies the utility's own fold formula (its lines
147 and 148, verbatim) to each pair it discovers:

weight.data.copy_(quantizer(weight.float()).to(weight.dtype))
quantizer.disable()

with two guards the upstream walk lacks: skip weight is None, and skip disabled
quantizers. A disabled quantizer's forward is the identity, so folding it is a no-op,
and cloning its weight for restore was pure waste (the disabled lm_head/embedding
pairs hold ~40% of the quantized-weight bytes on Qwen3-0.6B). A side benefit is that
keep_attrs becomes moot, since the _amax/_pre_quant_scale deletion lives in the
code path we no longer invoke.

Same math, verified: logits are bit-identical folded vs unfolded (W4A16 and W4A4,
bf16, real ModelOpt), the full state_dict round-trips bit-identical, and the
previously crashing E2E run now completes with correct fold/restore counts. Added a
regression test that builds the tied-output_layer shape (weight=None plus an
attached quantizer).

Worth an upstream ModelOpt issue for the missing guards at quant_module.py:147; if
that lands, this module can go back to delegating.

[babyplutokurt]

Filed an upstream issue related to mtq.fold_weight: NVIDIA/Model-Optimizer#2131

[babyplutokurt]

fixed ModelOpt upstream at: NVIDIA/Model-Optimizer#2132

[babyplutokurt]

@mxinO @sharonyu-115 @terrykong ^^

[mxinO]

Thanks, the fold_weight issue should be handled in modelopt side, rather than duplicate the logic here, I will make a pr for this NVIDIA/Model-Optimizer#2140

[babyplutokurt]

Thanks @mxinO, I looked at the diff in NVIDIA/Model-Optimizer#2140 and mtq.temporarily_fold_weights looks like exactly the right upstream home for this.

Question on sequencing: since this PR is opt-in and already verified against the ModelOpt commit NeMo-RL currently pins, can we merge it first and then migrate the internals to mtq.temporarily_fold_weights in a follow-up once the pin is bumped past #2140? Or should we wait for nemo_RL track updated modelOpt that contains your PR?