[mxinO]

What does this PR do?

Type of change: New feature and bug fix.

Part 1 of 2: expose exact per-weight quantizer-state capture and metadata-only export specs. Capture is non-mutating and handles numbered grouped weights and supported transposed layouts. Callers can inspect deployment metadata without retaining tensor state.

This also fixes shared format detection, grouped quantizer restoration and HF export metadata, including unquantized-layer exclusions in mixed configurations. Existing offline packing remains on its current path.

Existing scale helpers are reused. The opaque snapshot also records logical weight shape, quantization axes, block sizes and packing layout, which a quantizer-buffer dictionary alone does not provide. It is detached in-memory state, not a new serialized configuration model.

Usage

For a converted and calibrated module:

from modelopt.torch.export import (
    capture_quantized_weight_export_state,
    get_quantized_weight_export_spec,
)

spec = get_quantized_weight_export_spec(module, "weight")
state = capture_quantized_weight_export_state(module, "weight", cpu=False)

Capture returns None for an unquantized weight. The second PR adds consumers of these snapshots. Recapture state for each export when quantizer values may change.

Testing

Frozen-source container run: 169 focused CPU/GPU tests passed, no skips. Covers capture/spec, known scale values and axis/block metadata, detached-clone independence, required-calibration failures, ordinary offline export, GPU packing and TE grouped restoration. All changed-file pre-commit checks pass. This validates the extracted API and existing export paths, not training convergence.

Before your PR is "Ready for review"

• Backward compatible: Yes; additive APIs and shared export fixes.
• New dependency or externally copied code: None.
• Necessary tests: Capture/spec, exact grouped-weight resolution, detector, grouped restore and ordinary HF export regressions.
• Changelog: Updated.
• Claude approval: Not requested; independent agent review found no issues.

Additional Information

Split of #2251; the original PR is unchanged. Merge this first, then #2778 ([2/2]), which targets this branch. Megatron-Bridge and NeMo-RL integration require both.

Summary by CodeRabbit

• New Features
  • Added export specifications for supported quantized weights, including deployment metadata and non-mutating capture of quantizer state.
  • Expanded Hugging Face quantization metadata support for MXFP4, MXFP8, FP8 activation schemes, and block/group sizes, including mixed-format models and excluded layers.
  • Added support for the W4A16 MXFP4 format.
• Bug Fixes
  • Preserved grouped weight calibration values during model restore.
  • Unsupported non-square FP8 block layouts are now rejected explicitly.

[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]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

• Configuration used: Repository: NVIDIA/Model-Optimizer/.coderabbit.yaml
• Review profile: CHILL
• Plan: Enterprise
• Run ID: 205b714a-d6ce-4156-8ce8-23527223e267

📥 Commits

Reviewing files that changed from the base of the PR and between 5690015 and 847c9d2.

📒 Files selected for processing: 2

• modelopt/torch/export/quantized_weight_export.py
• tests/unit/torch/export/test_quantized_weight_state.py

Included review availability: This review used your included allowance. Your plan provides up to 12 included reviews per hour; 11 remain after this review.

📝 Walkthrough

The changes add quantized-weight export state capture and metadata-only deployment specs. They extend quantization format detection and Hugging Face config conversion, add export handling and FP8 block-layout validation, and skip grouped weight quantizers during post-restore amax recalibration.

Changes: Quantized-weight export

Layer / File(s)	Summary
Format detection and Hugging Face config conversion modelopt/torch/export/quant_format.py, modelopt/torch/export/quant_utils.py, modelopt/torch/export/convert_hf_config.py, tests/unit/torch/export/test_get_quantization.py, tests/unit/torch/export/test_quantized_weight_state.py	Format detection and Hugging Face config conversion cover additional FP8, MXFP8, MXFP4, and mixed-format algorithms. The conversion preserves a supplied group size.
Offline export formats and validation modelopt/torch/export/quant_utils.py, tests/gpu/torch/export/test_export.py, CHANGELOG.rst	Export handling adds MXFP4 scaling and quantization, preserves exclusions for unquantized layers, and validates supported FP8 block layouts. GPU tests cover export metadata and rejection cases.
Resolve and capture quantized-weight state modelopt/torch/export/quantized_weight_export.py, tests/unit/torch/export/test_quantized_weight_state.py	The capture API resolves per-weight quantizers and records validated, detached quantization state for supported formats. Tests cover calibration state, grouped weights, and packing layouts.
Deployment spec API modelopt/torch/export/quantized_weight_export.py, modelopt/torch/export/__init__.py, tests/unit/torch/export/test_quantized_weight_state.py	The deployment spec API reports format, block size, and input-quantizer status without capturing tensor state.

Grouped quantizer restore

