Add LSQ (Learned Scale Quantization) support and recipes

[realAsma]

Summary

This PR adds two scale-learning algorithms to Model Optimizer—LSQ and Dual-LSQ—and provides focused NVFP4 Quantization-Aware Distillation (QAD) recipes for them.

• LSQ (learnt scale quantization) corresponds to the original Learned Step Size Quantization paper. ModelOpt uses the name learnt scale quantization and learns one scale shared before and after quantization.
• Dual-LSQ learns separate pre-quantization and post-quantization scales.

ModelOpt implements both algorithms with an amax reparameterization: it learns amax rather than the scale directly, with scale = amax / max_bound. This allows scale parameters to be optimized during QAD while the quantized model weights remain fixed.

What changed

• Add an LSQConfig algorithm and calibration mode with configurable learned pre/post amax values, tied-scale support, and optional pre-scale quantization.
• Extend TensorQuantizer, tensor quantization, and the FP4 GEMM path so gradients propagate through LSQ scale parameters.
• Preserve LSQ parameter dtypes when training with FSDP2.
• Add two modular NVFP4 recipes:
  • LSQ: one shared learned pre/post amax value.
  • Dual-LSQ: independent learned pre/post amax values; the pre-scale is not quantized.
• Add a scale-only QAD training config that trains only LSQ amax parameters.
• Add focused CPU/GPU LSQ behavior, recipe, and FSDP2 coverage.

Validation

• Focused LSQ and recipe unit tests: 38 passed
• GPU tests: test_lsq_cuda.py 6 passed; FSDP2 LSQ test passed.
• Pre-commit checks passed for all changed files.
• git diff --check passed.

[copy-pr-bot]

Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

[coderabbitai]

Important

Review skipped

Draft detected.

Please check the settings in the CodeRabbit UI or the .coderabbit.yaml file in this repository. To trigger a single review, invoke the @coderabbitai review command.

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 4a907e1a-5570-40e3-b896-7e6d25736b71

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

Use the checkbox below for a quick retry:

• 🔍 Trigger review

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Commit unit tests in branch asma/laq-qad-recipes

---

_{Comment @coderabbitai help to get the list of available commands.}

[codecov]

Codecov Report

❌ Patch coverage is 56.56934% with 119 lines in your changes missing coverage. Please review.
✅ Project coverage is 77.39%. Comparing base (43fee0c) to head (6c85097).
⚠️ Report is 1 commits behind head on main.

Files with missing lines	Patch %	Lines
modelopt/torch/quantization/model_calib.py	14.10%	67 Missing ⚠️
...lopt/torch/kernels/quantization/gemm/fp4_kernel.py	0.00%	23 Missing ⚠️
.../torch/quantization/nn/modules/tensor_quantizer.py	82.85%	18 Missing ⚠️
modelopt/torch/quantization/tensor_quant.py	74.19%	8 Missing ⚠️
...torch/quantization/plugins/transformers_trainer.py	80.00%	3 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #1884      +/-   ##
==========================================
- Coverage   77.64%   77.39%   -0.25%     
==========================================
  Files         516      516              
  Lines       57465    58761    +1296     
==========================================
+ Hits        44618    45481     +863     
- Misses      12847    13280     +433     

Flag	Coverage Δ
unit	55.12% <56.56%> (+0.09%)	⬆️

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[realAsma]

BB: Can you write a good PR description for this PR? Summarize the changes in this PR - address it to external reviewers, not me.

[realAsma]

BB: rename this to qad_scale_only.yaml

[realAsma]

🤖 Bot comment.

🐝 I’ll prepare this rename in the PR worktree and report back here before any commit or push.

[realAsma]

🤖 Bot comment.

Prepared the rename locally:

examples/llm_qat/configs/train/qad_laq_scales.yaml → examples/llm_qat/configs/train/qad_scale_only.yaml

It is a content-preserving move. The renamed YAML parses successfully, git diff --check passes, and there are no references that need updating. I have not committed or pushed it; reply commit and push if you want me to publish the rename to this PR.

[realAsma]

okay commit and push.

[realAsma]

BB: Can you add one more train yaml file for full parameter QAD. Lets use LR 1e-5 for all parameters except amax. for learnable amax -> use 1e-4 LR.

