clean up

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

Author: kylesayrs

src/compressed_tensors/compressors/model_compressors/model_compressor.py

@@ -19,11 +19,20 @@
 import re
 from contextlib import contextmanager
 from copy import deepcopy
-from typing import TYPE_CHECKING, Any, Dict, List, Optional, Set, Tuple, TypeVar, Union, Callable
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Callable,
+    Dict,
+    List,
+    Optional,
+    Set,
+    Tuple,
+    TypeVar,
+    Union,
+)
 
 import compressed_tensors
-from compressed_tensors.linear.compressed_linear import CompressedLinear
-from compressed_tensors.utils.offload import update_offload_parameter
 import torch
 import transformers
 from compressed_tensors.base import (
@@ -34,6 +43,7 @@
 )
 from compressed_tensors.compressors.base import BaseCompressor
 from compressed_tensors.config import CompressionFormat, SparsityCompressionConfig
+from compressed_tensors.linear.compressed_linear import CompressedLinear
 from compressed_tensors.quantization import (
     DEFAULT_QUANTIZATION_METHOD,
     QuantizationConfig,
@@ -50,12 +60,14 @@
 from compressed_tensors.utils import (
     get_safetensors_folder,
     merge_names,
+    module_replace_dfs,
     update_parameter_data,
 )
 from compressed_tensors.utils.helpers import (
     fix_fsdp_module_name,
     is_compressed_tensors_config,
 )
+from compressed_tensors.utils.offload import update_offload_parameter
 from torch import Tensor
 from torch.nn import Module
 from tqdm import tqdm
@@ -67,31 +79,6 @@
 
 _LOGGER: logging.Logger = logging.getLogger(__name__)
 
-def module_replace_dfs(
-    module: Module,
-    func: Callable[[Module], Module],
-    pre: bool = True,
-    progress: Union[bool, tqdm] = False,
-) -> Module:
-    if progress is True:
-        total = len(list(module.modules()))
-        progress = tqdm(total=total)
-
-    if pre:
-        module = func(module)
-
-    for name, child in list(module.named_children()):
-        module.add_module(name, module_replace_dfs(child, func, pre, progress))
-
-    if not pre:
-        module = func(module)
-
-    if isinstance(progress, tqdm):
-        progress.update(1)
-
-    return module
-
-
 
 if TYPE_CHECKING:
     # dummy type if not available from transformers
@@ -385,36 +372,35 @@ def get_unexpected_file_keys(self, model: Module) -> List[str]:
                 )
 
         return list(unexpected_keys)
-    
+
     def apply_compression_status(self, model: Module) -> Module:
         quantization_format = self.quantization_config.format
 
         def replace_with_compressed(module: Module) -> Module:
             scheme = getattr(module, "quantization_scheme", None)
             if isinstance(module, torch.nn.Linear) and scheme is not None:
-                #compressed_state_dict_2 = self.compress(module)  # debug
+                # compressed_state_dict_2 = self.compress(module)  # debug
 
                 module = CompressedLinear.from_linear(
                     module,
                     quantization_scheme=scheme,
-                    quantization_format=quantization_format
+                    quantization_format=quantization_format,
                 )
-                state_dict = module.compressor.compress(module.state_dict(), {"": scheme})  # added by compressed linear
+                state_dict = module.compressor.compress(
+                    module.state_dict(), {"": scheme}
+                )  # added by compressed linear
 
                 for name, value in state_dict.items():
                     update_offload_parameter(module, name, value)
 
             return module
 
-
         progress = tqdm(total=len(list(model.modules())))
         return module_replace_dfs(model, replace_with_compressed, progress=progress)
 
     def compress(
         self, model: Module, state_dict: Optional[Dict[str, Tensor]] = None
     ) -> Dict[str, Tensor]:
-        from torch.profiler import profile, ProfilerActivity
-        from .track_tensor_memory import TrackTensorAllocations
         """
         Compresses a dense state dict or model with sparsity and/or quantization
 
@@ -427,21 +413,16 @@ def compress(
             state_dict = model.state_dict()
 
         if self.quantization_compressor is not None:
-            #with profile(activities=[ProfilerActivity.CUDA], profile_memory=True, record_shapes=True, with_stack=True) as prof:
-            #with TrackTensorAllocations() as prof:
             module_to_scheme = map_module_to_scheme(model)
             state_dict = self.quantization_compressor.compress(
                 state_dict, names_to_scheme=module_to_scheme
             )
-            # if self.quantization_config.format != CompressionFormat.dense.value:
-            #     self.quantization_config.quantization_status = (
-            #         QuantizationStatus.COMPRESSED
-            #     )
-
-            #prof.export_memory_timeline("memory.html")
-            #print(prof.key_averages().table(sort_by="self_device_memory_usage", row_limit=3))
-            #breakpoint()
-            return state_dict
+
+            # TODO: consider sparse compression to also be compression
+            if self.quantization_config.format != CompressionFormat.dense.value:
+                self.quantization_config.quantization_status = (
+                    QuantizationStatus.COMPRESSED
+                )
 
         if self.sparsity_compressor is not None:
             sparse_compression_targets: Set[str] = expand_target_names(

src/compressed_tensors/compressors/model_compressors/track_tensor_memory.py

@@ -76,11 +76,12 @@ def compress(
         """
         Compresses a dense state dict
 