Layer / File(s)	Summary
Preserve grouped weight amax on restore modelopt/torch/quantization/plugins/custom.py, tests/gpu_megatron/torch/quantization/plugins/test_transformer_engine.py	Post-restore recalibration skips grouped weight quantizers. A regression test checks that their calibrated amax values remain unchanged.

Priority: ➖ Normal

Estimated code review effort: 4 (Complex) | ~45 minutes

Change: Feature

Sequence Diagram(s)

sequenceDiagram
  participant Caller
  participant capture_quantized_weight_export_state
  participant WeightQuantizer
  participant get_quantized_weight_export_spec
  Caller->>capture_quantized_weight_export_state: Request state for a module weight
  capture_quantized_weight_export_state->>WeightQuantizer: Resolve quantizer and capture scale state
  WeightQuantizer-->>capture_quantized_weight_export_state: Provide quantization metadata and scales
  capture_quantized_weight_export_state-->>Caller: Return detached export state
  Caller->>get_quantized_weight_export_spec: Request deployment metadata
  get_quantized_weight_export_spec->>WeightQuantizer: Resolve format, block size, and input status
  WeightQuantizer-->>get_quantized_weight_export_spec: Provide quantizer metadata
  get_quantized_weight_export_spec-->>Caller: Return deployment spec

Loading

Possibly related PRs

• NVIDIA/Model-Optimizer#2251: Both changes implement functional quantized-weight export in modelopt.torch.export.quantized_weight_export.

Suggested labels: export/deploy

Merge Risk: ⚪ Minimal · up to 847c9

No actionable issue was established in the reviewed changes. The PR is mergeable after normal checks.

Pre-merge checks: 5 · 1

Failed checks: 1 warning

Check name	Status	Explanation	Resolution
Docstring Coverage		Docstring coverage is 29.41% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 51 functions across 10 files.	Write docstrings for the functions missing them to satisfy the coverage threshold.

Passed checks: 5 passed

Check name	Status	Explanation
Linked Issues check		Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check		Check skipped because no linked issues were found for this pull request.
Security Anti-Patterns		PASS. The authoritative diff changes six modelopt Python files and no examples, pyproject.toml, requirements, or lock files. Added lines contain none of the prohibited patterns: unsafe torch.load or n…
Description Check		Check skipped - CodeRabbit’s high-level summary is enabled.
Title check		The title clearly identifies the main change: unified export support for quantized-weight state capture. It is concise and accurately reflects the pull request objectives.

• Fix all pre-merge checks with AI

✨ Finishing Touches: 💡 1

📝 Generate docstrings: 💡

• Commit to this branch
• Create a new PR

🧪 Generate unit tests (beta)

• Commit to this branch
• Create a new PR

• Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts

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

[github-actions]

PR Preview Action v1.8.1
Preview removed because the pull request was closed.
2026-10-16 05:08 UTC

[codecov]

Codecov Report

❌ Patch coverage is 97.05882% with 8 lines in your changes missing coverage. Please review.
✅ Project coverage is 78.36%. Comparing base (b5f10f8) to head (847c9d2).
⚠️ Report is 5 commits behind head on main.

Files with missing lines	Patch %	Lines
modelopt/torch/export/quantized_weight_export.py	96.77%	6 Missing ⚠️
modelopt/torch/export/quant_utils.py	96.96%	2 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #2777      +/-   ##
==========================================
+ Coverage   69.48%   78.36%   +8.88%     
==========================================
  Files         648      649       +1     
  Lines       71791    72005     +214     
==========================================
+ Hits        49881    56426    +6545     
+ Misses      21910    15579    -6331     

Flag	Coverage Δ
examples-diffusers	20.28% <31.98%> (-0.01%)	⬇️
examples-gpt-oss	13.56% <19.11%> (+0.01%)	⬆️
examples-hf_ptq	23.37% <38.97%> (-0.05%)	⬇️
examples-llm_distill	13.64% <19.11%> (+0.01%)	⬆️
examples-llm_eval	17.33% <25.36%> (+<0.01%)	⬆️
examples-llm_qat	17.63% <29.04%> (-0.01%)	⬇️
examples-llm_sparsity	15.88% <19.11%> (+<0.01%)	⬆️
examples-megatron_bridge	26.69% <30.14%> (-0.16%)	⬇️
examples-specdec_bench	13.34% <19.11%> (+0.01%)	⬆️
examples-speculative_decoding	17.62% <24.63%> (-0.07%)	⬇️
examples-torch_onnx	21.65% <19.11%> (-0.03%)	⬇️
examples-torch_trt	15.32% <19.11%> (+<0.01%)	⬆️
examples-vllm_serve	14.04% <19.11%> (+0.04%)	⬆️
gpu	59.19% <47.79%> (+37.47%)	⬆️
regression	15.18% <19.11%> (+0.01%)	⬆️
unit	60.04% <87.86%> (+0.10%)	⬆️

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.