[realAsma]

🤖 Bot comment.

🐝 I’m drafting a reviewer-facing summary of the PR now.

[realAsma]

🤖 Bot comment.

I prepared the following reviewer-facing PR description. I have not changed the PR body because this request did not authorize publishing; reply apply if you want me to update it.

Summary

This PR adds Learnable Amax Quantization (LAQ) support to Model Optimizer and provides focused NVFP4 Quantization-Aware Distillation (QAD) recipes for learning quantization scales.

LAQ represents quantization with separate pre-quantization and post-dequantization amax values. Those values can be learned independently (Dual Scale) or tied to one shared parameter (LSQ), allowing scale parameters to be optimized during QAD while the quantized model weights remain fixed.

What changed

• Add an LAQConfig algorithm and calibration mode with configurable learned pre/post amax values, tied-scale support, and optional pre-scale quantization.
• Extend TensorQuantizer, tensor quantization, and the FP4 GEMM path so gradients propagate through LAQ scale parameters.
• Preserve LAQ parameter dtypes when training with FSDP2.
• Add two modular NVFP4 recipes:
  • Dual Scale: independent learned pre/post amax values; the pre-scale is not quantized.
  • LSQ: one shared learned pre/post amax value.
• Add a scale-only QAD training config that trains only LAQ amax parameters.
• Add focused CPU/GPU LAQ behavior, recipe, and FSDP2 coverage.

Validation

• Focused LAQ and recipe unit tests: 38 passed
• Pre-commit checks passed for all changed files.
• git diff --check passed.

[realAsma]

🤖 Bot comment.

🐝 Got it — I’ll add a full-parameter QAD train config with LR 1e-5 for regular parameters and 1e-4 for learnable amax parameters. I’ll prepare and validate it without pushing until you approve.

[realAsma]

🤖 Bot comment.

Prepared the full-parameter QAD config locally.

• configs/train/qad_full.yaml keeps all parameters trainable and uses the global 1e-5 learning rate.
• configs/train/qad_full_lr.yaml applies 1e-4 specifically to *weight_quantizer._amax_pre and *weight_quantizer._amax_post through the existing per-parameter LR-config mechanism.

The training config and LR mapping parse successfully, their exact LR values/patterns were checked, the YAML formatting hooks pass, and git diff --check is clean. I have not committed or pushed these files; reply commit and push if you want me to publish them together with the prepared qad_scale_only.yaml rename.

[realAsma]

BB: Please update the PR description to be external facing.

[realAsma]

🤖 Bot comment.

🐝 I’ll update the PR description to be external-facing now.

[realAsma]

🤖 Bot comment.

Updated the PR description to be external-facing, with a reviewer-oriented summary, the LAQ/LSQ and QAD changes, and validation results. I also verified the published body on the live PR. No source code or branch changes were made.

Requested in: #1884 (comment)

[realAsma]

BB: Please do this #1884 (comment)

[realAsma]

🤖 Bot comment.

🐝 Got it — I’ll publish the prepared QAD config updates and report back here with the commit and validation results.

[realAsma]

🤖 Bot comment.

Published in 39cc0fe060.

• Renamed qad_laq_scales.yaml to qad_scale_only.yaml without content changes.
• Added qad_full.yaml with all parameters trainable at the global 1e-5 rate.
• Added qad_full_lr.yaml to apply 1e-4 specifically to LAQ _amax_pre / _amax_post parameters.

Validated config parsing and exact LR mappings; YAML hooks and git diff --check pass. The commit is signed with DCO, and the live draft PR now points to this head.

[realAsma]

BB: This PR add two algorithms actually -> LSQ and Dual-LSQ -> can we update PR description and and the modelopt/torch/quantization/config.py with this information.

LSQ: Refer to the original https://arxiv.org/pdf/1902.08153 paper (we call it learnt scale quantization instead of learnt step size quantization - both are similar);

Dual-LSQ: We learn separate pre quantization scale and post quantization scale.

Then add a note that we actually do a reparameterized version of LSQ, learn amax instead of scale directly. scale is computed as scale = amax/max_bound.

[realAsma]

🤖 Bot comment.

🐝 Got it — I’ll update the PR description and quantization config to clearly distinguish LSQ from Dual-LSQ and document the learned-amax reparameterization.

