[Observers] Refactor for better FP4 support, static and memoryless observers

[kylesayrs]

Purpose

• FP4
  • Fix bug discovered here where dynamic="local" nvfp4 calculations would increment the observer twice as fast as normal
  • Enable MSE observer to be used with FP4

    mse_quant_error := mean((x - fake_quant(x))**2)
    global_scale <- min[min_vals, max_vals, global_scale](mse_quant_error(x))
    scale, zp <- min[min_vals, max_vals](mse_quant_error(x, global_scale))
    

• Simplification
  • Make supporting attention calibration easier by separating out weight/activation/attention reshaping
  • Improve readability of observer codes by removing many levels of function indirection
  • Drop support for calibration with non-divisible group sizes. This is not really a loss, since forward passes also make this assumption
• New observers
  • memoryless_minmax computes min and max values on the fly in a dynamic-quantization style. This observer is useful for PTQ weight quantization
  • static_minmax computes absolute min and max values across all observations. This observer is useful for PTQ activation quantization
  • memoryless_mse computes best qparams w.r.t. MSE loss for each observation. This observer is useful for PTQ weight quantization
• Memory improvements
  • All observers no longer store copies of scales and zero points, reducing the amount of required memory
  • Newly introduced "memoryless" observers do not store any quantization parameters, which greatly reduces the memory requirements for PTQ weight quantization of very large models

Diagrams
Before
After

Changes

• Standardize reshaping using flatten_for_calibration
  • This function reshapes all observed values to (num_observations, *qparams_shape, group_size)
  • This function the complexity associated with passing "reduce dims" and trying to handle weights, activations, and attention states all in the same function
  • In the future, this function could be applied to the quantization forward pass, although there's probably no need to outside of standardization
• Implement get_global_scale on Observer base
  • This function decouples minmax calculations from regular qparam calculations (avoiding the double increment bug)
  • This function enables the MSE observer to be used with FP4 global scales

Testing

• Added additional minmax tests which check exact values of scales. This test passes both on main and this branch, demonstrating that minmax observer behavior remains unchanged
• Added additional MSE tests which check exact values of mse losses. This test passes both on main and this branch, demonstrating that MSE observer behavior remains unchanged
• Added FP4 MSE test

Evaluation

nvfp4-static-minmax
| Tasks  |Version|Filter|n-shot| Metric |   |Value |   |Stderr|
|--------|------:|------|-----:|--------|---|-----:|---|------|
|mmmu_val|      0|none  |     0|mmmu_acc|↑  |0.6167|±  |   N/A|

nvfp4-minmax
| Tasks  |Version|Filter|n-shot| Metric |   |Value |   |Stderr|
|--------|------:|------|-----:|--------|---|-----:|---|------|
|mmmu_val|      0|none  |     0|mmmu_acc|↑  |0.6011|±  |   N/A|

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

[shanjiaz]

Looks good to me, thanks for fixing the weird observer_kwargs setup.

[brian-dellabetta]

The cleanup looks good, nice that we can prune so much code in src. Some clarifying questions/nits. Would be good to have @dsikka or @rahul-tuli check this out too, since they know the observer logic better than me.

Does this remove the need for #1840 ?

[kylesayrs]

@brian-dellabetta Yes, this adds support for FP4 + MSE in a more direct way, and has a test to validate this.

[brian-dellabetta]

One suggestion on adding tests for the new static observer classes

[dsikka]

Just an FYI - the Qwen3 VL MoE NVFP4 example has OOO errors on this branch, for both the static_minmax and minmax observer.

[kylesayrs]

@dsikka The OOO errors were due to gradients being calculated. Now, gradient calculations are disabled for observers.

The bad calibrations you saw with the MSE observer were due to the MSE observer using the uninitialized global_scale when computing MSE loss. I was able to replicate and fix the issue by recomputing the global scale for each iteration (note that this is only done when computing global_scales, not scales/zps).

[HDCharles]

I think the change to not have a static version of MSE makes sense given the lack of a clear algorithm. LGTM

[brian-dellabetta]

very nice! looks so much cleaner