File size: 7,002 Bytes
13a5289 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 | # Copyright (c) 2021 - present / Neuralmagic, Inc. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from abc import ABC, abstractmethod
from typing import Dict, Generator, Optional, Tuple, Union
import torch
from compressed_tensors.config import SparsityCompressionConfig
from compressed_tensors.quantization import QuantizationArgs, QuantizationConfig
from compressed_tensors.registry import RegistryMixin
from torch import Tensor
from torch.nn import Module
__all__ = ["BaseCompressor"]
class BaseCompressor(RegistryMixin, ABC):
"""
Base class representing a model compression algorithm. Each child class should
implement compression_param_info, compress_weight and decompress_weight.
Compressors support compressing/decompressing a full module state dict or a single
quantized PyTorch leaf module.
Model Load Lifecycle (run_compressed=False):
- ModelCompressor.decompress()
- apply_quantization_config()
- BaseCompressor.decompress()
Model Save Lifecycle:
- ModelCompressor.compress()
- BaseCompressor.compress()
Module Lifecycle (run_compressed=True):
- apply_quantization_config()
- compressed_module = CompressedLinear(module)
- initialize_module_for_quantization()
- BaseCompressor.compression_param_info()
- register_parameters()
- compressed_module.forward()
-compressed_module.decompress()
:param config: config specifying compression parameters
"""
def __init__(
self, config: Union[SparsityCompressionConfig, QuantizationConfig, None] = None
):
self.config = config
def compression_param_info(
self,
weight_shape: torch.Size,
quantization_args: Optional[QuantizationArgs] = None,
) -> Dict[str, Tuple[torch.Size, torch.dtype]]:
"""
Creates a dictionary of expected shapes and dtypes for each compression
parameter used by the compressor
:param weight_shape: uncompressed weight shape
:param quantization_args: quantization parameters for the weight
:return: dictionary mapping compressed parameter names to shape and dtype
"""
raise NotImplementedError()
@property
@abstractmethod
def compression_param_names(self) -> Tuple[str]:
"""
Returns a tuple of compression parameter names introduced by
the compressor during compression
"""
raise NotImplementedError()
@abstractmethod
def compress(
self,
model_state: Dict[str, Tensor],
**kwargs,
) -> Dict[str, Tensor]:
"""
Compresses a dense state dict
:param model_state: state dict of uncompressed model
:param kwargs: additional arguments for compression
:return: compressed state dict
"""
raise NotImplementedError()
@abstractmethod
def decompress(
self,
path_to_model_or_tensors: str,
device: str = "cpu",
**kwargs,
) -> Generator[Tuple[str, Tensor], None, None]:
"""
Reads a compressed state dict located at path_to_model_or_tensors
and returns a generator for sequentially decompressing back to a
dense state dict
:param path_to_model_or_tensors: path to compressed safetensors model (directory
with one or more safetensors files) or compressed tensors file
:param names_to_scheme: quantization args for each quantized weight
:param device: optional device to load intermediate weights into
:return: compressed state dict
"""
raise NotImplementedError()
def compress_module(self, module: Module) -> Optional[Dict[str, torch.Tensor]]:
"""
Compresses a single quantized leaf PyTorch module. If the module is not
quantized, this function has no effect.
:param module: PyTorch module to compress
:return: dictionary of compressed weight data, or None if module is not
quantized
"""
if not hasattr(module, "quantization_scheme"):
return None # module is not quantized
quantization_scheme = module.quantization_scheme
if not hasattr(quantization_scheme, "weights"):
return None # weights are not quantized
quantization_args = quantization_scheme.weights
weight = getattr(module, "weight", None)
weight_scale = getattr(module, "weight_scale", None)
weight_zero_point = getattr(module, "weight_zero_point", None)
return self.compress_weight(
weight=weight,
scale=weight_scale,
zero_point=weight_zero_point,
quantization_args=quantization_args,
)
def compress_weight(
self,
weight: Tensor,
**kwargs,
) -> Dict[str, torch.Tensor]:
"""
Compresses a single uncompressed weight
:param weight: uncompressed weight tensor
:param kwargs: additional arguments for compression
"""
raise NotImplementedError()
def decompress_module(self, module: Module):
"""
Decompresses a single compressed leaf PyTorch module. If the module is not
quantized, this function has no effect.
:param module: PyTorch module to decompress
:return: tensor of the decompressed weight, or None if module is not quantized
"""
if not hasattr(module, "quantization_scheme"):
return None # module is not quantized
quantization_scheme = module.quantization_scheme
if not hasattr(quantization_scheme, "weights"):
return None # weights are not quantized
quantization_args = quantization_scheme.weights
compressed_data = {}
for name, parameter in module.named_parameters():
compressed_data[name] = parameter
return self.decompress_weight(
compressed_data=compressed_data, quantization_args=quantization_args
)
def decompress_weight(
self, compressed_data: Dict[str, Tensor], **kwargs
) -> torch.Tensor:
"""
Decompresses a single compressed weight
:param compressed_data: dictionary of data needed for decompression
:param kwargs: additional arguments for decompression
:return: tensor of the decompressed weight
"""
raise NotImplementedError()
|