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

[babyplutokurt]

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.