implement memory decompression

Signed-off-by: Kyle Sayers <kylesayrs@gmail.com>

Author: kylesayrs

src/compressed_tensors/compressors/model_compressors/model_compressor.py

@@ -268,6 +268,8 @@ def __init__(
                 quantization_config.format, config=quantization_config
             )
 
+    # ----- used by hf quantizer ----- #
+
     def get_missing_module_keys(self, model: Module) -> List[str]:
         """
         Identifies the expected missing weight keys in the compressed state_dict.
@@ -369,7 +371,7 @@ def get_unexpected_file_keys(self, model: Module) -> List[str]:
 
         return list(unexpected_keys)
 
-    # ----- model compression/decompression pathways ----- #
+    # ----- model memory compression/decompression pathways ----- #
 
     def compress_model(self, model: Module):
         module_to_scheme = map_module_to_scheme(model)
@@ -418,13 +420,6 @@ def decompress_model(self, model: Module):
             ignore=self.sparsity_config.ignore if self.sparsity_config else [],
         )
 
-        # because decompressors are implemented to only generate new values (rather than
-        # generating new values and unused values), we must explicitly pass a list of
-        # keys to yield which are unused (but used in subsequent decompressors)
-        params_to_ignore = None
-        if self.quantization_compressor is not None:
-            params_to_ignore = self.quantization_compressor.compression_param_names
-
         for prefix, module in model.named_modules():
             if prefix in module_to_scheme or prefix in sparse_compression_targets:
                 state_dict = module.state_dict(prefix=f"{prefix}.")
@@ -433,37 +428,37 @@ def decompress_model(self, model: Module):
                     # sparse_compression_targets are automatically inferred by this fn
                     generator = self.sparsity_compressor.decompress_from_state_dict(
                         state_dict,
-                        params_to_ignore=params_to_ignore,
                     )
-                    decompressed = dict()
-                    for _, decompressed_data in generator:
-                        decompressed.update(decompressed_data)
-                    state_dict = decompressed
+                    # generates (param_path, param_val)
+                    # of compressed and unused params
+                    state_dict = {key: value for key, value in generator}
 
                 # quantization second
                 if prefix in module_to_scheme:
                     generator = self.quantization_compressor.decompress_from_state_dict(
-                        state_dict,  # asdf
+                        state_dict,
                         names_to_scheme=module_to_scheme,
                     )
-                    decompressed = dict()
-                    for _, decompressed_data in generator:
-                        decompressed.update(decompressed_data)
-                    state_dict = decompressed
+                    # generates (mod_path, {param_name, param_val})
+                    # of compressed params only (ignores unused params)
+                    state_dict = {
+                        merge_names(module_path, param_name): param_value
+                        for module_path, compressed_data in generator
+                        for param_name, param_value in compressed_data.items()
+                    }
 
                 # remove any existing parameters
                 for name, _ in list(module.named_parameters()):
                     delattr(module, name)
 
                 # replace with decompressed parameters
                 for name, value in state_dict.items():
+                    name = name.removeprefix(f"{prefix}.")
                     param = torch.nn.Parameter(value, requires_grad=False)
                     register_offload_parameter(module, name, param)
 
                 module.quantization_status = QuantizationStatus.FROZEN
 
-        # apparently we only have logic for decompressing from a file...
-
     # ----- state dict compression pathways ----- #
 
     def compress(

src/compressed_tensors/compressors/sparse_compressors/base.py

@@ -16,7 +16,11 @@
 from typing import Dict, Generator, Optional, Set, Tuple
 
 from compressed_tensors.compressors.base import BaseCompressor
-from compressed_tensors.utils import get_nested_weight_mappings, merge_names, get_nested_mappings_from_state_dict
+from compressed_tensors.utils import (
+    get_nested_mappings_from_state_dict,
+    get_nested_weight_mappings,
+    merge_names,
+)
 from safetensors import safe_open
 from torch import Tensor
 from tqdm import tqdm
@@ -154,10 +158,14 @@ def decompress(
     def decompress_from_state_dict(
         self,
         state_dict: Dict[str, Tensor],
-        params_to_skip_load: Optional[Tuple] = None,
     ) -> Generator[Tuple[str, Dict[str, Tensor]], None, None]:
         """
