QAT training recomputes identical weight fake-quantization on every gradient-accumulation microbatch
- 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
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