Add layerwise decompression and compression to sequential pipeline - #2995
Add layerwise decompression and compression to sequential pipeline#2995kylesayrs wants to merge 2 commits into
Conversation
Enable per-layer decompress→calibrate→compress cycles in the sequential pipeline, controlled by `layerwise_decompression` and `layerwise_compression` DatasetArguments. This allows recalibrating pre-compressed models without fully decompressing the entire model at once. When layerwise_decompression is enabled, each subgraph's compressed modules are decompressed with leave_decompressed=False (stripping all quantization metadata), then quantization config, observers, and hooks are re-applied scoped to just those modules before the calibration batch loop runs. When layerwise_compression is enabled, quantized modules are compressed back after calibration and error propagation, keeping peak memory low. Companion PR: vllm-project/compressed-tensors#811 Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
|
👋 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. |
|
Important Review skippedAuto reviews are disabled on this repository. Please check the settings in the CodeRabbit UI or the ⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Pro Plus Run ID: You can disable this status message by setting the Use the checkbox below for a quick retry:
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
|
The quality checks have failed. Please run |
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces layerwise decompression and compression options to the sequential calibration pipeline, allowing pre-compressed models to be decompressed layer-by-layer for recalibration and compressed back afterward to reduce peak memory usage. The changes span dataset arguments, the oneshot entrypoint, quantization modifiers, and the sequential pipeline itself. The review feedback highlights two important improvements: first, adding a validation check to ensure these layerwise options are only used with the 'sequential' pipeline to prevent silent failures; second, clearing the _calibration_hooks set after removing hooks during sequential epoch ends to avoid O(N^2) accumulation of hook handles and potential memory leaks.
| session.state.layerwise_decompression = ( | ||
| self.dataset_args.layerwise_decompression | ||
| ) | ||
| session.state.layerwise_compression = ( | ||
| self.dataset_args.layerwise_compression | ||
| ) |
There was a problem hiding this comment.
When layerwise_decompression or layerwise_compression is enabled, the sequential pipeline is strictly required. If another pipeline (such as basic or independent) is used, the global initialization and calibration steps in QuantizationModifier will be silently skipped, resulting in a model that is never actually quantized or calibrated. Adding a validation check here prevents this silent failure.
if (
self.dataset_args.layerwise_decompression
or self.dataset_args.layerwise_compression
) and self.dataset_args.pipeline != "sequential":
raise ValueError(
"Layerwise decompression and compression are only supported "
"with the 'sequential' pipeline."
)
session.state.layerwise_decompression = (
self.dataset_args.layerwise_decompression
)
session.state.layerwise_compression = (
self.dataset_args.layerwise_compression
)| if getattr(state, "layerwise_decompression", False): | ||
| self.remove_hooks(self._calibration_hooks) | ||
| for module in modules: | ||
| freeze_module_quantization(module) | ||
| enable_quantization(module) |
There was a problem hiding this comment.
The self._calibration_hooks set is not cleared after removing the hooks. Since on_sequential_epoch_end is called per-layer in the sequential pipeline, this leads to an
| if getattr(state, "layerwise_decompression", False): | |
| self.remove_hooks(self._calibration_hooks) | |
| for module in modules: | |
| freeze_module_quantization(module) | |
| enable_quantization(module) | |
| if getattr(state, "layerwise_decompression", False): | |
| self.remove_hooks(self._calibration_hooks) | |
| self._calibration_hooks.clear() | |
| for module in modules: | |
| freeze_module_quantization(module) | |
| enable_quantization(module) |
Merge Protections🔴 1 of 1 protections blocking · waiting on 👀 reviews
🔴 Require one maintainer reviewWaiting for any of
This rule is failing.All PRs must have at least one approving review from a maintainer before merging.
|
Signed-off-by: Kyle Sayers <kylesayrs@gmail.com>
|
The quality checks have failed. Please run |
1 participant
Summary
layerwise_decompressionandlayerwise_compressionflags toDatasetArgumentsfor the sequential pipelinelayerwise_decompression=True, each subgraph's compressed modules are decompressed (stripping all quantization metadata), then quantization config, observers, and hooks are re-applied per-layer before calibrationlayerwise_compression=True, quantized modules are compressed back after calibration and error propagation to reduce peak memoryQuantizationModifierconditionally skips global setup (on_initialize,on_calibration_start,on_calibration_end) when layerwise mode is active, deferring to per-layer setup viastart_layerwise_calibrationfreeze_module_quantizationto preserveCOMPRESSEDstatus (not overwrite withFROZEN)Companion PR
leave_decompressedparameter +allowed_modulessupport inapply_quantization_config)Test plan
freeze_module_quantizationpreserves COMPRESSED status🤖 Generated with Claude Code