-        Implemented to copy the pattern of 
+        Unlike `self.decompress`, this function does not need to explicitly skip params
+        via params_to_skip_load because it is more convenient for its only caller
+        (ModelCompressor.decompress_model) to retrieve all unused param keys
+
+        :param state_dict: state dict containing parameters to decompress
+        :return: Generator of (param_path, param_val)
         """
         weight_mappings, ignored_params = get_nested_mappings_from_state_dict(
             state_dict, self.compression_param_names, return_unmatched_params=True
@@ -168,17 +176,11 @@ def decompress_from_state_dict(
             for param_name, param_value in weight_mappings[module_path].items():
                 weight_data[param_name] = param_value
 
-            yield module_path, self.decompress_weight(weight_data)
-
-        for ignored_param_name, safe_path in ignored_params.items():
-            should_skip = False
-            if params_to_skip_load is not None:
-                for param_to_skip in params_to_skip_load:
-                    if param_to_skip in ignored_param_name:
-                        should_skip = True
+            decompressed = self.decompress_weight(weight_data)
+            yield merge_names(module_path, "weight"), decompressed
 
-            if not should_skip:
-                yield ignored_param_name, state_dict[ignored_param_name]
+        for ignored_param_path, ignored_param_value in ignored_params.items():
+            yield ignored_param_path, ignored_param_value
 
     @staticmethod
     def should_compress(name: str, expanded_targets: Optional[Set[str]] = None) -> bool:

src/compressed_tensors/compressors/sparse_compressors/sparse_24_bitmask.py

@@ -13,7 +13,7 @@
 # limitations under the License.
 
 from dataclasses import dataclass
-from typing import Dict, List, Tuple, Union
+from typing import Dict, Generator, List, Tuple, Union
 
 import torch
 from compressed_tensors.compressors.base import BaseCompressor
@@ -202,11 +202,7 @@ def sparse24_bitmask_decompress(
     decompressed_tensor = torch.zeros(original_shape, dtype=values.dtype)
     decompressed_tensor = decompressed_tensor.to(values.device)
     values = values.flatten()
-    if decompressed_tensor.dtype == FP8_DTYPE:
-        decompressed_tensor[bytemasks_unpacked] = values
-        decompressed_tensor = decompressed_tensor.cuda()
-    else:
-        decompressed_tensor[bytemasks_unpacked] = values
+    decompressed_tensor[bytemasks_unpacked] = values
     return decompressed_tensor
 
 

src/compressed_tensors/utils/safetensors_load.py

@@ -250,7 +250,9 @@ def get_nested_weight_mappings(
 
 
 def get_nested_mappings_from_state_dict(
-    state_dict, params_to_nest: Iterable[str], return_unmatched_params: bool = False,
+    state_dict,
+    params_to_nest: Iterable[str],
+    return_unmatched_params: bool = False,
 ) -> Union[NestedStateDictType, Tuple[NestedStateDictType, Dict[str, Tensor]]]:
     """
     Takes a state dict and returns a nested mapping from uncompressed
