paddlepaddle--paddle
wehub-resource-sync
495 行
18 KiB
Python
495 行
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
|