wehub-resource-sync
495 lines
18 KiB
Python
495 lines
18 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.
|
||
"""Define some layers used to export quantization model with ONNX style."""
|
||
|
||
from __future__ import annotations
|
||
|
||
import abc
|
||
|
||
import paddle
|
||
from paddle import _C_ops, _legacy_C_ops
|
||
from paddle.base import unique_name
|
||
from paddle.framework import in_dynamic_mode, in_pir_mode
|
||
|
||
from ..layer.layers import Layer
|
||
|
||
|
||
def fake_fp8_quant(input, scale, axis=-1, type='e4m3'):
|
||
# only support channelwise or tensorwise
|
||
if axis >= 0:
|
||
shape = [1] * len(input.shape)
|
||
shape[axis] = scale.numel()
|
||
scale = scale.reshape(shape)
|
||
inp = input.astype("float32")
|
||
|
||
if type == 'e4m3':
|
||
return paddle.cast(
|
||
(inp * 448 / scale).clip(-448, 448), "float8_e4m3fn"
|
||
).astype(input.dtype) # clip then cast
|
||
elif type == 'e5m2':
|
||
return paddle.cast(
|
||
(inp * 57344 / scale).clip(-57344, 57344), "float8_e5m2"
|
||
).astype(input.dtype) # clip then cast
|
||
else:
|
||
raise NotImplementedError("only support e4m3 or e5m2 now")
|
||
|
||
|
||
def fake_fp8_dequant(input, scale, axis=-1, type='e4m3'):
|
||
# only support channelwise or tensorwise
|
||
if axis >= 0:
|
||
shape = [1] * len(input.shape)
|
||
shape[axis] = scale.numel()
|
||
scale = scale.reshape(shape)
|
||
if type == 'e4m3':
|
||
return (input.astype("float32") / 448 * scale).astype(input.dtype)
|
||
elif type == 'e5m2':
|
||
return (input.astype("float32") / 57344 * scale).astype(input.dtype)
|
||
else:
|
||
raise NotImplementedError("only support e4m3 or e5m2 now")
|
||
|
||
|
||
class LinearQuanterDequanter(Layer):
|
||
def __init__(self, quanter, dequanter):
|
||
super().__init__()
|
||
self._quanter = quanter
|
||
self._dequanter = dequanter
|
||
|
||
def forward(self, input):
|
||
out = input
|
||
if self._quanter is not None:
|
||
out = self._quanter(out)
|
||
if self._dequanter is not None:
|
||
out = self._dequanter(out)
|
||
return out
|
||
|
||
@staticmethod
|
||
def from_quanter(quanter):
|
||
assert quanter is not None
|
||
return LinearQuanterDequanter(
|
||
LinearQuanter.from_quanter(quanter),
|
||
LinearDequanter.from_quanter(quanter),
|
||
)
|
||
|
||
|
||
class LinearQuanter(Layer):
|
||
def __init__(
|
||
self,
|
||
scales,
|
||
zero_point=None,
|
||
quant_axis=None,
|
||
bit_length=8,
|
||
group_size=128,
|
||
):
|
||
super().__init__()
|
||
scales = paddle.to_tensor(scales, dtype="float32")
|
||
scale_attr = paddle.framework.ParamAttr(
|
||
name=paddle.utils.unique_name.generate('quant_dequant.scale'),
|
||
initializer=paddle.nn.initializer.Constant(1.0),
|
||
trainable=False,
|
||
)
|
||
self._scales = self.create_parameter(
|
||
shape=scales.shape, attr=scale_attr, dtype="float32"
|
||
)
|
||
self._scales.set_value(scales)
|
||
self.in_accum = paddle.to_tensor(0.0, dtype="float32")
|
||
self.in_state = paddle.to_tensor(0.0, dtype="float32")
|
||
zero_point = zero_point if zero_point is not None else 0.0
|
||
zero_point = paddle.to_tensor(zero_point, dtype="float32")
|
||
zp_attr = paddle.framework.ParamAttr(
|
||
name=paddle.utils.unique_name.generate('quant_dequant.zero_point'),
|
||
initializer=paddle.nn.initializer.Constant(0.0),
|
||
trainable=False,
|
||
)
|
||
self._zero_point = self.create_parameter(
|
||
shape=zero_point.shape, attr=zp_attr, dtype="float32"
|
||
)
|
||
self._zero_point.set_value(zero_point)
|
||
self._quant_axis = -1 if quant_axis is None else quant_axis
|
||
self._bit_length = bit_length
|
||
self._group_size = group_size
|
||
if isinstance(self._bit_length, tuple):
|
||
if (
|
||
self._bit_length[0] == 4
|
||
and self._bit_length[1] == 3
|
||
and len(self._bit_length) == 2
|
||
):
|
||
self._qmin = -1 * 448
|
||
self._qmax = 448
|
||
elif (
|
||
self._bit_length[0] == 5
|
||
and self._bit_length[1] == 2
|
||
and len(self._bit_length) == 2
|
||
):
|
||
self._qmin = -1 * 57344
|
||
self._qmax = 57344
|
||
else:
|
||
raise NotImplementedError(
|
||
"Currently, only float8_e4m3 and float8_e5m2 formats are supported. Please set quant_bits to (4,3) or (5,2) for the corresponding format."
|
||
)
|
||
else:
|
||
self._qmax = (1 << (self._bit_length - 1)) - 1
|
||
self._qmin = -1 * self._qmax - 1
|
||
if isinstance(self._bit_length, tuple):
|
||
self._bit_length = self._bit_length[0] + self._bit_length[1] + 1
|
||
|
||
def forward(self, input):
|
||
if in_dynamic_mode():
|
||
if self._qmax == 448:
|
||
return fake_fp8_quant(
|
||
input, self._scales, self._quant_axis, type='e4m3'
|
||
)
|
||
elif self._qmax == 57344:
|
||
return fake_fp8_quant(
|
||
input, self._scales, self._quant_axis, type='e5m2'
|
||
)
|
||
elif len(self._scales.shape) > 1:
|
||
if self._zero_point.sum() != 0:
|
||
quant_weight = paddle.clip(
|
||
paddle.round(input.cast('float32') / self._scales)
|
||
+ self._zero_point,
|
||
self._qmin,
|
||
self._qmax,
|
||
)
|
||
else:
|
||
new_s = paddle.repeat_interleave(
|
||
self._scales, self._group_size, 0
|
||
)
|
||
new_zp = paddle.repeat_interleave(
|
||
self._zero_point, self._group_size, 0
|
||
)
|
||
quant_weight = paddle.clip(
|
||
paddle.round(input.cast('float32') / new_s * self._qmax)
|
||
+ new_zp,
|
||
self._qmin,
|
||
self._qmax,
|
||
)
|
||
return quant_weight.cast(input.dtype)
|
||
|
||
return _legacy_C_ops.quantize_linear(
|
||
input.cast('float32'),
|
||
self._scales,
|
||
self._zero_point,
|
||
"quant_axis",
|
||
self._quant_axis,
|
||
"bit_length",
|
||
self._bit_length,
|
||
"qmin",
|
||
self._qmin,
|
||
"qmax",
|
||
self._qmax,
|
||
).cast(input.dtype)
|
||
if in_pir_mode():
|
||
input.stop_gradient = True
|
||
quant_out = paddle.pir.core.create_persistable_value(
|
||
dtype='float32',
|
||
shape=input.shape,
|
||
name=unique_name.generate("quant_out"),
|
||
initializer=paddle.nn.initializer.Constant(0.0),
|
||
stop_gradient=True,
|
||
)
|
||
# TODO(xiaoluomi): need to add only observer pass for quantize_linear
|
||
quant_out, out_state, out_accum, out_scale = _C_ops.quantize_linear(
|
||
input,
|
||
self._scales,
|
||
self._zero_point,
|
||
self.in_accum,
|
||
self.in_state,
|
||
self._quant_axis,
|
||
self._bit_length,
|
||
self._qmin,
|
||
self._qmax,
|
||
0,
|
||
True,
|
||
False,
|
||
)
|
||
return quant_out
|
||
else:
|
||
out = self._helper.create_variable_for_type_inference(input.dtype)
|
||
self._helper.append_op(
|
||
type='quantize_linear',
|
||
inputs={
|
||
'X': input,
|
||
'Scale': self._scales,
|
||
'ZeroPoint': self._zero_point,
|
||
},
|
||
outputs={'Y': out},
|
||
attrs={
|
||
'quant_axis': self._quant_axis,
|
||
'bit_length': self._bit_length,
|
||
'qmin': self._qmin,
|
||
'qmax': self._qmax,
|
||
},
|
||
)
|
||
return out
|
||
|
||
@staticmethod
|
||
def from_quanter(quanter):
|
||
return LinearQuanter(
|
||
quanter.scales(),
|
||
zero_point=quanter.zero_points(),
|
||
quant_axis=quanter.quant_axis(),
|
||
bit_length=quanter.bit_length(),
|
||
)
|
||
|
||
|
||
class LinearDequanter(Layer):
|
||
def __init__(
|
||
self,
|
||
scales,
|
||
zero_point=None,
|
||
quant_axis=None,
|
||
bit_length=8,
|
||
group_size=128,
|
||
):
|
||
super().__init__()
|
||
scales = paddle.to_tensor(scales, dtype="float32")
|
||
scale_attr = paddle.framework.ParamAttr(
|
||
name=paddle.utils.unique_name.generate('quant_dequant.scale'),
|
||
initializer=paddle.nn.initializer.Constant(1.0),
|
||
trainable=False,
|
||
)
|
||
self._scales = self.create_parameter(
|
||
shape=scales.shape, attr=scale_attr, dtype="float32"
|
||
)
|
||
self._scales.set_value(scales)
|
||
self.in_accum = paddle.to_tensor(0.0, dtype="float32")
|
||
self.in_state = paddle.to_tensor(0.0, dtype="float32")
|
||
zero_point = zero_point if zero_point is not None else 0.0
|
||
zero_point = paddle.to_tensor(zero_point, dtype="float32")
|
||
zp_attr = paddle.framework.ParamAttr(
|
||
name=paddle.utils.unique_name.generate('quant_dequant.zero_point'),
|
||
initializer=paddle.nn.initializer.Constant(0.0),
|
||
trainable=False,
|
||
)
|
||
self._zero_point = self.create_parameter(
|
||
shape=zero_point.shape, attr=zp_attr, dtype="float32"
|
||
)
|
||
self._zero_point.set_value(zero_point)
|
||
self._quant_axis = -1 if quant_axis is None else quant_axis
|
||
self._bit_length = bit_length
|
||
self._group_size = group_size
|
||
if isinstance(self._bit_length, tuple):
|
||
if (
|
||
self._bit_length[0] == 4
|
||
and self._bit_length[1] == 3
|
||
and len(self._bit_length) == 2
|
||
):
|
||
self._qmin = -1 * 448
|
||
self._qmax = 448
|
||
elif (
|
||
self._bit_length[0] == 5
|
||
and self._bit_length[1] == 2
|
||
and len(self._bit_length) == 2
|
||
):
|
||
self._qmin = -1 * 57344
|
||
self._qmax = 57344
|
||
else:
|
||
raise NotImplementedError(
|
||
"Currently, only float8_e4m3 and float8_e5m2 formats are supported. Please set quant_bits to (4,3) or (5,2) for the corresponding format."
|
||
)
|
||
else:
|
||
self._qmax = (1 << (self._bit_length - 1)) - 1
|
||
self._qmin = -1 * self._qmax - 1
|
||
if isinstance(self._bit_length, tuple):
|
||
self._bit_length = self._bit_length[0] + self._bit_length[1] + 1
|
||
|
||
def forward(self, input):
|
||
if in_dynamic_mode():
|
||
if self._qmax == 448:
|
||
return fake_fp8_dequant(
|
||
input, self._scales, self._quant_axis, type='e4m3'
|
||
)
|
||
elif self._qmax == 57344:
|
||
return fake_fp8_dequant(
|
||
input, self._scales, self._quant_axis, type='e5m2'
|
||
)
|
||
elif len(self._scales.shape) > 1:
|
||
if self._zero_point.sum() != 0:
|
||
quant_dequant_weight = (
|
||
input.cast('float32') - self._zero_point
|
||
) * self._scales
|
||
else:
|
||
new_s = paddle.repeat_interleave(
|
||
self._scales, self._group_size, 0
|
||
)
|
||
new_zp = paddle.repeat_interleave(
|
||
self._zero_point, self._group_size, 0
|
||
)
|
||
quant_dequant_weight = (
|
||
(input.cast('float32') - new_zp) / self._qmax * new_s
|
||
)
|
||
return quant_dequant_weight.cast(input.dtype)
|
||
|
||
return _legacy_C_ops.dequantize_linear(
|
||
input.cast('float32'),
|
||
self._scales,
|
||
self._zero_point,
|
||
"quant_axis",
|
||
self._quant_axis,
|
||
"bit_length",
|
||
self._bit_length,
|
||
"qmin",
|
||
self._qmin,
|
||
"qmax",
|
||
self._qmax,
|
||
).cast(input.dtype)
|
||
if in_pir_mode():
|
||
input.stop_gradient = True
|
||
dequant_out = paddle.pir.core.create_persistable_value(
|
||
dtype='float32',
|
||
shape=input.shape,
|
||
name=unique_name.generate("quant_out"),
|
||
initializer=paddle.nn.initializer.Constant(0.0),
|
||
stop_gradient=True,
|
||
)
|
||
# TODO(xiaoluomi): need to add only observer pass for dequantize_linear
|
||
dequant_out, out_state, out_accum, out_scale = (
|
||
_C_ops.dequantize_linear(
|
||
input,
|
||
self._scales,
|
||
self._zero_point,
|
||
self.in_accum,
|
||
self.in_state,
|
||
self._quant_axis,
|
||
self._bit_length,
|
||
self._qmin,
|
||
self._qmax,
|
||
0,
|
||
True,
|
||
False,
|
||
)
|
||
)
|
||
return dequant_out
|
||
else:
|
||
out = self._helper.create_variable_for_type_inference(input.dtype)
|
||
self._helper.append_op(
|
||
type='dequantize_linear',
|
||
inputs={
|
||
'X': input,
|
||
'Scale': self._scales,
|
||
'ZeroPoint': self._zero_point,
|
||
},
|
||
outputs={'Y': out},
|
||
attrs={
|
||
'quant_axis': self._quant_axis,
|
||
'bit_length': self._bit_length,
|
||
'qmin': self._qmin,
|
||
'qmax': self._qmax,
|
||
},
|
||
)
|
||
return out
|
||
|
||
@staticmethod
|
||
def from_quanter(quanter):
|
||
return LinearDequanter(
|
||
quanter.scales(),
|
||
zero_point=quanter.zero_points(),
|
||
quant_axis=quanter.quant_axis(),
|
||
bit_length=quanter.bit_length(),
|
||
)
|
||
|
||
|
||
class ConvertibleQuantedLayer(Layer, metaclass=abc.ABCMeta):
|
||
r"""Abstract class to help convert quantized layer to inference model.
|
||
It defines some functions to convert quantizers and observers to quantize
|
||
or dequantize operators that maintain the quantization parameters used
|
||
during inference.
|
||
|
||
Examples:
|
||
.. code-block:: pycon
|
||
|
||
>>> # Given codes in ./customized_quanter.py
|
||
>>> class CustomizedQuantedLayer(ConvertibleQuantedLayer):
|
||
... def __init__(self):
|
||
... super().__init__()
|
||
... self.weight_a = paddle.create_parameter(shape=[1], dtype='float32')
|
||
... self.weight_b = paddle.create_parameter(shape=[1], dtype='float32')
|
||
... self.quanter_for_weight_a = None
|
||
... self.activation_weight = None
|
||
...
|
||
... def forward(self, input):
|
||
... qweight_a = self.quanter_for_weight_a(self.weight_a)
|
||
... weight_b = self.weight_b
|
||
... qinput = self.activation_weight(input)
|
||
... # compute with qweight_a, weight_b and qinput.
|
||
... return qweight * qinput + weight_b
|
||
...
|
||
... def weights_to_quanters(self):
|
||
... return [('weight_a', 'quanter_for_weight_a')]
|
||
...
|
||
... def activation_quanters(self):
|
||
... return ['activation_weight']
|
||
"""
|
||
|
||
def __init__(self):
|
||
super().__init__()
|
||
self.converted = False
|
||
|
||
@abc.abstractmethod
|
||
def weights_to_quanters(self) -> list[tuple[str, str]]:
|
||
r"""Get the name pairs of weights to be quantized and their corresponding
|
||
quantizers. In the convert function of this abstract class, it will call
|
||
the ‘weights_to_quanters’ function and do something as follows:
|
||
For each pair, the quantizer will be converted to a quantize operator and
|
||
a dequantize operator. Then, the weight will be quantized by the quantize
|
||
operator. Finally, the quantize operator will be removed and the weights
|
||
will be stored in integer data type.
|
||
|
||
Returns: A list of name pairs. Each pair contains two names. The first is name of weight
|
||
to be quantized and the second is name of corresponding quanter.
|
||
"""
|
||
pass
|
||
|
||
@abc.abstractmethod
|
||
def activation_quanters(self) -> list[str]:
|
||
r"""Get the names of quanters used to quantize activations.
|
||
All the quanters or observers returned by this function will be converted to quantize
|
||
and dequantize operators for deployment.
|
||
Returns: A list of quanter names.
|
||
"""
|
||
pass
|
||
|
||
def _convert_quanter_to_qdq(self, quanter_name) -> LinearQuanterDequanter:
|
||
r"""Convert quanter to an instance of LinearQuanterDequanter."""
|
||
if not hasattr(self, quanter_name):
|
||
return None
|
||
quanter = getattr(self, quanter_name)
|
||
if quanter is None:
|
||
return None
|
||
quanter = LinearQuanterDequanter.from_quanter(quanter)
|
||
setattr(self, quanter_name, quanter)
|
||
self._sub_layers[quanter_name] = quanter
|
||
return quanter
|
||
|
||
def _quant_weights(self, weight_name, quanter):
|
||
r"""Quantize the weight by given quanter."""
|
||
weight = getattr(self, weight_name)
|
||
qweight = quanter(weight)
|
||
weight.set_value(qweight)
|
||
|
||
def _convert(self, remain_weight=False):
|
||
r"""Convert current layer to onnx style for inference."""
|
||
assert not self.converted, "The model should be converted only once."
|
||
for weight_name, quanter_name in self.weights_to_quanters():
|
||
qdq = self._convert_quanter_to_qdq(quanter_name)
|
||
if qdq is not None and remain_weight is False:
|
||
self._quant_weights(weight_name, qdq._quanter)
|
||
qdq._quanter = None
|
||
qdq._sub_layers['_quanter'] = None
|
||
|
||
for quanter_name in self.activation_quanters():
|
||
self._convert_quanter_to_qdq(quanter_name)
|
||
|
||
self.converted = True
|