NVIDIA-NeMo / NVIDIA-NeMo/RL

Avoid repeated weight fake-quantization during QAT logprob rescoring

Open
#3,440 0 comments 0 reactions 0 assignees View on GitHub
bug community-request Speed waiting-on-maintainers
Dominant language
Python
Stars
2k
Forks
561
Avg merge
4d 5h
Merged PRs (30d)
145

Description

## Summary

With ModelOpt QAT enabled, Megatron GRPO repeatedly fake-quantizes unchanged weights during `get_logprobs`.

`get_logprobs` runs without optimizer updates, but every microbatch still executes the ModelOpt quantized forward wrappers and recomputes the quantized weights. On Qwen2.5-32B-Instruct with W4A16, this makes logprob rescoring approximately 1.76x slower than bf16.

## Measurements

Configuration:

- Qwen2.5-32B-Instruct
- GRPO with the Megatron backend
- ModelOpt `nvfp4_a16`
- TP=8, PP=1, DP=4
- 32 H100-80GB GPUs
- Sequence length 16384
- Global batch size 256
- Microbatch size 1

`policy_and_reference_logprobs`, measured over two consecutive steps:

| configuration | step 1 | step 2 |
|---|---:|---:|
| bf16 | 148.43 s | 160.26 s |
| W4A16 QAT | 261.69 s (1.763x) | 281.30 s (1.755x) |

## Where the repeated work occurs

The repeated work comes from the interaction of these two paths:

1. NeMo-RL's [`MegatronPolicyWorkerImpl.get_logprobs`](https://github.com/babyplutokurt/RL/blob/60b1e3085150b8868760d32234881085a76164ed/nemo_rl/models/policy/workers/megatron_policy_worker.py#L1495-L1572) runs the forward-only microbatch loop under `torch.no_grad()`.
2. In the pinned ModelOpt revision, [`_QuantFunctionalMixin.forward`](https://github.com/NVIDIA/Model-Optimizer/blob/c3b913b9cc1d82d5a0af9fa77b4db87829e6f158/modelopt/torch/quantization/plugins/custom.py#L49-L73) installs the quantized linear functional on every forward. That functional calls [`self.weight_quantizer(weight)`](https://github.com/NVIDIA/Model-Optimizer/blob/c3b913b9cc1d82d5a0af9fa77b4db87829e6f158/modelopt/torch/quantization/nn/modules/quant_linear.py#L45-L53) each time.

The weights do not change between these microbatches, so the weight quantization result is recomputed for the same parameter throughout the stage.

## Proposed direction

Avoid recomputing frozen-weight snaps during `get_logprobs`:

- Cache each weight snap for recipes that still require the ModelOpt wrappers for activation quantization.
- For weight-only recipes, materialize each snapped weight once for the stage and bypass eligible per-forward wrappers.

The optimization must remain limited to frozen-weight evaluation. It must not change the training forward or bypass enabled activation or attention quantizers.

## Scope

Megatron ModelOpt QAT logprob rescoring only. The training forward and its straight-through gradient path are out of scope.

Contributor guide

Open the contributing guide

Research direction

Start in nemo_rl/models/policy/workers/megatron_policy_worker.py, especially MegatronPolicyWorkerImpl.get_logprobs, then inspect the referenced ModelOpt _QuantFunctionalMixin.forward and quant_linear.py paths. Reproduce the Megatron ModelOpt QAT logprob measurements if the benchmark is available, and verify that frozen-weight rescoring avoids repeated weight quantization without changing training, activation quantization, or attention quantization.

Written by the indexing model from the issue text.

Assessment

Tech stack
python, pytorch
Domain
backend, machine-learning, performance
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Quiet
Clarity
Mostly clear
Newbie friendliness
48/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.