-        :param model_state: state dict of uncompressed model, consumed by compression
+        :param model_state: state dict of uncompressed model
         :param names_to_scheme: quantization args for each quantized weight, needed for
             quantize function to calculate bit depth
         :return: compressed state dict
         """
+        compressed_dict = {}
         save_device = "cpu"
 
         uncompressed_names = list(model_state.keys())
@@ -98,7 +99,7 @@ def compress(
 
                 # is scale does not exist, then weight cannot be compressed
                 if scale is None:
-                    model_state[name] = value.to(save_device)
+                    compressed_dict[name] = value.to(save_device)
                     continue
 
                 # compress values on cpu (memory movement too expensive)
@@ -116,22 +117,22 @@ def compress(
                 # update state dict
                 del model_state[name]
                 for key, value in compressed_values.items():
-                    model_state[prefix + key] = value.to(save_device)
+                    compressed_dict[prefix + key] = value.to(save_device)
 
             else:
                 # omit saving zero points for symmetric quantization
                 if name.endswith("zero_point") and _is_symmetric(name, names_to_scheme):
-                    del model_state[name]
+                    continue
 
                 # omit saving for g_idx if uninitialized
                 # TODO: does this case actually occur?
                 elif name.endswith("g_idx") and torch.any(value <= -1):
-                    del model_state[name]
+                    continue
 
                 else:
-                    model_state[name] = value.to(save_device)
+                    compressed_dict[name] = value.to(save_device)
 
-        return model_state
+        return compressed_dict
 
     def decompress(
         self,

src/compressed_tensors/compressors/quantized_compressors/base.py

@@ -187,7 +187,7 @@ def pack_to_int32(value: torch.Tensor, num_bits: int) -> torch.Tensor:
 
     # convert back to signed and torch
     packed = np.ascontiguousarray(packed).view(np.int32)
-    return torch.Tensor(torch.from_numpy(packed))
+    return torch.from_numpy(packed)
 
 
 def unpack_from_int32(

src/compressed_tensors/compressors/quantized_compressors/pack_quantized.py

@@ -15,7 +15,6 @@
 import warnings
 from typing import Dict, Tuple
 
-from compressed_tensors.utils.offload import get_execution_device
 import torch
 from compressed_tensors.compressors.base import BaseCompressor
 from compressed_tensors.quantization import (
@@ -24,6 +23,7 @@
     initialize_module_for_quantization,
 )
 from compressed_tensors.utils import register_offload_parameter
+from compressed_tensors.utils.offload import get_execution_device
 from torch import Tensor
 from torch.nn import Parameter
 from torch.nn.functional import linear

src/compressed_tensors/linear/compressed_linear.py

@@ -322,7 +322,6 @@ def calculate_compression_ratio(model: Module) -> float:
     :param model: pytorch module to calculate compression ratio for
     :return: compression ratio of the whole model
     """
-    return 0.0
     total_compressed = 0.0
     total_uncompressed = 0.0
     for name, submodule in tqdm(

src/compressed_tensors/quantization/utils/helpers.py

@@ -14,10 +14,11 @@
 
 import warnings
 from functools import wraps
-from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional
+from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Union
 
 import numpy
 import torch
+import tqdm
 from transformers import AutoConfig
 
 
@@ -39,6 +40,7 @@
     "pack_bitmasks",
     "unpack_bitmasks",
     "remove_suffix",
+    "module_replace_dfs",
 ]
 
 FSDP_WRAPPER_NAME = "_fsdp_wrapped_module"
@@ -335,3 +337,28 @@ def remove_suffix(value: str, suffix: str) -> str:
     # can replace with str.removesuffix in python3.9+
     assert value.endswith(suffix)
     return value[: -len(suffix)]
+
+
+def module_replace_dfs(
+    module: torch.nn.Module,
+    func: Callable[[torch.nn.Module], torch.nn.Module],
+    pre: bool = True,
+    progress: Union[bool, tqdm.tqdm] = False,
+) -> torch.nn.Module:
+    if progress is True:
+        total = len(list(module.modules()))
+        progress = tqdm.tqdm(total=total)
+
+    if pre:
+        module = func(module)
+
+    for name, child in list(module.named_children()):
+        module.add_module(name, module_replace_dfs(child, func, pre, progress))
+
+    if not pre:
+        module = func(module)
+
+    if isinstance(progress, tqdm.tqdm):
+        progress.update(1)
+
+    return module