234 lines
8.7 KiB
Python
234 lines
8.7 KiB
Python
# Copyright (c) 2023 PaddlePaddle Authors. 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.
|
||
|
||
import copy
|
||
import json
|
||
from dataclasses import dataclass
|
||
|
||
quant_inference_mapping = {"avg": "abs_max", "abs_max_channel_wise": "abs_max_channel_wise", "abs_max": "abs_max"}
|
||
fp8_format_mapping = {
|
||
"hybrid": {"weight": "float8_e4m3fn", "activation": "float8_e4m3fn", "grad_output": "float8_e5m2"},
|
||
"e4m3": {"weight": "float8_e4m3fn", "activation": "float8_e4m3fn", "grad_output": "float8_e4m3fn"},
|
||
}
|
||
|
||
|
||
@dataclass
|
||
class QuantizationConfig:
|
||
"""
|
||
This is the configuration class to store quantization configuration.
|
||
Args:
|
||
weight_quantize_algo: Weight quantization algorithm.
|
||
quant_type: Quantization type applied to weight and activation, weight may still keep in float tensor.
|
||
shift: Whether the model applied the shift strategy.
|
||
smooth: Whether the model applied the smooth strategy.
|
||
shift_smooth_all_linears: Whether the model applied shift or smooth strategy for all linears.
|
||
quant_round_type: The quant round type, 0:-rounding to nearest ties to even, 1: -rounding to nearest ties away from zero.
|
||
llm_int8_threshold: The threshold for llm.int8 quantization.
|
||
qlora_weight_double_quant: Whether quant weight scale.
|
||
qlora_weight_blocksize: Block size for weight quantization.
|
||
qlora_weight_double_quant_block_size: Block size for quant_scale of weight quant_scale.
|
||
weight_quant_method: The method for weight quantization.
|
||
act_quant_method: The method for activation quantization.
|
||
apply_online_actscale_step: Use online (per-step) activation scales for the first N steps. During these steps, activation scales are also collected to compute their mean for later use.
|
||
"""
|
||
|
||
def __init__(
|
||
self,
|
||
weight_quantize_algo=None,
|
||
quant_type=None,
|
||
shift=False,
|
||
smooth=False,
|
||
shift_smooth_all_linears=False,
|
||
quant_round_type=0,
|
||
llm_int8_threshold=6.0,
|
||
qlora_weight_double_quant=False,
|
||
qlora_weight_blocksize=64,
|
||
qlora_weight_double_quant_block_size=256,
|
||
weight_quant_method="abs_max_channel_wise",
|
||
act_quant_method="abs_max",
|
||
activation_scheme=None,
|
||
fmt=None,
|
||
quant_method=None,
|
||
weight_block_size=None,
|
||
dtype=None,
|
||
ignore_modules=None,
|
||
group_size=-1,
|
||
apply_hadamard=False,
|
||
hadamard_block_size=32,
|
||
quant_input_grad=False,
|
||
quant_weight_grad=False,
|
||
apply_online_actscale_step=200,
|
||
actscale_moving_rate=0.01,
|
||
fp8_format_type="hybrid",
|
||
scale_epsilon=1e-8,
|
||
**kwargs,
|
||
):
|
||
if weight_quantize_algo is not None:
|
||
if isinstance(weight_quantize_algo, dict):
|
||
if any(
|
||
algo
|
||
not in [
|
||
"weight_only_int8",
|
||
"weight_only_int4",
|
||
"llm.int8",
|
||
"a8w8",
|
||
"nf4",
|
||
"fp4",
|
||
"a8w8linear",
|
||
"a8w4linear",
|
||
"fp8linear",
|
||
]
|
||
for algo in weight_quantize_algo
|
||
):
|
||
raise ValueError(
|
||
f"weight_quantize_algo:{weight_quantize_algo.keys()} not in supported list ['weight_only_int8', 'weight_only_int4', 'llm.int8', 'a8w8', 'nf4', 'fp4']"
|
||
)
|
||
elif weight_quantize_algo not in [
|
||
"weight_only_int8",
|
||
"weight_only_int4",
|
||
"llm.int8",
|
||
"a8w8",
|
||
"nf4",
|
||
"fp4",
|
||
"a8w8linear",
|
||
"a8w4linear",
|
||
"fp8linear",
|
||
]:
|
||
raise ValueError(
|
||
f"weight_quantize_algo:{weight_quantize_algo} not in supported list ['weight_only_int8', 'weight_only_int4', 'llm.int8', 'a8w8', 'nf4', 'fp4']"
|
||
)
|
||
if quant_type is not None and quant_type not in [
|
||
"weight_only_int8",
|
||
"weight_only_int4",
|
||
"a8w8",
|
||
"a8w8c8",
|
||
"a8w8_fp8",
|
||
"a8w8c8_fp8",
|
||
]:
|
||
raise ValueError(
|
||
f"quant_type:{quant_type} not in supported list ['weight_only_int8', 'weight_only_int4', 'a8w8', 'a8w8c8', 'a8w8_fp8', 'a8w8c8_fp8']"
|
||
)
|
||
self.weight_quantize_algo = weight_quantize_algo
|
||
self.quant_type = quant_type
|
||
self.shift = shift
|
||
self.smooth = smooth
|
||
self.shift = shift
|
||
self.shift_smooth_all_linears = shift_smooth_all_linears
|
||
self.quant_round_type = quant_round_type
|
||
self.llm_int8_threshold = llm_int8_threshold
|
||
self.qlora_weight_double_quant = qlora_weight_double_quant
|
||
self.qlora_weight_blocksize = qlora_weight_blocksize
|
||
self.weight_quant_method = weight_quant_method
|
||
self.act_quant_method = quant_inference_mapping[act_quant_method]
|
||
self.qlora_weight_double_quant_block_size = qlora_weight_double_quant_block_size
|
||
self.activation_scheme = activation_scheme
|
||
self.fmt = fmt
|
||
self.quant_method = quant_method
|
||
self.weight_block_size = weight_block_size
|
||
self.dtype = dtype
|
||
self.ignore_modules = ignore_modules
|
||
self.group_size = group_size
|
||
self.apply_hadamard = apply_hadamard
|
||
self.hadamard_block_size = hadamard_block_size
|
||
self.quant_input_grad = quant_input_grad
|
||
self.quant_weight_grad = quant_weight_grad
|
||
self.apply_online_actscale_step = apply_online_actscale_step
|
||
self.actscale_moving_rate = actscale_moving_rate
|
||
self.fp8_format_type = fp8_format_type
|
||
self.scale_epsilon = scale_epsilon
|
||
|
||
@property
|
||
def fp8_format(self):
|
||
return fp8_format_mapping[self.fp8_format_type]
|
||
|
||
def is_weight_quantize(self):
|
||
if isinstance(self.weight_quantize_algo, dict):
|
||
return True
|
||
elif self.weight_quantize_algo in [
|
||
"weight_only_int8",
|
||
"weight_only_int4",
|
||
"llm.int8",
|
||
"nf4",
|
||
"fp4",
|
||
"a8w8",
|
||
"a8w8linear",
|
||
"a8w4linear",
|
||
"fp8linear",
|
||
]:
|
||
return True
|
||
else:
|
||
return False
|
||
|
||
def is_support_merge_tensor_parallel(self):
|
||
if self.weight_quantize_algo in ["weight_only_int8", "weight_only_int4", "llm.int8", "a8w8"]:
|
||
return False
|
||
else:
|
||
return True
|
||
|
||
@classmethod
|
||
def from_dict(cls, config_dict, return_unused_kwargs=False, **kwargs):
|
||
"""
|
||
Instantiates QuantizationConfig from dict
|
||
"""
|
||
config = cls(**config_dict)
|
||
|
||
to_remove = []
|
||
for key, value in kwargs.items():
|
||
if hasattr(config, key):
|
||
setattr(config, key, value)
|
||
to_remove.append(key)
|
||
for key in to_remove:
|
||
kwargs.pop(key, None)
|
||
|
||
if return_unused_kwargs:
|
||
return config, kwargs
|
||
else:
|
||
return config
|
||
|
||
def to_json_file(self, json_file_path):
|
||
"""
|
||
Save this instance to a JSON file.
|
||
"""
|
||
with open(json_file_path, "w", encoding="utf-8") as f:
|
||
f.write(json.dumps(self.to_dict(), indent=2, sort_keys=True) + "\n")
|
||
|
||
def to_dict(self):
|
||
return copy.deepcopy(self.__dict__)
|
||
|
||
def __repr__(self):
|
||
return f"{self.__class__.__name__} {self.to_json_string()}"
|
||
|
||
def to_json_string(self, use_diff=True):
|
||
if use_diff is True:
|
||
config_dict = self.to_diff_dict()
|
||
else:
|
||
config_dict = self.to_dict()
|
||
return json.dumps(config_dict, indent=2, sort_keys=True) + "\n"
|
||
|
||
def to_diff_dict(self):
|
||
config_dict = self.to_dict()
|
||
|
||
# get the default config dict
|
||
|
||
default_config_dict = QuantizationConfig().to_dict()
|
||
|
||
serializable_config_dict = {}
|
||
|
||
# only serialize values that differ from the default config
|
||
for key, value in config_dict.items():
|
||
if value != default_config_dict[key]:
|
||
serializable_config_dict[key] = value
|
||
|
||
return serializable_config_dict
|