[mxinO]

Follow-ups for this two-PR series, carried over from the scope comment on #2251:

• Offline checkpoint saving: integrate these APIs through Megatron-Bridge. The current offline save path is not migrated by this series.
• Real KV-cache export: add K/V scale capture/export and FP8/NVFP4 KV metadata, including mixed per-layer configuration and the FP8 scale-clamp option.
• MLA scale unification: reconcile calibrated scales for separate projections that inference engines fuse. This belongs in quantization, not per-weight packing.
• AWQ/SmoothQuant: add channel pre-quant scales and shared-input scale reconciliation.
• SVDQuant: coordinate export of low-rank residuals alongside packed weights.
• INT4/INT8/IQ: add packing/state/schema coverage and runtime qualification. Their exclusion is initial scope, not an architectural restriction; plain INT8 weight-only does not require smoothing.
• Weight-only MXFP4: existing offline packing/metadata are supported, but functional API support remains deferred. Export support alone does not imply engine support; vLLM 0.29.0's ModelOpt loader does not recognize MXFP4/W4A16_MXFP4.
• Validation: extend architecture/parity coverage and qualify native runtime support separately from export.

These exclusions apply to the new functional API, not ModelOpt's existing offline exporter. Follow-up integration is tracked in Megatron-Bridge #4804; Bridge #5812 and NeMo-RL #3945 consume the complete two-part API.

[cjluo-nv]

Bot review (gpt-6.1-sol) — DM the bot to share feedback.

Nudge: 657 added core-logic lines remain over budget, and the snapshot abstraction needs design rationale and broader capture tests.

Needs action:

• ✂️ Split export/ (649 lines, itself over budget) by file: [1/3] shared quant_utils.py/quant_format.py/convert_hf_config.py fixes; [2/3] quantization/plugins/custom.py restoration; [3/3] quantized_weight_export.py capture APIs and export/__init__.py.
• Keep the series acyclic and linearly ordered; document merge order, link siblings, and ensure each slice builds with its own tests and passing CI without later slices.
• Explain in the PR body why extending quant_utils.py scale helpers or adapting get_quantizer_state_dict() is insufficient, and why dedicated dataclass snapshots are preferred over existing Pydantic state representations.
• Add capture tests in tests/unit/torch/export/test_quantized_weight_state.py for per-channel/block FP8, dynamic NVFP4, and MXFP4/MXFP8; verify scale values, axis/block metadata, detached-copy independence, and missing-calibration failures.

No action needed:

• The existing mixed-format assertion change is justified by the new exclusion behavior; existing coverage is retained.
• New-file headers match LICENSE_HEADER.
• Tests were inspected, not executed.

[mxinO]

Addressed the capture coverage and rationale feedback from this review. The new cases cover scale values, axis/block metadata, detached-copy independence and required-calibration failures; the independent capture snapshot passes 169 CPU/GPU tests with no skips. The PR body explains the extra layout metadata and use of dataclasses for internal state rather than persisted configuration.

Keeping the agreed two-part capture/consumer split for this series instead of another three-way split. Both parts are independently buildable and tested.

[cjluo-nv]

Bot review (gpt-6.1-sol) — DM the bot to share feedback.

Nudge: capture coverage and design rationale are resolved, but 659 added core-logic lines still exceed the review budget.

Needs action:

• ✂️ Split modelopt/torch/export (651 lines, itself over budget) by file: [1/3] shared quant_utils.py/quant_format.py/convert_hf_config.py fixes; [2/3] modelopt/torch/quantization/plugins/custom.py restoration; [3/3] quantized_weight_export.py capture APIs and export/__init__.py.
• 💬 Author retains the two-part capture/consumer series — further subdivide capture; keep dependencies acyclic and linearly ordered, document merge order, link siblings, and give each slice its own tests and passing CI without later slices.

No action needed:

• ✔️ Resolved since the last review: requested format/scale/axis/block, detached-copy and calibration-failure coverage; snapshot design rationale.
• The problem is independent per-weight export capture with layout metadata. Reusing scale helpers, adapting get_quantizer_state_dict(), and existing Pydantic representations were considered; the body reasonably explains detached dataclass state versus persisted configuration.
• The existing mixed-format assertion edit remains justified by retained exclusions; existing coverage is preserved. New headers match LICENSE_HEADER.
• Tests were inspected, not executed.

[sugunav14]

My understanding is that this file is capturing the quantization state. Could we rename the file to better reflect that? Maybe something like quantization_export_state.py or quantization_weight_state.py

[sugunav14]

Left a minor readability comment. Otherwise LGTM!