Add awq activation fp8 support in loss compute

[Bluedyson]

SUMMARY:
Introduced FP8 quantization error into loss calculation with new config options, validation checks, and tests.
See #1657 (comment)

TEST PLAN:

"Qwen/Qwen3-32B", On our internal dataset, the KL divergence

W4A16: 81
W4AFP8 : 89
TensorRT-Model-Optimizer NVFP4-FP8: 88

I'll quickly add general dataset tests result to this PR.

Keep tracing:

• W4fp8 AWQ #1657

[gemini-code-assist]

Summary of Changes

Hello @Bluedyson, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request enhances the AWQ modifier by integrating FP8 activation quantization support, which allows for the simulation of FP8 quantization errors directly into the loss calculation. This change involves adding new configuration options, implementing robust validation to enforce float-type activation quantization, and updating the core scaling computation to reflect FP8 precision. The goal is to improve the accuracy of loss calculations in quantized models and expand the capabilities of the AWQ modifier.

Highlights

• FP8 Activation Support: Introduced support for FP8 activation quantization within the AWQ modifier, allowing for more granular control over quantization error in loss calculations.
• Configuration and Validation: Added a new private attribute _activation_bits to the AWQModifier and updated validation logic to ensure that lower-precision activation quantization in AWQ is of type FLOAT, raising an error otherwise.
• Loss Calculation Integration: Modified the _compute_best_scale method to simulate FP8 activation quantization by converting scaled weights to torch.float8_e4m3fn and then back to torch.float16 when FP8 activations are configured.
• Testing: New test cases were added to validate the correct behavior of FP8 float activation quantization and to ensure that invalid integer-type activation quantization configurations raise appropriate errors.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[gemini-code-assist]

Code Review

This pull request introduces support for FP8 activation quantization within the AWQ modifier. The changes primarily involve updating the validation logic to accommodate activation quantization settings and modifying the loss computation in _compute_best_scale to simulate FP8 precision loss. Additionally, new tests are included to verify the new configuration options. My review focuses on enhancing the robustness of the validation logic to prevent non-deterministic behavior when multiple, differing activation quantization configurations are provided.

[brian-dellabetta]

@Bluedyson thanks for this! I will take a look after our release this week.