Avoid repeated weight fake-quantization during QAT logprob rescoring
- 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
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