perf(modelopt): cache fake-quantized weights across training microbatches - #3556
Open
babyplutokurt wants to merge 2 commits into
Open
perf(modelopt): cache fake-quantized weights across training microbatches#3556babyplutokurt wants to merge 2 commits into
babyplutokurt wants to merge 2 commits into
Conversation
In ModelOpt QAT the weight quantizer sits inside the linear, so every forward recomputes weight_quantizer(weight). During the get_logprobs re-scoring stage the weights are frozen (no_grad, no optimizer step between microbatches), so that result is identical across every microbatch and is pure wasted work. Add an opt-in policy.quant_fold_frozen_weight_snap that wraps the stage in a context manager which folds each enabled fake-quant weight quantizer using ModelOpt's fold formula (QuantModule.fold_weight): the fake-quantized value is written into the parameter and the weight quantizer is disabled -- exactly the frozen-weight steady state. Forwards during the stage then read an already-snapped weight instead of re-snapping per microbatch. Weights and quantizers are restored on exit, including on exception. Fold per discovered pair rather than delegating to mtq.fold_weight: the upstream utility selects on fake_quant alone and dereferences weight.data unconditionally, so it crashes with AttributeError on Megatron models with tied word embeddings, where the tied output_layer exposes a weight_quantizer but carries weight = None (the embedding weight is borrowed at forward time). Verified end-to-end on Qwen3-0.6B QA-GRPO, which ties embeddings. Skip disabled quantizers during discovery: their forward is the identity, so folding them is a no-op and cloning their weights for restore is pure memory waste (the disabled lm_head/embedding quantizers in the standard recipes hold ~40% of the quantized-weight bytes). Calibration state (_amax/_pre_quant_scale) is never touched. Discover quantizers the same way fold_weight does, by the *_weight_quantizer attribute-name suffix, rather than looking up a plain module.weight_quantizer. Fused and MoE modules expose names like w13_weight_quantizer and gate_up_proj_weight_quantizer; a narrower lookup would leave them folded and disabled for the rest of training. The option applies to any quantization format. Only the weight quantizer is disabled, so activation-quantized recipes (W4A4) keep their input/output quantizers running and produce unchanged logprobs. Default off, and scoped to the no-grad frozen-weight stage. Includes config examples, documentation, and unit coverage that exercises the real ModelOpt library: forward output is bit-identical folded vs unfolded, weights restore through their original storage, amax survives, restoration holds on exception, fused *_weight_quantizer modules are restored, tied-embedding modules with weight=None are skipped, and disabled quantizers are left untouched. Signed-off-by: babyplutokurt <attaboykurt.yang@gmail.com>
…ches The frozen-weight redundancy addressed for get_logprobs by quant_fold_frozen_weight_snap also exists inside training: within one global batch, the gradient-accumulation microbatches all run forwards against identical weights (the optimizer steps once, after all of them), yet every microbatch forward recomputes weight_quantizer(weight). Add an opt-in policy.quant_cache_train_weight_snap that fake-quantizes each weight once per global batch. Folding cannot be reused here because training needs the weight quantizer in the autograd graph: ModelOpt's backward is straight-through estimation that can carry an amax clip mask (pass_through_bwd=false), which a disabled quantizer would silently drop. Instead, each enabled weight quantizer's forward is patched for the duration of one megatron_forward_backward call to replay a precomputed quantized weight, with a backward replicating ModelOpt's exactly: pass-through by default, or where(|w| <= amax, grad, 0) when the config disables pass-through. Forward outputs and gradients are bit-identical to the uncached path (unit-verified for INT8 per-channel and NVFP4 dynamic block quantization, in both backward modes). The wrap happens per megatron_forward_backward call -- one per global batch, strictly between zero_grad and optimizer.step() -- so the cache is rebuilt from fresh weights after every optimizer step and can never go stale. Parameters and quantizer state are never mutated; the patch is an instance-level forward override removed on exit, including on exception. Quantizers whose forward chain the replica cannot reproduce exactly (smoothquant pre_quant_scale, rotation, static block quantization, bias quantization, calibration mode) and calls with any tensor other than the module's weight (e.g. refit exports quantizing .float() copies) fall back to the original quantizer: correct, just not accelerated. Default off. Costs one cached copy of each quantized weight shard, held for the duration of one global batch. Includes documentation and unit coverage against the real ModelOpt library. Also minimizes the w4a4-real recipe YAML (drops a quant_fold_frozen_weight_snap already set by its parent config), flagged by the configs-minimize-check hook. Signed-off-by: babyplutokurt <attaboykurt.yang@gmail.com>
github-actions
Bot
added
the
Documentation
svcnvidia-nemo-ci
added
the
waiting-on-maintainers
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What does this PR do ?
Adds opt-in
policy.quant_cache_train_weight_snap: during QAT training, eachweight is fake-quantized once per global batch instead of on every
gradient-accumulation microbatch forward, with bit-identical forward outputs
and gradients.
Stacked on top of #3441, which is currenrt under review
Issues
Closes #3555
Details
Weights are frozen within one global batch (the optimizer steps once, after
all microbatches), so every microbatch forward recomputes the identical
weight_quantizer(weight)tensor. The existing logprobs fold(
quant_fold_frozen_weight_snap) cannot be reused here: it disables thequantizer, and ModelOpt's backward is STE that can carry an amax clip mask
(
pass_through_bwd: false), which a disabled quantizer would silently dropfrom weight gradients.
Instead,
temporarily_cache_weight_quantizationkeeps the quantizer in theautograd graph: it precomputes
Q(W)via the quantizer's own forward, patchesthe quantizer to replay it, and replicates ModelOpt's backward exactly
(pass-through STE by default,
where(|w| <= amax, grad, 0)when the configdisables pass-through). The quant worker wraps each
megatron_forward_backwardcall, one per global batch, strictly between
zero_gradandoptimizer.step(), so the cache is rebuilt from fresh weights after everyoptimizer step and can never go stale. Parameters and quantizer state are
never mutated; the patch is removed on exit, including on exception.
Safety fallbacks: quantizers whose forward chain the replica cannot reproduce
exactly (smoothquant
pre_quant_scale, rotation, static block quantization,bias quantization, calibration mode) and any call with a tensor other than the
module's weight (e.g. refit exports quantizing
.float()copies) fall back tothe original quantizer.
Verification
tests/unit/models/policy/test_weight_folding.py, 22 passing):forward and gradients bit-identical to the real quantizer for INT8
per-channel and NVFP4 dynamic block quantization on GPU, in both backward
modes, including a forced-active clip mask; stale-cache rebuild across a
weight update; exception restore; foreign-tensor fallback; worker routing
(one cache window per forward-backward call, never spanning an optimizer
step).
logs 112 quantizers cached with 1792 hits and 0 fallbacks per global batch
(112 x 16 microbatches, full coverage through the TE path). Steps 1 and 2
match the cache-off arm exactly (loss, rewards, generation lengths); step 3
diverges only via vLLM generation nondeterminism, outside the patched
window.
policy_trainingimproved 1.47s -> 1.33s and 1.51s -> 1.44s,matching the expected removal of 15 of 16 weight-quant passes.
Cost: one cached copy of each quantized weight shard, held for the duration of
one global batch. Default off.
Cost: one cached copy of each quantized weight shard, held for the duration of
one global batch. Default off.
Usage
Before your PR is "Ready for review"
Pre checks:
Additional Information
(fix/qat-logprobs-frozen-weight-snap); both features share the discovery
helper in weight_folding.py but hold no shared state.
(drops a quant_fold_frozen_weight_snap already set by its parent config),
flagged by the configs-minimize-check pre-commit hook.