@@ -271,14 +273,15 @@ def get_nested_mappings_from_state_dict(
     """
     nested_weight_mappings = {}
     unmatched_params = {}
-    
+
     for key in state_dict.keys():
+        matched = False
         for param_name in params_to_nest:
-            dense_param = match_param_name(key, param_name)
-            if dense_param:
-                if dense_param not in nested_weight_mappings:
-                    nested_weight_mappings[dense_param] = {}
-                nested_weight_mappings[dense_param][param_name] = state_dict[key]
+            module_path = match_param_name(key, param_name)
+            if module_path:
+                if module_path not in nested_weight_mappings:
+                    nested_weight_mappings[module_path] = {}
+                nested_weight_mappings[module_path][param_name] = state_dict[key]
                 matched = True
         if return_unmatched_params and not matched:
             unmatched_params[key] = state_dict[key]

tests/test_compressors/model_compressors/test_model_compressor.py

@@ -16,12 +16,12 @@
 from copy import deepcopy
 from pathlib import Path
 
-from compressed_tensors.config.sparse_24_bitmask import Sparse24BitMaskConfig
 import pytest
 import torch
 import torch.nn as nn
 from compressed_tensors.compressors import ModelCompressor
 from compressed_tensors.config import SparsityCompressionConfig
+from compressed_tensors.config.sparse_24_bitmask import Sparse24BitMaskConfig
 from compressed_tensors.linear.compressed_linear import CompressedLinear
 from compressed_tensors.quantization import QuantizationConfig, QuantizationStatus
 from safetensors.torch import save_file
@@ -386,11 +386,11 @@ def _get_combined_config(s_config, q_config):
         (
             "nm-testing/llama2.c-stories42M-gsm8k-stacked-uncompressed",
             "float-quantized",
-            Sparse24BitMaskConfig(targets=["Linear"]),
+            "sparse-24-bitmask",
         ),
     ],
 )
-def test_compress_decompress_model(model_stub, q_format, s_config, tmpdir):
+def test_compress_model(model_stub, q_format, s_config, tmpdir):
     model = AutoModelForCausalLM.from_pretrained(model_stub, torch_dtype=torch.float32)
     compressor = ModelCompressor.from_pretrained_model(model, s_config, q_format)
 
@@ -407,64 +407,32 @@ def test_compress_decompress_model(model_stub, q_format, s_config, tmpdir):
     for key in compressed.keys():
         assert torch.all(compressed[key] == true_compressed[key]), f"{key}"
 
-    del compressed
-    del true_compressed
-
-    # -- decompress -- #
-
-    # reinstantiate compressor to mimic LLM Compressor flows
-    model.save_pretrained(tmpdir)
-    model = AutoModelForCausalLM.from_pretrained(tmpdir, torch_dtype=torch.float32)
-    compressor = ModelCompressor.from_pretrained_model(model, s_config, q_format)
-    
-    true_decompressed_model = AutoModelForCausalLM.from_pretrained(model_stub, torch_dtype=torch.float32)
-    compressor.decompress(tmpdir, true_decompressed_model)
-    true_decompressed = dict(true_decompressed_model.state_dict())
-
-    # decompress model
-    compressor.decompress_model(model)
-    decompressed = dict(model.state_dict())
-
-    # equivalent to decompressing from disk
-    assert decompressed.keys() == true_decompressed.keys()
-    for key in decompressed.keys():
-        mask = ~torch.isclose(decompressed[key], true_decompressed[key], rtol=1e-3, atol=1e-5)
-        print("Mismatched indices:", mask.nonzero(as_tuple=True))
-        print("a values:", decompressed[key][mask])
-        print("b values:", true_decompressed[key][mask])
-        assert torch.allclose(decompressed[key], true_decompressed[key], rtol=1e-3, atol=1e-5), f"{key}"
-    del true_decompressed
-
 
 @pytest.mark.parametrize(
-    "comp_stub,q_format,s_config",
+    "model_stub,comp_stub",
     [
-        # (
-        #     "nm-testing/llama2.c-stories42M-gsm8k-quantized-only-compressed",
-        #     "float-quantized",
-        #     None,
-        # ),
-        # (
-        #     "nm-testing/llama2.c-stories42M-gsm8k-sparse-only-compressed",
-        #     None,
-        #     "sparse-24-bitmask",
-        # ),
         (
+            "nm-testing/llama2.c-stories42M-gsm8k-quantized-only-uncompressed",
+            "nm-testing/llama2.c-stories42M-gsm8k-quantized-only-compressed",
+        ),
+        (
+            "nm-testing/llama2.c-stories42M-gsm8k-sparse-only-uncompressed",
+            "nm-testing/llama2.c-stories42M-gsm8k-sparse-only-compressed",
+        ),
+        (
+            "nm-testing/llama2.c-stories42M-gsm8k-stacked-uncompressed",
             "nm-testing/llama2.c-stories42M-gsm8k-stacked-compressed",
-            "float-quantized",
-            Sparse24BitMaskConfig(targets=["Linear"]),
         ),
     ],
 )
-def test_decompress_model(comp_stub, q_format, s_config):
+def test_decompress_model(model_stub, comp_stub):
+    from transformers.utils.quantization_config import CompressedTensorsConfig
+
+    # decompress from disk
     # NOTE: transformers adds extra zero points if run_compressed=False or w/ sparsity
     # https://github.com/huggingface/transformers/blob/main/src/transformers/quantizers/quantizer_compressed_tensors.py#L131-L133
     # however, decompression does not add zero points in non-asymmetric cases
     # in order to normalize for this effect in this test, we remove empty weight zps
-
-    from transformers.utils.quantization_config import CompressedTensorsConfig
-
-    # decompress from disk
     true_decompressed_model = AutoModelForCausalLM.from_pretrained(
         comp_stub,
         quantization_config=CompressedTensorsConfig(run_compressed=False),
@@ -474,24 +442,19 @@ def test_decompress_model(comp_stub, q_format, s_config):
     true_decompressed = remove_empty_weight_zero_points(true_decompressed)  # see above
 
     # decompress from memory
-    model = AutoModelForCausalLM.from_pretrained(
-        comp_stub,
-        quantization_config=CompressedTensorsConfig(run_compressed=True),
-        torch_dtype=torch.float32,
-    )
-    compressor = ModelCompressor.from_pretrained_model(model, s_config, q_format)
+    # NOTE there is no other way to load a compressed model into memory, since
+    # there is no way to turn off decompression for sparse models
+    # https://github.com/huggingface/transformers/blob/main/src/transformers/quantizers/quantizer_compressed_tensors.py#L133
+    model = AutoModelForCausalLM.from_pretrained(model_stub, torch_dtype=torch.float32)
+    compressor = ModelCompressor.from_pretrained(comp_stub)
+    compressor.compress_model(model)
     compressor.decompress_model(model)
     decompressed = dict(model.state_dict())
-    if "sparse" in str(s_config):
-        decompressed = remove_empty_weight_zero_points(decompressed)  # see above
 
     # equivalent to decompressing from disk
-    breakpoint()
     assert decompressed.keys() == true_decompressed.keys()
     for key in decompressed.keys():
-        if not torch.allclose(decompressed[key], true_decompressed[key]):
-            breakpoint()
-    del true_decompressed
+        assert torch.allclose(decompressed[key], true_decompressed[key])
 
 
 def remove_empty_weight_zero_points(state_dict):

tests/test_compressors/sparse_compressors/test_sparse_24_bitmask.py

@@ -47,6 +47,8 @@ def _validate_shard_shapes(sharded_values, sharded_bitmask, expected_shapes):
 
 def validate_compression(dense_matrix, decompressed_tensor):
     """Validate that the decompressed tensor matches the original dense matrix."""
+    if decompressed_tensor.device == FP8_DTYPE:
+        decompressed_tensor = decompressed_tensor.to("cuda")
     dense_matrix = dense_matrix.to(decompressed_tensor.device)
     assert dense_matrix.dtype == decompressed_tensor.dtype, "Dtype mismatch"
     assert dense_matrix.shape == decompressed_tensor.shape, "Shape mismatch"