[realAsma]

🤖 Bot comment.

Done 🐝

Updated the PR description and modelopt/torch/quantization/config.py to:

• distinguish LSQ (learnt scale quantization, tied pre/post scale) from Dual-LSQ (separate pre/post scales),
• reference the original Learned Step Size Quantization paper, and
• document ModelOpt’s reparameterization: learn amax and compute scale = amax / max_bound.

Published as signed/DCO commit fb9dee9a5c. The affected-file pre-commit hooks and git diff --check passed; this is documentation-only with no behavior change.

[realAsma]

BB: please investigate and fix the CICD errors - you can restart pipelines as needed. If there are any major design changes notify me via slack for approval. Minor changes are approved - keep a summary and commits diff I can review later.

[realAsma]

🤖 Bot comment.

🐝 I’m on it. I’ll inspect the failing CI/CD checks, fix and publish minor issues as needed, restart pipelines where useful, and report back with the validation and commit diff. I’ll pause and message you privately before any major design change.

[realAsma]

BB: Can you launch these NVFP4 PTQ+QAD jobs from this feature branch folder:

1. full parameters + scale learning (Dual LSQ), mse init
2. Scale only learning, mse init -> both Dual LSQ and LSQ
3. Dynamic max full parameter.

Model:
Qwen3 1.7B

Dataset:
The default blend (I think 20k size)

Use local GPUs (may be 4 GPU job)

Create BeeBot timers to monitor each job, once everything is done, update the PR description with train and eval loss (2 column 1 row). Solid lines - each variant in a distinct color.

Notify me in slack once done.

[realAsma]

🤖 Bot comment.

🐝 I’ve verified this request and queued the four Qwen3 1.7B NVFP4 PTQ+QAD variants with the active PR work. I’ll monitor them, add the train/eval loss plots to the PR description, and notify you in Slack when everything is complete.

[realAsma]

🤖 Bot comment.

🐝 I’ve got the four-run experiment request too. I’ll launch the specified Qwen3 1.7B NVFP4 PTQ+QAD variants on local GPUs with captured launch evidence, monitor them through completion, add the train/eval loss curves to the PR description, and notify you on Slack.

[github-actions]

PR Preview Action v1.8.1
🚀 View preview at https://NVIDIA.github.io/Model-Optimizer/pr-preview/pr-1884/
Built to branch gh-pages at 2026-07-07 04:04 UTC. Preview will be ready when the GitHub Pages deployment is complete.

[realAsma]

BB:

Make sure you launch one job at a time on the local a8 server. Otherwise they might go OOM.

[realAsma]

🤖 Bot comment.

🐝 Got it. I’ll keep the local a8 workload strictly serialized—one job at a time—and verify the active launch plan follows that limit to avoid OOM.

[realAsma]

🤖 Bot comment.

Outcome: the local a8 OOM-safety constraint is recorded on the active experiment task: no more than one job may run at a time. I did not start any duplicate work from this handler; the existing launcher owner is responsible for serializing the remaining PTQ/QAD stages and verifying the limit before each launch.

[realAsma]

🤖 Bot comment.

CI/CD fix published in signed/DCO commit 6c85097f38 (Fix LSQ CI checks).

• Fixed docs: rendered the FP4 clip expression as an inline literal, removing Sphinx’s undefined |x| substitution.
• Fixed dependency-minimal tests: scoped the optional Transformers plugin import to its one LSQ test, so torch-only collection skips it cleanly when Transformers is absent.
• Diff: 2 files, 6 insertions, 2 deletions.
• Local validation: 33 LSQ tests passed; docs passed with warnings-as-errors; torch partial-install suite passed (1,447 passed, 76 skipped); pre-commit and diff checks passed.
• Fresh GitHub results: both previously failing gates (build-docs and partial-install (torch)) are green, along with Linux, Windows, deploy preview, code quality, DCO, and the other completed matrix jobs. One final multi-version job is still running with no failures reported.

[realAsma]

🤖 Bot comment.

CI follow-up: the final multi-version job and aggregate required check are now green. The fresh Unit tests workflow finished with all 19 jobs successful at 6c85097f38; no pipeline restart is needed.