File size: 6,010 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 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 | # 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 copy import deepcopy
from typing import Any, Dict, List, Optional
from compressed_tensors.quantization.quant_args import (
QuantizationArgs,
QuantizationStrategy,
QuantizationType,
)
from pydantic import BaseModel, model_validator
__all__ = [
"QuantizationScheme",
"preset_name_to_scheme",
"is_preset_scheme",
]
class QuantizationScheme(BaseModel):
"""
Set of QuantizationArgs defining how the weights, inputs and outputs of target list
of modules should be quantized
:param targets: list of modules to apply the QuantizationArgs to, can be layer
names, layer types or a regular expression, typically ["Linear"]
:param weights: quantization config for layer weights
:param input_activations: quantization config for layer inputs
:param output_activations: quantization config for layer outputs
"""
targets: List[str]
weights: Optional[QuantizationArgs] = None
input_activations: Optional[QuantizationArgs] = None
output_activations: Optional[QuantizationArgs] = None
@model_validator(mode="after")
def validate_model_after(model: "QuantizationArgs") -> Dict[str, Any]:
inputs = model.input_activations
outputs = model.output_activations
if inputs is not None:
if inputs.actorder is not None:
raise ValueError("Cannot apply actorder to input activations")
if outputs is not None:
if outputs.actorder is not None:
raise ValueError("Cannot apply actorder to output activations")
return model
"""
Pre-Set Quantization Scheme Args
"""
def preset_name_to_scheme(name: str, targets: List[str]) -> QuantizationScheme:
"""
:param name: preset quantization settings name. must exist in upper case in
PRESET_SCHEMES
:param targets: list of quantization targets to be passed to the Scheme
:return: new QuantizationScheme for a given name with the given targets
"""
name = name.upper()
if name not in PRESET_SCHEMES:
raise KeyError(
f"Unknown preset scheme name {name}, "
f"available names: {list(PRESET_SCHEMES.keys())}"
)
scheme_args = deepcopy(PRESET_SCHEMES[name]) # deepcopy to avoid args references
return QuantizationScheme(
targets=targets,
**scheme_args,
)
def is_preset_scheme(name: str) -> bool:
"""
:param name: preset quantization settings name
:return: True if the name is a preset scheme name
"""
return name.upper() in PRESET_SCHEMES
UNQUANTIZED = dict()
# 8 bit integer weights and 8 bit activations quantization
INT8_W8A8 = dict(
weights=QuantizationArgs(
num_bits=8,
type=QuantizationType.INT,
strategy=QuantizationStrategy.CHANNEL,
symmetric=True,
dynamic=False,
),
input_activations=QuantizationArgs(
num_bits=8,
type=QuantizationType.INT,
strategy=QuantizationStrategy.TOKEN,
symmetric=True,
dynamic=True,
observer=None,
),
)
# 8 bit integer weights only quantization
W8A16 = dict(
weights=QuantizationArgs(
num_bits=8,
type=QuantizationType.INT,
strategy=QuantizationStrategy.CHANNEL,
symmetric=True,
dynamic=False,
),
)
# 4 bit integer weights only quantization
W4A16 = dict(
weights=QuantizationArgs(
num_bits=4,
type=QuantizationType.INT,
strategy=QuantizationStrategy.GROUP,
group_size=128,
symmetric=True,
dynamic=False,
),
)
# 4 bit integer weights and 8 bit activations quantization
INT8_W4A8 = dict(
weights=QuantizationArgs(
num_bits=4,
type=QuantizationType.INT,
group_size=128,
strategy=QuantizationStrategy.GROUP,
symmetric=True,
dynamic=False,
),
input_activations=QuantizationArgs(
num_bits=8,
type=QuantizationType.INT,
strategy=QuantizationStrategy.TOKEN,
symmetric=True,
dynamic=True,
observer=None,
),
)
# FP8 weights and FP8 activations quantization
FP8 = dict(
weights=QuantizationArgs(
num_bits=8,
type=QuantizationType.FLOAT,
strategy=QuantizationStrategy.TENSOR,
symmetric=True,
dynamic=False,
),
input_activations=QuantizationArgs(
num_bits=8,
type=QuantizationType.FLOAT,
strategy=QuantizationStrategy.TENSOR,
symmetric=True,
dynamic=False,
),
)
# FP8 weights and FP8 dynamic activations quantization
FP8_DYNAMIC = dict(
weights=QuantizationArgs(
num_bits=8,
type=QuantizationType.FLOAT,
strategy=QuantizationStrategy.CHANNEL,
symmetric=True,
dynamic=False,
),
input_activations=QuantizationArgs(
num_bits=8,
type=QuantizationType.FLOAT,
strategy=QuantizationStrategy.TOKEN,
symmetric=True,
dynamic=True,
observer=None,
),
)
PRESET_SCHEMES = {
# Unquantized (no-op)
"UNQUANTIZED": UNQUANTIZED,
# Integer weight only schemes
"W8A16": W8A16,
"W4A16": W4A16,
# Integer weight and activation schemes
"W8A8": INT8_W8A8,
"INT8": INT8_W8A8, # alias for W8A8
"W4A8": INT8_W4A8,
# Float weight and activation schemes
"FP8": FP8,
"FP8_DYNAMIC": FP8_DYNAMIC,
}
|