[Observer Restructure]: Add Observers; Add calibration and frozen steps to QuantizationModifier

[dsikka]

SUMMARY:

• PR to add observers to llm-compressor
• Adds the required hooks needed to run calibration as part of the QuantizationModifier. All required calibration lifecycle steps can now be found in calibration.py
• Also adds the KV Cache object such that calibration can be done to update k_scale and v_scale for kv_cache quantization
• Requires the following PR to land in compressed-tensors: Observer Restructure: Remove Observers, calibration, and applying frozen steps from lifecycle neuralmagic/compressed-tensors#189
• Updated Calibration lifecycle (also shown in the docstrings). This will run as part of the calibration step within the QuantizationModifier

Run calibration if running input/output activation quantization or kv_cache quantization.

Calibration Lifecycle for a single torch.nn.Module:

      1. initialize_observer():
          if input/output activation:
              - observer = Observer.load_from_registry(...)
              - module.register_module(f"{base_name}_observer", observer)
              
      2. register_calibration_hooks():
          if input activation and not dynamic quant (used to call observers before intput QDQ):
              - pre_hook_handle = module.register_forward_pre_hook(calibrate_input_hook())
          if output activation and not dynamic quant (used to call observers before output QDQ):
              - post_hook_handle = module.register_forward_hook(calibrate_kv_cache_output_hook())
          if kv_cache quantization (used to set kv_cache to QuantizedKVParameterCache and update k_scale/v_scale)
              - pre_hook_handle = module.register_forward_pre_hook(calibrate_kv_cache_input_hook(), with_kwargs=True)
              - post_hook_handle = module.register_forward_hook(calibrate_kv_cache_output_hook())
          self.calibration_hooks.append(pre_hook_handle)
          self.calibration_hooks.append(post_hook_handle)

      3. self._calibrate(module) # run forward pass through model using calibration data
      4. set_unset_kv_cache() # remove kv_cache objects attached to attention layers  initially set in _apply_modifier_to_model
      5. remove calibration hooks in self.calibration_hooks_
      6. remove observers

Testing:

• Tested w4a16, quantized kv_cache, and w8a8 int8 workflows

[github-actions]

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

[dsikka]

With the corresponding remove-observers branch checked out

python3 examples/quantization_w4a16/llama3_example.py                                                                                                                         
Traceback (most recent call last):                                                                                         
  File "/home/ksayers/llm-compressor/examples/quantization_w4a16/llama3_example.py", line 4, in <module>                   
    from llmcompressor.modifiers.quantization import GPTQModifier                                                          
  File "/home/ksayers/llm-compressor/src/llmcompressor/modifiers/quantization/__init__.py", line 3, in <module>            
    from .cache import *                                                                                                   
  File "/home/ksayers/llm-compressor/src/llmcompressor/modifiers/quantization/cache.py", line 18, in <module>              
    from compressed_tensors.quantization.lifecycle import KVCacheScaleType 

Can you confirm you're using the most recent commit for both? I do not get this error. And in general you should not as the scales are defined under src/compressed_tensors/quantization/lifecycle/initialize.py

[rahul-tuli]

I really like the new structure, Great work!

Left a few nits, would recommend revisiting the docstrings and updating them for consistency:
-> Start docstrings with a Capital Letter
-> Include param info in :params over just writing a description in the main docstring

Otherwise no big red flags! Good tests as well.