Added support to export for BF16 weight and amax for vLLM fakequant QAT

[kinjalpatel27]

What does this PR do?

Type of change: New Feature

Overview:

Support for vLLM fakequantize QAT/QAD checkpoint evaluation. This MR adds function to export checkpoint as BF16 weights and amax using export_hf_checkpoint for HF and export_mcore_gpt_to_hf for MCore using export_bf16_weights_amax option. The exported weights and amax can be used with vllm_serve_fakequant.py script to run saved checkpoint.

Usage

Refer to README.md

Testing

• Tested HF approach by exporting bf16 model using QAT script and running vllm server, verified amax values match
• Tested MCore approach by quantizing and exporting bf16 model using quantize.sh and export.sh script and running vllm server, verified amax values match

Before your PR is "Ready for review"

• Make sure you read and follow Contributor guidelines and your commits are signed.
• Is this change backward compatible?: Yes
• Did you write any new necessary tests?: No
• Did you add or update any necessary documentation?: Yes
• Did you update Changelog?: Yes

Additional Information

MCore export script doesn't have the option to export enable currently

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

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 74.47%. Comparing base (1d0ee04) to head (13f6bcd).
⚠️ Report is 7 commits behind head on main.

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #579      +/-   ##
==========================================
+ Coverage   74.43%   74.47%   +0.04%     
==========================================
  Files         182      182              
  Lines       18234    18255      +21     
==========================================
+ Hits        13572    13596      +24     
+ Misses       4662     4659       -3     

☔ View full report in Codecov by Sentry.
📢 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]

Thanks Kinjal for documenting this. Create a jira ticket to address this - https://jirasw.nvidia.com/browse/OMNIML-3051

[kinjalpatel27]

Thank you for creating the ticket

[realAsma]

Is it possible to export the entire TensorQuantizer state - This way we can seamlessly support PTQ/QAT and fake quantization.

The current support does not work in cases such as Mixed precision quantization (some layers in FP4, some in FP8, some disabled etc.) - we would need manual work arounds for this case. This support also does not work for other quantization such as AWQ.

We are relying on the fact that we are quantizing the model with the same quantization formats as that of PTQ/QAT during vllm_serve.

[kinjalpatel27]

Is it possible to export the entire TensorQuantizer state - This way we can seamlessly support PTQ/QAT and fake quantization.

The current support does not work in cases such as Mixed precision quantization (some layers in FP4, some in FP8, some disabled etc.) - we would need manual work arounds for this case. This support also does not work for other quantization such as AWQ.

We are relying on the fact that we are quantizing the model with the same quantization formats as that of PTQ/QAT during vllm_serve.

I have created two tickets to explore mixed precision and other quantization algorithm support. Exporting Tensorquantizer state and loading may require additional effort since vLLM model also combines multiple layers etc.

[meenchen]

Can you also add unit tests for modelopt.torch.export.unified_export_hf.export_hf_checkpoint and modelopt.torch.export.unified_export_megatron.export_mcore_gpt_to_hf

[meenchen]

Does this work with MoE models that use this quant module:

TensorRT-Model-Optimizer/modelopt/torch/quantization/plugins/vllm.py

Line 134 in 01e24fd

class _QuantFusedMoEBase(QuantModule):

[kinjalpatel27]

Can you give an example for which model you are talking about?

[meenchen]

Sure, you can try Qwen/Qwen3-30B-A3B-Instruct-2507

[kinjalpatel27]

Can you also add unit tests for modelopt.torch.export.unified_export_hf.export_hf_checkpoint and modelopt.torch.export.unified_export_megatron.export_mcore_gpt_to_hf

@meenchen Added tests for both in separate file: tests/gpu/torch/export/test_vllm_fakequant_export.py

[meenchen]

LGTM. Could you also try if Qwen/Qwen3-30B-A3B-Instruct-2507 works with this change?

[kinjalpatel27]

LGTM. Could you also try if Qwen/Qwen3-30B-A3B-Instruct-2507 works with this change?

@meenchen Thank you. I checked Qwen/Qwen3-30B-A3B-Instruct-2507, It doesn't work with FuseMoE yet. I am looking into fixing it.