Implement QuantizationMixin

[kylesayrs]

Purpose

• Abstract functionality which allows modifiers to act as quantization configs into a mixin called QuantizationMixin
  • This gives Pipeline Extraction #1279 an interface to properly infer which pipeline to use based on the recipe (if a recipe contains modifiers requires calibration, then use the "basic" or "sequential" pipelines)
  • This enables future modifiers to act as quantization modifiers (in the same way that GPTQ does now)
• Related to problem cause when doing w8a8kvfp8 for qwen2.5-vl #1354 where previous logic would attempt to add a QuantizedKVCache for dynamic kv_quant

Changes

• Implement QuantizationMixin which implements five public methods

  • Lifecycle methods
    • initialize_quantization is used to apply a config and attach observers to a model
      • quantization is disabled so that modules aren't quantized before they're calibrated
    • start_calibration is used to initialize calibration hooks and status
      • quantization is enabled, since we currently quantize as we calibrate, although this decision is somewhat arbitrary
    • end_calibration is used to remove calibration hooks and apply the frozen status
      • quantization remains enabled, since we want future forward passes to simulate quantization
  • Recipe-related methods
    • has_config returns true if a config was specified, used for checking against duplicate configs in the recipe
    • resolve_quantization_config returns the quantization config specified by the modifier fields

• QuantizationModifier inherits from QuantizationMixin

• GPTQModifier inherits from QuantizationMixin

  • Unlike QMod, GPTQ disables quantization during calibration. As noted before, this is a somewhat arbitrary choice but one which matches the current implementation

• Calibration utils

  • Replace set_unset_kv_cache with initialize_quantized_kv_cache and freeze_module_quantization
    • Treat the QuantizedKVCache as analogous to another observer
  • Pull setting the calibration status out ofupdate_weight_zp_scale
    • This better matches the lifecycle detailed in QuantizationMixin description
  • Implement reset_quantization_status which is used to remove any existing quantization configs before the current config is applied by initialize_quantization

Remove Support

• Removing support for recipe with multiple quantization modifiers active at the same time (a check for this will be added by Pipeline Extraction #1279)
• Remove num_calibration_steps, quantize, disable_quantization_observer_epoch and min_tokens_per_module
  • num_calibration_steps is already controlled by

    llm-compressor/src/llmcompressor/datasets/utils.py

    Line 106 in 42b62f5

    num_calibration_samples=dataset_args.num_calibration_samples,

  • quantize was implemented as a workaround for GPTQ's modifier builder. Similar functionality may be require to support SpinQuant + GPTQ, but such functionality should exist at a higher level
  • disable_quantization_observer_epoch seems to implement functionality where a model's observers are removed but quantization remains active. This functionality is maintained by setting an "end" epoch for qmod
  • min_tokens_per_module requires that the modifier have references to the calibration dataset, which is disallowed by Pipeline Extraction #1279. This information is already printed in GPTQ's logs. If research still wants this tool specifically for QuantizationModifier, then it can be reimplemented to avoid using references to the calibration dataset

Testing

• Updated tests to reflect new mixin
• Ran a set of GPTQ and QuantizationModifier examples to completion
• CI tests pass

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

[brian-dellabetta]

This LGTM! slowly understanding more

[coolKeen]

Purpose

• Abstract functionality which allows modifiers to act as quantization configs into a mixin called QuantizationMixin

  • This gives Pipeline Extraction #1279 an interface to properly infer which pipeline to use based on the recipe (if a recipe contains modifiers requires calibration, then use the "basic" or "sequential" pipelines)
  • This enables future modifiers to act as quantization modifiers (in the same way that GPTQ does now)

• Related to problem cause when doing w8a8kvfp8 for qwen2.5-vl #1354 where previous logic would attempt to add a QuantizedKVCache for dynamic kv_quant

Changes

• Implement QuantizationMixin which implements four public methods

  • attach_scheme_and_observers is used to apply the quantization config to the model and attach observers
  • register_calibration_hooks add calibration hooks which calibrate the observers which calibrate the scales
  • resolve_quantization_config returns the quantization config specified by the modifier fields
  • has_config returns true if a config was specified, used for checking against duplicate configs

• QuantizationModifier inherits from QuantizationMixin

  • The scheme and observers are attached on initialization

  • The activations and weights are calibrated on start via hooks and update_weight_zp_scale

    • Add a tqdm for weight calibration, which is useful for larger models

  • The observers and hooks are removed on finalize

• GPTQModifier inherits from QuantizationMixin

  • Implements a similar lifecycle to QuantizationModifier
  • Remove attached QuantizationModifier logic

• Replace set_unset_kv_cache with initialize_quantized_kv_cache and freeze_module_quantization

  • Treat the QuantizedKVCache as analogous to another observer

• Pull setting the calibration status out ofupdate_weight_zp_scale

  • This better matches the lifecycle detailed in QuantizationMixin description

Remove Support

• Removing support for recipe with multiple quantization modifiers active at the same time (a check for this will be added by Pipeline Extraction #1279)

• Remove num_calibration_steps, quantize, disable_quantization_observer_epoch and min_tokens_per_module

  • num_calibration_steps is already controlled by

    llm-compressor/src/llmcompressor/datasets/utils.py

    Line 106 in 42b62f5

    num_calibration_samples=dataset_args.num_calibration_samples,

  • quantize was implemented as a workaround for GPTQ's modifier builder. Similar functionality may be require to support SpinQuant + GPTQ, but such functionality should exist at a higher level
  • disable_quantization_observer_epoch seems to implement functionality where a model's observers are removed but quantization remains active. This is currently implemented by setting an "end" epoch
  • min_tokens_per_module requires that the modifier have references to the calibration dataset, which is disallowed by Pipeline Extraction #1279. This information is already printed in GPTQ's logs. If research still wants this tool specifically for QuantizationModifier, then it can be reimplemented to avoid using references to the calibration dataset

Testing

• Updated tests to reflect new mixin
• Ran a set of GPTQ and QuantizationModifier examples to completion
• CI tests pass

Hi, I see that you mention SpinQuant + GPTQ? Does llm compressor support such recipe now?

[kylesayrs]

Hi @coolKeen! Transforms support is actively being worked on, you can see the WIP PRs here!

[rahul-tuli]

LGTM pending comment; Thank you, nice work!

[brian-dellabetta]

Just make sure your other PR goes in first, or drop the changes to kv cache tests from this one

[brian-dellabetta]

good stuff. always nice when adding an abstraction leads to more lines removed than added 🔥

[rahul-tuli]

Synced offline on scope of testing; Changes introduced by this diff are tested by #1279 by virtue of lm_eval tests! Changes look good.

[kylesayrs]

https://github.com/neuralmagic/llm-compressor-testing/actions/runs/14779544427

Addressed