Merge branch 'kylesayrs/map_module_to_scheme' into kylesayrs/reduce-quantized-compression-memory

Author: kylesayrs

src/compressed_tensors/compressors/model_compressors/model_compressor.py

@@ -19,18 +19,7 @@
 import re
 from contextlib import contextmanager
 from copy import deepcopy
-from typing import (
-    TYPE_CHECKING,
-    Any,
-    Callable,
-    Dict,
-    List,
-    Optional,
-    Set,
-    Tuple,
-    TypeVar,
-    Union,
-)
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, Set, TypeVar, Union
 
 import compressed_tensors
 import torch
@@ -44,7 +33,6 @@
 from compressed_tensors.compressors.base import BaseCompressor
 from compressed_tensors.compressors.sparse_compressors import DenseCompressor
 from compressed_tensors.config import CompressionFormat, SparsityCompressionConfig
-from compressed_tensors.linear.compressed_linear import CompressedLinear
 from compressed_tensors.quantization import (
     DEFAULT_QUANTIZATION_METHOD,
     QuantizationConfig,

src/compressed_tensors/compressors/quantized_compressors/base.py

@@ -14,7 +14,7 @@
 
 import logging
 from pathlib import Path
-from typing import Any, Dict, Generator, Optional, Tuple, Union
+from typing import Any, Dict, Generator, Tuple, Union
 
 import torch
 from compressed_tensors.compressors.base import BaseCompressor
@@ -119,10 +119,8 @@ def compress(
                     compressed_dict[prefix + key] = value.to(save_device)
 
             else:
-                # omit saving zero points for symmetric quantization
-                if name.endswith("zero_point") and not self._should_save_zp(
-                    name, names_to_scheme
-                ):
+                # omit saving zero points for symmetric or packed quantization
+                if name.endswith("zero_point") and self._skip_zp(name, names_to_scheme):
                     continue
 
                 # omit saving for g_idx if uninitialized
@@ -134,7 +132,7 @@ def compress(
 
         return compressed_dict
 
-    def _should_save_zp(
+    def _skip_zp(
         self, name: str, names_to_scheme: Dict[str, QuantizationScheme]
     ) -> bool:
         from compressed_tensors.compressors import PackedQuantizationCompressor
@@ -150,16 +148,16 @@ def _should_save_zp(
             args = scheme.output_activations
 
         symmetric = args.symmetric
-        packable_strats = [
+        packable_strategies = [
             QuantizationStrategy.GROUP.value,
             QuantizationStrategy.CHANNEL.value,
         ]
         packed = (
             isinstance(self, PackedQuantizationCompressor)
-            and args.strategy in packable_strats
+            and args.strategy in packable_strategies
         )
 
-        return not symmetric and not packed
+        return symmetric or packed
 
     def decompress(
         self,

src/compressed_tensors/utils/helpers.py

@@ -14,11 +14,10 @@
 
 import warnings
 from functools import wraps
-from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Union
+from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional
 
 import numpy
 import torch
-import tqdm
 from transformers import AutoConfig