NVIDIA-NeMo / NVIDIA-NeMo/RL

QAT training recomputes identical weight fake-quantization on every gradient-accumulation microbatch

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

Description

Prerequisite issue and PR:

#3440
#3441

## Problem

In ModelOpt QAT the weight quantizer sits inside the linear, so every forward
recomputes `weight_quantizer(weight)`. During training, weights only change at
`optimizer.step()`, which runs once per global batch, after all
gradient-accumulation microbatches. Within one global batch, every microbatch
forward therefore fake-quantizes the exact same weights and produces the exact
same tensor. With `train_global_batch_size=32` and `train_micro_batch_size=1`
on 2 GPUs, that is 16 identical quantization passes per rank where 1 would do.

This is the training-stage counterpart of the logprobs-stage redundancy
addressed by `policy.quant_fold_frozen_weight_snap`.

## Why the existing fold cannot be reused

The logprobs fold writes the quantized value into the parameter and disables
the quantizer. That is only safe under `no_grad`: ModelOpt's fake-quant
backward is straight-through estimation that can carry an amax clip mask
(`where(|w| <= amax, grad, 0)` when a config sets `pass_through_bwd: false`).
A disabled quantizer has no backward at all, so folding during training would
silently drop the clip mask from weight gradients.

## Proposal

Cache instead of fold: precompute `Q(W)` once per global batch, patch the
quantizer forward to replay it while keeping the quantizer in the autograd
graph, and replicate ModelOpt's backward exactly (pass-through by default,
clip mask when configured). Scope the cache to one `megatron_forward_backward`
call, strictly between `zero_grad` and `optimizer.step()`, so it is rebuilt
from fresh weights after every weight update. Forward outputs and gradients
must be bit-identical to the uncached path.

Contributor guide

Open the contributing guide

Research direction

Start at the ModelOpt QAT weight quantizer inside the linear and trace one megatron_forward_backward call from zero_grad through optimizer.step. Compare cached and uncached paths, including the configured backward clip-mask behavior. Done means the cache rebuilds after each weight update while forward outputs and gradients remain bit-identical.

Written by the indexing model from the issue text.

Assessment

Tech stack
python
Domain
machine-learning, performance
Issue type
Bug
Difficulty
5/5
Estimated time
Over a week
Activity status
Quiet
Clarity
Mostly clear
Newbie friendliness
42/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.