Add layerwise decompression and compression to sequential pipeline

src/llmcompressor/args/dataset_arguments.py

@@ -304,6 +304,22 @@ class DatasetArguments(CustomDatasetArguments):
             "calibration. Default False."
         },
     )
+    layerwise_decompression: bool = field(
+        default=False,
+        metadata={
+            "help": "When True, each layer is decompressed before calibrating in "
+            "the sequential pipeline. Use when loading a pre-compressed model "
+            "that needs recalibration. Default False."
+        },
+    )
+    layerwise_compression: bool = field(
+        default=False,
+        metadata={
+            "help": "When True, each layer is compressed after propagation in the "
+            "sequential pipeline. Reduces peak memory by keeping only the "
+            "active layer decompressed. Default False."
+        },
+    )
 
     def is_dataset_provided(self) -> bool:
         return self.dataset is not None or self.dataset_path is not None

src/llmcompressor/entrypoints/oneshot.py

@@ -244,6 +244,13 @@ def apply_recipe_modifiers(
             if self.dataset_args.moe_calibrate_all_experts:
                 stack.enter_context(moe_calibration_context())
 
+            session.state.layerwise_decompression = (
+                self.dataset_args.layerwise_decompression
+            )
+            session.state.layerwise_compression = (
+                self.dataset_args.layerwise_compression
+            )
+
             session.initialize(
                 model=self.model,
                 start=-1,

src/llmcompressor/modifiers/quantization/calibration.py

@@ -217,8 +217,10 @@ def freeze_module_quantization(module: Module):
         # no quantization scheme nothing to do
         return
 
-    if module.quantization_status == QuantizationStatus.FROZEN:
-        # nothing to do, already frozen
+    if module.quantization_status in (
+        QuantizationStatus.FROZEN,
+        QuantizationStatus.COMPRESSED,
+    ):
         return
 
     # remove observers

src/llmcompressor/modifiers/quantization/quantization/base.py

@@ -4,11 +4,13 @@
     greedy_bin_packing,
     is_distributed,
 )
+from compressed_tensors.quantization import enable_quantization
 from compressed_tensors.quantization.utils import is_module_quantized
 
 from llmcompressor.core import Event, State
 from llmcompressor.modifiers import Modifier
 from llmcompressor.modifiers.quantization.calibration import (
+    freeze_module_quantization,
     observe,
     update_qparams,
 )
@@ -65,20 +67,28 @@ def on_initialize(self, state: State, **kwargs) -> bool:
 
         Then, according to the module's quantization scheme, observers and calibration
         hooks are added. These hooks are disabled until the modifier starts.
+
+        Skipped when layerwise_decompression is enabled because modules are still
+        compressed; per-layer setup happens in start_layerwise_calibration instead.
         """
         if not QuantizationMixin.has_config(self):
             raise ValueError(
                 "QuantizationModifier requires that quantization fields be specified"
             )
-        QuantizationMixin.initialize_quantization(self, state.model)
+        if not getattr(state, "layerwise_decompression", False):
+            QuantizationMixin.initialize_quantization(self, state.model)
 
         return True
 
     def on_calibration_start(self, state: State, event: Event, **kwargs):
         """
         Begin calibrating activations.
+
+        Skipped when layerwise_decompression is enabled; per-layer calibration
+        setup is handled by start_layerwise_calibration in the pipeline.
         """
-        QuantizationMixin.start_calibration(self, state.model)
+        if not getattr(state, "layerwise_decompression", False):
+            QuantizationMixin.start_calibration(self, state.model)
 
     def on_sequential_epoch_end(
         self, state: State, event: Event, modules: list[torch.nn.Module], **kwargs
@@ -91,24 +101,35 @@ def on_sequential_epoch_end(
         if not is_distributed():
             observe(modules, "weight")
             update_qparams(modules, "weight")
-            return
-
-        ### Distributed
-        rank = dist.get_rank()
-        world_size = dist.get_world_size()
+        else:
+            ### Distributed
+            rank = dist.get_rank()
+            world_size = dist.get_world_size()
+
+            module_list, rank_to_modules, module_to_rank = greedy_bin_packing(
+                modules,
+                world_size,
+                item_weight_fn=lambda mod: mod.weight.numel(),
+            )
 
-        module_list, rank_to_modules, module_to_rank = greedy_bin_packing(
-            modules,
-            world_size,
-            item_weight_fn=lambda mod: mod.weight.numel(),
-        )
+            observe(rank_to_modules[rank], "weight")
+            update_qparams(rank_to_modules[rank], "weight")
+            broadcast_qparams_and_cleanup(
+                module_list, module_to_rank, _WEIGHT_Q_PARAMS
+            )
 
-        observe(rank_to_modules[rank], "weight")
-        update_qparams(rank_to_modules[rank], "weight")
-        broadcast_qparams_and_cleanup(module_list, module_to_rank, _WEIGHT_Q_PARAMS)
+        if getattr(state, "layerwise_decompression", False):
+            self.remove_hooks(self._calibration_hooks)
+            for module in modules:
+                freeze_module_quantization(module)
+                enable_quantization(module)
 
     def on_calibration_end(self, state: State, event: Event, **kwargs):
         """
-        Finish calibrating by removing observers and calibration hooks
+        Finish calibrating by removing observers and calibration hooks.
+
+        Skipped when layerwise_decompression is enabled; per-layer cleanup
+        is handled in on_sequential_epoch_end.
         """
-        QuantizationMixin.end_calibration(self, state.model)
+        if not getattr(state, "layerwise_decompression", False):
+            QuantizationMixin.end_calibration(self, state.model)

src/llmcompressor/modifiers/quantization/quantization/mixin.py

@@ -24,6 +24,7 @@
     preset_name_to_scheme,
 )
 from compressed_tensors.quantization.utils import KV_CACHE_TARGETS
+from compressed_tensors.quantization.utils.helpers import is_module_quantized
 from compressed_tensors.utils import match_named_modules
 from pydantic import Field, PrivateAttr, field_validator, model_validator
 from torch.utils.hooks import RemovableHandle
@@ -269,6 +270,28 @@ def end_calibration(self, model: torch.nn.Module):
 
         model.apply(enable_quantization)  # keep quantization enabled
 
+    def start_layerwise_calibration(
+        self, model: torch.nn.Module, modules: list[torch.nn.Module]
+    ):
+        """
+        Set up quantization for a subset of modules after layerwise
+        decompression. Applies quantization config scoped to the given
+        modules, then initializes observers, calibration hooks, and fuses
+        weight observers.
+
+        :param model: the full model (needed for config application)
+        :param modules: modules in the current subgraph to prepare
+        """
+        apply_quantization_config(
+            model, self.resolved_config, allowed_modules=modules
+        )
+        for module in modules:
+            if is_module_quantized(module):
+                self._initialize_observers(module)
+                self._calibration_hooks |= self._initialize_hooks(module)
+                apply_calibration_status(module)
+        fuse_weight_observers(model)
+
     def sync_obs_act_stats(self, modules: Iterator[torch.nn.Module]):
         """
         Synchronize the activation statistics for observers

src/llmcompressor/pipelines/sequential/pipeline.py

@@ -2,7 +2,10 @@
 from typing import TYPE_CHECKING, Iterator
 
 import torch
+from compressed_tensors.compressors import compress_module, decompress_module
 from compressed_tensors.offload import disable_offloading, set_onload_device
+from compressed_tensors.distributed import is_distributed, replace_module_parallel
+from compressed_tensors.quantization.utils import is_module_quantized
 from torch.utils.data.dataloader import DataLoader
 from tqdm import tqdm
 
@@ -141,6 +144,22 @@ def __call__(
                 # reduce memory movement by keeping modules onloaded
                 num_batches = len(dataloader)
                 with disable_offloading():
+                    modules = subgraph.submodules(model)
+
+                    # layerwise decompression: strip compression and
+                    # re-apply quantization config for this subgraph
+                    if dataset_args.layerwise_decompression:
+                        compressed = [
+                            m for m in modules if is_module_quantized(m)
+                        ]
+                        for module in compressed:
+                            decompress_module(module, leave_decompressed=False)
+                        for modifier in modifiers:
+                            if hasattr(modifier, "start_layerwise_calibration"):
+                                modifier.start_layerwise_calibration(
+                                    model, modules
+                                )
+
                     # do a preliminary pass to trigger modifier hooks
                     for batch_idx, inputs in _get_batches(
                         activations,
@@ -157,7 +176,7 @@ def __call__(
                                 activations.update(batch_idx, outputs)
                                 activations.delete(batch_idx, subgraph.consumed_names)
 
-                    LifecycleCallbacks.sequential_epoch_end(subgraph.submodules(model))
+                    LifecycleCallbacks.sequential_epoch_end(modules)
 
                     if dataset_args.propagate_error:
                         # this pass does not trigger modifier hooks
@@ -177,5 +196,17 @@ def __call__(
                                         batch_idx, subgraph.consumed_names
                                     )
 
+                    # layerwise compression: pack weights back after
+                    # calibration and error propagation
+                    if dataset_args.layerwise_compression:
+                        quantized = [
+                            m for m in modules if is_module_quantized(m)
+                        ]
+                        if not is_distributed():
+                           for module in tqdm(quantized, desc="Compressing modules"):
+                                compress_module(module)
+                        else:
+                            replace_module_parallel(quantized, compress_module, desc="Compressing modules")
+
             # redundant, finish any remaining compression
             LifecycleCallbacks.calibration_end()