MSE observer for NVFP4

[shubhra]

SUMMARY:
"MSE Observer implemented for NVFP4 for both weights and activations. Originally, I had added a new observer mseact for input activations global scale calculation but moved it under the same observer to keep it clean."

[github-actions]

👋 Hi! Thank you for contributing to llm-compressor. Please add the ready label when the PR is ready for review.

Note: This is required to complete the testing suite, please only add the label once the PR is code complete and local testing has been performed.

[gemini-code-assist]

Summary of Changes

Hello @shubhra, 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 significantly enhances the quantization capabilities by introducing a specialized Mean Squared Error (MSE) observer for NVFP4 activations. The new mseact observer is tailored to accurately determine quantization parameters, including a dynamic global scale, for input activations. Concurrently, the base MSE observer has been updated to support global scale calculation and includes detailed debug logging, providing better visibility into the quantization process. These changes collectively improve the precision and debuggability of FP4 quantization for both weights and activations.

Highlights

• New mseact Observer for Activations: Introduced MovingAverageMSEObserverActivations specifically designed for NVFP4 quantization of input activations, enabling dynamic global scale calculation.
• Global Scale Calculation for MSE Observer: The existing MSEObserver (mse.py) now includes a calculate_gparam method to generate global scales based on MSE-optimized min/max values.
• Enhanced Debugging: Added extensive debug print statements across key methods in mse.py to provide better visibility into the quantization parameter calculation process.

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 an MSE observer for NVFP4, including a new mseact observer for activations. The overall direction is good, but there are several areas for improvement. The most significant issues are the presence of numerous debug print statements across the modified and new files, which should be removed, and substantial code duplication in the new mseact.py file that could be resolved through inheritance. I've also noted some minor issues regarding import locations and type hints.

[brian-dellabetta]

Thanks for the contribution! I've not worked on Observer logic much, so will leave approval to others, but code looks clean.

I wonder if this might be better served as a completely separate class that subclasses MovingAverageMSEObserver, sounds like that is what you had in mind in your summary, an "mseact" observer

[brian-dellabetta]

if this is the only place we use base_name, it might just be better to expose is_activation: bool = False to the constructor instead of base_name

[brian-dellabetta]

will have to confirm with @shanjiaz if this needs to be added elsewhere, since there are a couple different places Observers are instantiated

[shanjiaz]

Seems like the initialization step in calibration is covered, looks good to me!

[kylesayrs]

I think this is going to be a simpler approach

#1903