chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,44 @@
|
||||
# Copyright (c) 2021 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.
|
||||
|
||||
from . import qat # noqa: F401
|
||||
from .functional_layers import ( # noqa: F401
|
||||
FloatFunctionalLayer,
|
||||
add,
|
||||
concat,
|
||||
divide,
|
||||
flatten,
|
||||
matmul,
|
||||
multiply,
|
||||
reshape,
|
||||
subtract,
|
||||
transpose,
|
||||
)
|
||||
from .quant_layers import QuantStub # noqa: F401
|
||||
from .quantized_linear import ( # noqa: F401
|
||||
apply_per_channel_scale,
|
||||
llm_int8_linear,
|
||||
weight_dequantize,
|
||||
weight_only_linear,
|
||||
weight_quantize,
|
||||
)
|
||||
from .stub import Stub
|
||||
|
||||
__all__ = [
|
||||
"Stub",
|
||||
"weight_only_linear",
|
||||
"llm_int8_linear",
|
||||
"weight_quantize",
|
||||
"weight_dequantize",
|
||||
]
|
||||
@@ -0,0 +1,494 @@
|
||||
# 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
|
||||
@@ -0,0 +1,95 @@
|
||||
# Copyright (c) 2021 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.
|
||||
|
||||
from ...tensor import linalg, manipulation, math
|
||||
from ..layer.layers import Layer
|
||||
|
||||
__all__ = []
|
||||
|
||||
|
||||
class FloatFunctionalLayer(Layer):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
|
||||
|
||||
class add(FloatFunctionalLayer):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
|
||||
def forward(self, x, y, name=None):
|
||||
return math.add(x, y, name=name)
|
||||
|
||||
|
||||
class subtract(FloatFunctionalLayer):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
|
||||
def forward(self, x, y, name=None):
|
||||
return math.subtract(x, y, name=name)
|
||||
|
||||
|
||||
class multiply(FloatFunctionalLayer):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
|
||||
def forward(self, x, y, name=None):
|
||||
return math.multiply(x, y, name=name)
|
||||
|
||||
|
||||
class divide(FloatFunctionalLayer):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
|
||||
def forward(self, x, y, name=None):
|
||||
return math.divide(x, y, name=name)
|
||||
|
||||
|
||||
class reshape(FloatFunctionalLayer):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
|
||||
def forward(self, x, shape, name=None):
|
||||
return manipulation.reshape(x, shape, name=name)
|
||||
|
||||
|
||||
class transpose(FloatFunctionalLayer):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
|
||||
def forward(self, x, perm, name=None):
|
||||
return manipulation.transpose(x, perm, name=name)
|
||||
|
||||
|
||||
class concat(FloatFunctionalLayer):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
|
||||
def forward(self, x, axis=0, name=None):
|
||||
return manipulation.concat(x, axis, name=name)
|
||||
|
||||
|
||||
class flatten(FloatFunctionalLayer):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
|
||||
def forward(self, x, start_axis=0, stop_axis=-1, name=None):
|
||||
return manipulation.flatten(x, start_axis, stop_axis, name=name)
|
||||
|
||||
|
||||
class matmul(FloatFunctionalLayer):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
|
||||
def forward(self, x, y, transpose_x=False, transpose_y=False, name=None):
|
||||
return linalg.matmul(x, y, transpose_x, transpose_y, name=name)
|
||||
@@ -0,0 +1,370 @@
|
||||
# Copyright (c) 2022 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 math
|
||||
|
||||
import paddle
|
||||
from paddle.autograd import PyLayer
|
||||
from paddle.framework import ParamAttr
|
||||
from paddle.nn.initializer import Constant
|
||||
from paddle.utils import unique_name
|
||||
|
||||
from ..layer.layers import Layer
|
||||
|
||||
|
||||
def round(x):
|
||||
sign = paddle.sign(x)
|
||||
x = sign * paddle.floor(paddle.abs(x) + 0.5)
|
||||
return x
|
||||
|
||||
|
||||
class LsqFunc(PyLayer):
|
||||
@staticmethod
|
||||
def forward(ctx, weight, alpha, g, Qn, Qp, per_channel=False, quant_axis=0):
|
||||
ctx.save_for_backward(weight, alpha)
|
||||
ctx.other = g, Qn, Qp, per_channel, quant_axis
|
||||
if per_channel:
|
||||
sizes = weight.shape
|
||||
weight = weight.reshape((weight.shape[quant_axis], -1))
|
||||
weight = weight.transpose((1, 0))
|
||||
alpha = paddle.broadcast_to(alpha, weight.shape)
|
||||
quant_w = round(paddle.divide(weight, alpha)).clip(Qn, Qp)
|
||||
quant_w = quant_w * alpha
|
||||
quant_w = quant_w.transpose((1, 0))
|
||||
quant_w = quant_w.reshape(sizes)
|
||||
else:
|
||||
quant_w = round(paddle.divide(weight, alpha)).clip(Qn, Qp)
|
||||
quant_w = quant_w * alpha
|
||||
return quant_w
|
||||
|
||||
@staticmethod
|
||||
def backward(ctx, grad_weight):
|
||||
weight, alpha = ctx.saved_tensor()
|
||||
g, Qn, Qp, per_channel, quant_axis = ctx.other
|
||||
if per_channel:
|
||||
sizes = weight.shape
|
||||
weight = weight.reshape((weight.shape[quant_axis], -1))
|
||||
weight = weight.transpose((1, 0))
|
||||
alpha = paddle.broadcast_to(alpha, weight.shape)
|
||||
q_w = paddle.divide(weight, alpha)
|
||||
q_w = q_w.transpose((1, 0))
|
||||
q_w = q_w.reshape(sizes)
|
||||
else:
|
||||
q_w = paddle.divide(weight, alpha)
|
||||
lower_flag = paddle.cast((q_w < Qn), 'float32')
|
||||
upper_flag = paddle.cast((q_w > Qp), 'float32')
|
||||
middle_flag = 1.0 - lower_flag - upper_flag
|
||||
if per_channel:
|
||||
grad_alpha = (
|
||||
(
|
||||
lower_flag * Qn
|
||||
+ upper_flag * Qp
|
||||
+ middle_flag * round(q_w)
|
||||
- middle_flag * q_w
|
||||
)
|
||||
* grad_weight
|
||||
* g
|
||||
)
|
||||
grad_alpha = grad_alpha.reshape(
|
||||
(grad_alpha.shape[quant_axis], -1)
|
||||
).sum(axis=1)
|
||||
else:
|
||||
grad_alpha = (
|
||||
(
|
||||
(
|
||||
lower_flag * Qn
|
||||
+ upper_flag * Qp
|
||||
+ middle_flag * round(q_w)
|
||||
- middle_flag * q_w
|
||||
)
|
||||
* grad_weight
|
||||
* g
|
||||
)
|
||||
.sum()
|
||||
.unsqueeze(axis=0)[0]
|
||||
)
|
||||
grad_weight = middle_flag * grad_weight
|
||||
return grad_weight, grad_alpha
|
||||
|
||||
|
||||
class LsqPlusActFunc(PyLayer):
|
||||
@staticmethod
|
||||
def forward(ctx, x, alpha, beta, g, Qn, Qp):
|
||||
ctx.save_for_backward(x, alpha, beta)
|
||||
ctx.other = g, Qn, Qp
|
||||
quant_x = round(paddle.divide((x - beta), alpha)).clip(Qn, Qp)
|
||||
return quant_x * alpha + beta
|
||||
|
||||
@staticmethod
|
||||
def backward(ctx, grad_x):
|
||||
x, alpha, beta = ctx.saved_tensor()
|
||||
g, Qn, Qp = ctx.other
|
||||
q_x = (x - beta) / alpha
|
||||
lower_flag = paddle.cast((q_x < Qn), 'float32')
|
||||
upper_flag = paddle.cast((q_x > Qp), 'float32')
|
||||
middle_flag = 1.0 - lower_flag - upper_flag
|
||||
grad_alpha = (
|
||||
(
|
||||
(
|
||||
lower_flag * Qn
|
||||
+ upper_flag * Qp
|
||||
+ middle_flag * round(q_x)
|
||||
- middle_flag * q_x
|
||||
)
|
||||
* grad_x
|
||||
* g
|
||||
)
|
||||
.sum()
|
||||
.unsqueeze(axis=0)[0]
|
||||
)
|
||||
grad_beta = (
|
||||
((lower_flag + upper_flag) * grad_x * g).sum().unsqueeze(axis=0)[0]
|
||||
)
|
||||
grad_x = middle_flag * grad_x
|
||||
return grad_x, grad_alpha, grad_beta
|
||||
|
||||
|
||||
class FakeQuantActLSQPlus(Layer):
|
||||
def __init__(
|
||||
self,
|
||||
quant_bits,
|
||||
all_positive=False,
|
||||
symmetric=False,
|
||||
batch_init=20,
|
||||
dtype='float32',
|
||||
name=None,
|
||||
reduce_type=None,
|
||||
):
|
||||
super().__init__()
|
||||
'''
|
||||
Args:
|
||||
quant_bits(int): quantization bit number for weights.
|
||||
all_positive(bool): whether unsigned or signed quantization, where True for unsigned quantization and False for signed quantization.
|
||||
symmetric(bool): whether symmetric or asymmetric quantization.
|
||||
batch_init(int): number of batches that collect Gaussian approximation for the weight distribution in each layer.
|
||||
dtype(str): data type.
|
||||
name(str): the name of the weight.
|
||||
reduce_type(str): the reduce type which is needed when parallel training.
|
||||
'''
|
||||
self.bits = quant_bits
|
||||
self.all_positive = all_positive
|
||||
self.symmetric = symmetric
|
||||
self.batch_init = batch_init
|
||||
self.name = name
|
||||
self.reduce_type = reduce_type
|
||||
|
||||
if self.all_positive:
|
||||
# unsigned activation
|
||||
self.Qn = 0
|
||||
self.Qp = 2**self.bits - 1
|
||||
else:
|
||||
# signed activation
|
||||
self.Qn = -(2 ** (self.bits - 1))
|
||||
self.Qp = 2 ** (self.bits - 1) - 1
|
||||
|
||||
scale_prefix = f"{name}.scale" if name else 'quant_dequant.scale'
|
||||
self._scale_name = unique_name.generate(scale_prefix)
|
||||
|
||||
s_attr = ParamAttr(
|
||||
name=self._scale_name, initializer=Constant(1.0), trainable=True
|
||||
)
|
||||
self.s = self.create_parameter(shape=[], attr=s_attr, dtype='float32')
|
||||
self.s.stop_gradient = False
|
||||
|
||||
if not self.symmetric:
|
||||
beta_prefix = f"{name}.beta" if name else 'quant_dequant.beta'
|
||||
self._beta_name = unique_name.generate(beta_prefix)
|
||||
|
||||
beta_attr = ParamAttr(
|
||||
name=self._beta_name, initializer=Constant(0.0), trainable=True
|
||||
)
|
||||
self.beta = self.create_parameter(
|
||||
shape=[], attr=beta_attr, dtype='float32'
|
||||
)
|
||||
self.beta.stop_gradient = False
|
||||
|
||||
self.init_state = 0
|
||||
|
||||
def forward(self, activation):
|
||||
if self.reduce_type == "max":
|
||||
paddle.distributed.all_reduce(
|
||||
self.s, op=paddle.distributed.ReduceOp.MAX
|
||||
)
|
||||
|
||||
if not self.symmetric and self.reduce_type == "max":
|
||||
paddle.distributed.all_reduce(
|
||||
self.beta, op=paddle.distributed.ReduceOp.MAX
|
||||
)
|
||||
|
||||
if self.init_state == 0:
|
||||
self.g = paddle.to_tensor(
|
||||
1.0 / math.sqrt(activation.numel() * self.Qp)
|
||||
)
|
||||
min_a = paddle.min(activation.detach())
|
||||
max_a = paddle.max(activation.detach())
|
||||
self.s.set_value((max_a - min_a) / (self.Qp - self.Qn))
|
||||
if not self.symmetric:
|
||||
self.beta.set_value(min_a - self.s * self.Qn)
|
||||
self.init_state += 1
|
||||
elif self.init_state < self.batch_init:
|
||||
min_a = paddle.min(activation.detach())
|
||||
max_a = paddle.max(activation.detach())
|
||||
self.s.set_value(
|
||||
self.s * 0.9 + 0.1 * (max_a - min_a) / (self.Qp - self.Qn)
|
||||
)
|
||||
if not self.symmetric:
|
||||
self.beta.set_value(
|
||||
self.s * 0.9 + 0.1 * (min_a - self.s * self.Qn)
|
||||
)
|
||||
self.init_state += 1
|
||||
else:
|
||||
self.init_state += 1
|
||||
activation.stop_gradient = False
|
||||
if not self.symmetric:
|
||||
q_a = LsqPlusActFunc.apply(
|
||||
activation, self.s, self.beta, self.g, self.Qn, self.Qp
|
||||
)
|
||||
else:
|
||||
q_a = LsqFunc.apply(
|
||||
activation, self.s, self.g, self.Qn, self.Qp, per_channel=False
|
||||
)
|
||||
return q_a
|
||||
|
||||
|
||||
class FakeQuantWeightLSQPlus(Layer):
|
||||
def __init__(
|
||||
self,
|
||||
quant_bits,
|
||||
all_positive=False,
|
||||
per_channel=False,
|
||||
batch_init=20,
|
||||
channel_num=None,
|
||||
quant_linear=False,
|
||||
dtype='float32',
|
||||
name=None,
|
||||
reduce_type=None,
|
||||
):
|
||||
super().__init__()
|
||||
'''
|
||||
Args:
|
||||
quant_bits(int): quantization bit number for weights.
|
||||
all_positive(bool): whether unsigned or signed quantization, where True for unsigned quantization and False for signed quantization.
|
||||
per_channel(bool): whether layer-wise or channel-wise quantization, where True for layer-wise quantization and False for channel-wise quantization.
|
||||
batch_init(int): number of batches that collect Gaussian approximation for the weight distribution in each layer.
|
||||
channel_num(int): the channel number of the weight which is needed when per_channel is True.
|
||||
quant_linear(bool): whether the weight is from Linear.
|
||||
dtype(str): data type.
|
||||
name(str): the name of the weight.
|
||||
reduce_type(str): the reduce type which is needed when parallel training.
|
||||
'''
|
||||
|
||||
self.bits = quant_bits
|
||||
self.all_positive = all_positive
|
||||
self.per_channel = per_channel
|
||||
self.quant_linear = quant_linear
|
||||
self.batch_init = batch_init
|
||||
self.name = name
|
||||
self.quant_axis = 1 if quant_linear else 0
|
||||
self.collect_axis = 0 if quant_linear else 1
|
||||
self.reduce_type = reduce_type
|
||||
|
||||
if self.all_positive:
|
||||
# unsigned weight
|
||||
self.Qn = 0
|
||||
self.Qp = 2**self.bits - 1
|
||||
else:
|
||||
# signed weight
|
||||
self.Qn = -(2 ** (self.bits - 1))
|
||||
self.Qp = 2 ** (self.bits - 1) - 1
|
||||
|
||||
self.init_state = 0
|
||||
scale_prefix = f"{name}.scale" if name else 'quant_dequant.scale'
|
||||
self._scale_name = unique_name.generate(scale_prefix)
|
||||
s_attr = ParamAttr(
|
||||
name=self._scale_name, initializer=Constant(1.0), trainable=True
|
||||
)
|
||||
self.s = self.create_parameter(
|
||||
shape=[channel_num], attr=s_attr, dtype=dtype
|
||||
)
|
||||
self.s.stop_gradient = False
|
||||
|
||||
def forward(self, weight):
|
||||
if self.reduce_type == "max":
|
||||
paddle.distributed.all_reduce(
|
||||
self.s, op=paddle.distributed.ReduceOp.MAX
|
||||
)
|
||||
|
||||
if self.init_state == 0:
|
||||
self.g = paddle.to_tensor(1.0 / math.sqrt(weight.numel() * self.Qp))
|
||||
self.div = 2**self.bits - 1
|
||||
if self.per_channel:
|
||||
weight_tmp = weight.detach().reshape((weight.shape[0], -1))
|
||||
mean = paddle.mean(weight_tmp, axis=self.collect_axis)
|
||||
std = paddle.std(weight_tmp, axis=self.collect_axis)
|
||||
s = paddle.max(
|
||||
paddle.stack(
|
||||
[paddle.abs(mean - 3 * std), paddle.abs(mean + 3 * std)]
|
||||
),
|
||||
axis=0,
|
||||
)
|
||||
self.s.set_value(s / self.div)
|
||||
else:
|
||||
mean = paddle.mean(weight.detach())
|
||||
std = paddle.std(weight.detach())
|
||||
self.s.set_value(
|
||||
max(
|
||||
[paddle.abs(mean - 3 * std), paddle.abs(mean + 3 * std)]
|
||||
)
|
||||
/ self.div
|
||||
)
|
||||
self.init_state += 1
|
||||
elif self.init_state < self.batch_init:
|
||||
self.div = 2**self.bits - 1
|
||||
if self.per_channel:
|
||||
weight_tmp = weight.detach().reshape((weight.shape[0], -1))
|
||||
mean = paddle.mean(weight_tmp, axis=self.collect_axis)
|
||||
std = paddle.std(weight_tmp, axis=self.collect_axis)
|
||||
s = paddle.max(
|
||||
paddle.stack(
|
||||
[paddle.abs(mean - 3 * std), paddle.abs(mean + 3 * std)]
|
||||
),
|
||||
axis=0,
|
||||
)
|
||||
self.s.set_value(s * 0.9 + 0.1 * s / self.div)
|
||||
else:
|
||||
mean = paddle.mean(weight.detach())
|
||||
std = paddle.std(weight.detach())
|
||||
self.s.set_value(
|
||||
self.s * 0.9
|
||||
+ 0.1
|
||||
* max(
|
||||
[paddle.abs(mean - 3 * std), paddle.abs(mean + 3 * std)]
|
||||
)
|
||||
/ self.div
|
||||
)
|
||||
self.init_state += 1
|
||||
elif self.init_state == self.batch_init:
|
||||
self.init_state += 1
|
||||
|
||||
weight.stop_gradient = False
|
||||
w_q = LsqFunc.apply(
|
||||
weight,
|
||||
self.s,
|
||||
self.g,
|
||||
self.Qn,
|
||||
self.Qp,
|
||||
self.per_channel,
|
||||
self.quant_axis,
|
||||
)
|
||||
return w_q
|
||||
@@ -0,0 +1,15 @@
|
||||
# Copyright (c) 2022 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.
|
||||
from .conv import QuantedConv2D # noqa: F401
|
||||
from .linear import QuantedLinear # noqa: F401
|
||||
@@ -0,0 +1,88 @@
|
||||
# Copyright (c) 2022 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.
|
||||
"""
|
||||
Layers used for QAT.
|
||||
"""
|
||||
|
||||
from paddle.nn import functional as F
|
||||
|
||||
from ...layer.layers import Layer
|
||||
from ..format import ConvertibleQuantedLayer
|
||||
|
||||
|
||||
class QuantedConv2D(ConvertibleQuantedLayer):
|
||||
"""
|
||||
The computational logic of QuantizedConv2D is the same as Conv2D.
|
||||
The only difference is that its inputs are all fake quantized.
|
||||
"""
|
||||
|
||||
def __init__(self, layer: Layer, q_config):
|
||||
super().__init__()
|
||||
|
||||
# For Conv2D
|
||||
self._groups = layer._groups
|
||||
self._stride = layer._stride
|
||||
self._padding = layer._padding
|
||||
self._padding_mode = layer._padding_mode
|
||||
if self._padding_mode != 'zeros':
|
||||
self._reversed_padding_repeated_twice = (
|
||||
layer._reversed_padding_repeated_twice
|
||||
)
|
||||
self._dilation = layer._dilation
|
||||
self._data_format = layer._data_format
|
||||
self.weight = layer.weight
|
||||
self.bias = layer.bias
|
||||
|
||||
self.weight_quanter = None
|
||||
self.activation_quanter = None
|
||||
if q_config.weight is not None:
|
||||
self.weight_quanter = q_config.weight._instance(layer)
|
||||
if q_config.activation is not None:
|
||||
self.activation_quanter = q_config.activation._instance(layer)
|
||||
|
||||
def forward(self, input):
|
||||
quant_input = input
|
||||
quant_weight = self.weight
|
||||
if self.activation_quanter is not None:
|
||||
quant_input = self.activation_quanter(input)
|
||||
if self.weight_quanter is not None:
|
||||
quant_weight = self.weight_quanter(self.weight)
|
||||
return self._conv_forward(quant_input, quant_weight)
|
||||
|
||||
def _conv_forward(self, inputs, weights):
|
||||
if self._padding_mode != 'zeros':
|
||||
inputs = F.pad(
|
||||
inputs,
|
||||
self._reversed_padding_repeated_twice,
|
||||
mode=self._padding_mode,
|
||||
data_format=self._data_format,
|
||||
)
|
||||
self._padding = 0
|
||||
|
||||
return F.conv2d(
|
||||
inputs,
|
||||
weights,
|
||||
bias=self.bias,
|
||||
padding=self._padding,
|
||||
stride=self._stride,
|
||||
dilation=self._dilation,
|
||||
groups=self._groups,
|
||||
data_format=self._data_format,
|
||||
)
|
||||
|
||||
def weights_to_quanters(self):
|
||||
return [('weight', 'weight_quanter')]
|
||||
|
||||
def activation_quanters(self):
|
||||
return ['activation_quanter']
|
||||
@@ -0,0 +1,60 @@
|
||||
# Copyright (c) 2022 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.
|
||||
|
||||
|
||||
from paddle.nn import functional as F
|
||||
|
||||
from ...layer.layers import Layer
|
||||
from ..format import ConvertibleQuantedLayer
|
||||
|
||||
|
||||
class QuantedLinear(ConvertibleQuantedLayer):
|
||||
"""
|
||||
The computational logic of QuantizedLinear is the same as Linear.
|
||||
The only difference is that its inputs are all fake quantized.
|
||||
"""
|
||||
|
||||
def __init__(self, layer: Layer, q_config):
|
||||
super().__init__()
|
||||
# For Linear
|
||||
self.weight = layer.weight
|
||||
self.bias = layer.bias
|
||||
self.name = layer.name
|
||||
# For FakeQuant
|
||||
|
||||
self.weight_quanter = None
|
||||
self.activation_quanter = None
|
||||
if q_config.weight is not None:
|
||||
self.weight_quanter = q_config.weight._instance(layer)
|
||||
if q_config.activation is not None:
|
||||
self.activation_quanter = q_config.activation._instance(layer)
|
||||
|
||||
def forward(self, input):
|
||||
quant_input = input
|
||||
quant_weight = self.weight
|
||||
if self.activation_quanter is not None:
|
||||
quant_input = self.activation_quanter(input)
|
||||
if self.weight_quanter is not None:
|
||||
quant_weight = self.weight_quanter(self.weight)
|
||||
return self._linear_forward(quant_input, quant_weight)
|
||||
|
||||
def _linear_forward(self, input, weight):
|
||||
out = F.linear(x=input, weight=weight, bias=self.bias, name=self.name)
|
||||
return out
|
||||
|
||||
def weights_to_quanters(self):
|
||||
return [('weight', 'weight_quanter')]
|
||||
|
||||
def activation_quanters(self):
|
||||
return ['activation_quanter']
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,395 @@
|
||||
# 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.
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import TYPE_CHECKING, Literal
|
||||
|
||||
import paddle
|
||||
from paddle import _C_ops
|
||||
from paddle.base.data_feeder import check_dtype
|
||||
from paddle.device import (
|
||||
is_compiled_with_cuda,
|
||||
)
|
||||
from paddle.device.cuda import get_device_capability
|
||||
from paddle.framework import (
|
||||
LayerHelper,
|
||||
in_dynamic_or_pir_mode,
|
||||
)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from typing import TypeAlias
|
||||
|
||||
from paddle import Tensor
|
||||
from paddle._typing import DTypeLike
|
||||
|
||||
_Algo: TypeAlias = Literal[
|
||||
'weight_only_int8', 'weight_only_int4', 'llm.int8'
|
||||
]
|
||||
_GroupSize: TypeAlias = Literal[-1, 64, 128]
|
||||
|
||||
|
||||
def _get_arch_info():
|
||||
# Get SMVersion from device.
|
||||
if is_compiled_with_cuda():
|
||||
cuda_version = paddle.version.cuda()
|
||||
if (
|
||||
cuda_version is not None and cuda_version != 'False'
|
||||
) or paddle.is_compiled_with_rocm():
|
||||
major, minor = get_device_capability()
|
||||
arch = int(major * 10 + minor)
|
||||
return arch
|
||||
else:
|
||||
raise ValueError(
|
||||
"Paddle is not compiled with CUDA, we cannot get SMVersion from device, please try to compile Paddle with CUDA"
|
||||
)
|
||||
else:
|
||||
# Default arch value for type checking.
|
||||
return 0
|
||||
|
||||
|
||||
def weight_quantize(
|
||||
x: Tensor,
|
||||
algo: _Algo = "weight_only_int8",
|
||||
arch: int | None = None,
|
||||
group_size: _GroupSize = -1,
|
||||
) -> tuple[Tensor, Tensor]:
|
||||
"""
|
||||
Quantization function for weight_only and llm.int8's weight.
|
||||
|
||||
Args:
|
||||
x (Tensor): The input Tensor to be quantized, the data type is float16 or bfloat16.
|
||||
algo (str): The algo that is x will be apply, must be one of 'weight_only_int8',
|
||||
'weight_only_int4', 'llm.int8', 'w4a8' and 'w4afp8, default: 'weight_only_int8'.
|
||||
arch (int): The compute arch for target device. For example, A100 is 80, v100 is 70, if you do not assign arch, we will get arch from your device, default: None.
|
||||
group_size (int): The group size for weight quantization. -1 stands for default per-channel mode. Currently only support 64 or 128.
|
||||
|
||||
Returns:
|
||||
out (Tensor): The Tensor which is the quantitative results, the data type is int8, the shape is transposition of x.
|
||||
scale (Tensor): The scale Tensor which is the scale of pre-channel, the data type is float32.
|
||||
Examples:
|
||||
.. code-block:: pycon
|
||||
|
||||
>>> # doctest: +SKIP('No testing required')
|
||||
>>> import paddle
|
||||
>>> from paddle.nn.quant import weight_quantize
|
||||
|
||||
>>> paddle.seed(2023)
|
||||
>>> x = paddle.rand(shape=[64, 32], dtype=paddle.float16)
|
||||
>>> out, scale = weight_quantize(x, algo='weight_only_int8')
|
||||
>>> print(out.shape)
|
||||
paddle.Size([32, 64])
|
||||
>>> print(scale.shape)
|
||||
paddle.Size([32])
|
||||
"""
|
||||
if arch is None:
|
||||
arch = _get_arch_info()
|
||||
|
||||
if is_compiled_with_cuda():
|
||||
assert (
|
||||
arch == 70
|
||||
or arch == 75
|
||||
or arch == 80
|
||||
or arch == 86
|
||||
or arch == 89
|
||||
or arch == 90
|
||||
or arch == 92
|
||||
or arch == 100
|
||||
), (
|
||||
f"Currently weight_quantize only support SM70/75/80/86/89/90/92/100. but got {arch} "
|
||||
)
|
||||
|
||||
assert group_size == -1 or group_size == 64 or group_size == 128, (
|
||||
f"Currently group_size only support -1/64/128. but got {group_size} "
|
||||
)
|
||||
if in_dynamic_or_pir_mode():
|
||||
return _C_ops.weight_quantize(x, algo, arch, group_size)
|
||||
else:
|
||||
type = "weight_quantize"
|
||||
helper = LayerHelper(type, **locals())
|
||||
out = helper.create_variable_for_type_inference('int8')
|
||||
scale = helper.create_variable_for_type_inference('float')
|
||||
|
||||
helper.append_op(
|
||||
type=type,
|
||||
inputs={"x": x},
|
||||
outputs={'out': out, "scale": scale},
|
||||
attrs={"algo": algo, "arch": arch, "group_size": group_size},
|
||||
)
|
||||
return (out, scale)
|
||||
|
||||
|
||||
def weight_dequantize(
|
||||
x: Tensor,
|
||||
scale: Tensor,
|
||||
algo: _Algo = "weight_only_int8",
|
||||
out_dtype: DTypeLike = "float16",
|
||||
group_size: _GroupSize = -1,
|
||||
) -> Tensor:
|
||||
"""
|
||||
Dequantization function for weight_only and llm.int8's weight.
|
||||
|
||||
Args:
|
||||
x (Tensor): The input Tensor to be dequantized, the data type is int8.
|
||||
scale (Tensor): The scale Tensor which is the output of weight_quantize, the data type is float32.
|
||||
algo (str): The algo that is x will be apply, must be one of 'weight_only_int8',
|
||||
'weight_only_int4' and 'llm.int8', default: 'weight_only_int8'.
|
||||
out_dtype (str|np.dtype): [Deprecated][Not used] The output Tensor's data type, must be one of 'float16' and 'bfloat16', default: 'float16'.
|
||||
|
||||
Returns:
|
||||
out (Tensor): The Tensor which is the dequantitative results, the data type is float16 or bfloat16, the shape is transposition of x.
|
||||
|
||||
Examples:
|
||||
.. code-block:: pycon
|
||||
|
||||
>>> # doctest: +SKIP('No testing required')
|
||||
>>> import paddle
|
||||
>>> from paddle.nn.quant import weight_quantize, weight_dequantize
|
||||
|
||||
>>> paddle.seed(2023)
|
||||
>>> x = paddle.rand(shape=[64, 32], dtype=paddle.float16)
|
||||
>>> out, scale = weight_quantize(x, algo='weight_only_int8')
|
||||
>>> x_dequant = weight_dequantize(out, scale)
|
||||
"""
|
||||
assert group_size == -1 or group_size == 64 or group_size == 128, (
|
||||
f"Currently group_size only support -1/64/128. but got {group_size} "
|
||||
)
|
||||
|
||||
if in_dynamic_or_pir_mode():
|
||||
return _C_ops.weight_dequantize(x, scale, algo, group_size)
|
||||
else:
|
||||
type = "weight_dequantize"
|
||||
helper = LayerHelper(type, **locals())
|
||||
out_dtype = scale.dtype
|
||||
out = helper.create_variable_for_type_inference(out_dtype)
|
||||
|
||||
helper.append_op(
|
||||
type=type,
|
||||
inputs={"x": x, "scale": scale},
|
||||
outputs={'out': out},
|
||||
attrs={
|
||||
"algo": algo,
|
||||
"group_size": group_size,
|
||||
},
|
||||
)
|
||||
return out
|
||||
|
||||
|
||||
def weight_only_linear(
|
||||
x: Tensor,
|
||||
weight: Tensor,
|
||||
bias: Tensor | None = None,
|
||||
weight_scale: Tensor | None = None,
|
||||
weight_dtype: DTypeLike = "int8",
|
||||
arch: int | None = None,
|
||||
group_size: _GroupSize = -1,
|
||||
) -> Tensor:
|
||||
"""
|
||||
Applies matrix multiplication of two tensors and then bias addition if provided.
|
||||
This method requires CUDA version >= 11.2.
|
||||
|
||||
Args:
|
||||
x (Tensor): The first input Tensor to be multiplied, the data type is float16 or bfloat16.
|
||||
weight (Tensor): The second input Tensor to be multiplied. Its rank must be 2.
|
||||
bias (Tensor|None): The input bias Tensor. If it is None, no bias addition would
|
||||
be performed. Otherwise, The bias is added to the matrix multiplication result.
|
||||
weight_scale (Tensor|None): The input scale Tensor Provided to weight for dequantization. Its rank must be 1.
|
||||
weight_dtype(str): The dtype of weight Tensor, must be one of 'int8', 'int4', Defaulted to 'int8'.
|
||||
arch (int): The compute arch for target device. For example, A100 is 80, v100 is 70, if you do not assign arch, we will get arch from your device, default: None.
|
||||
group_size (int): The group size for weight quantization. -1 stands for default per-channel mode. Currently only support 64 or 128.
|
||||
Returns:
|
||||
Tensor: the output Tensor, the data type is the same as that of x.
|
||||
|
||||
Examples:
|
||||
.. code-block:: pycon
|
||||
|
||||
>>> # doctest: +SKIP('No testing required')
|
||||
>>> import paddle
|
||||
>>> from paddle.nn.quant import weight_only_linear
|
||||
|
||||
>>> x = paddle.cast(paddle.randn([1, 2, 64]), dtype='float16')
|
||||
>>> weight = paddle.cast(paddle.randint(0, 127, [32, 64]), dtype='int8')
|
||||
>>> scale = paddle.randn([32], dtype='float32')
|
||||
>>> bias = paddle.cast(paddle.randn([32]), dtype='float16')
|
||||
>>> if paddle.device.cuda.get_device_capability()[0] >= 8:
|
||||
... out = weight_only_linear(
|
||||
... x,
|
||||
... weight,
|
||||
... bias=bias,
|
||||
... weight_scale=scale,
|
||||
... weight_dtype='int8',
|
||||
... )
|
||||
... print(out.shape)
|
||||
paddle.Size([1, 2, 32])
|
||||
"""
|
||||
if arch is None:
|
||||
arch = _get_arch_info()
|
||||
|
||||
if is_compiled_with_cuda():
|
||||
assert (
|
||||
arch == 70
|
||||
or arch == 75
|
||||
or arch == 80
|
||||
or arch == 86
|
||||
or arch == 89
|
||||
or arch == 90
|
||||
or arch == 92
|
||||
or arch == 100
|
||||
), (
|
||||
f"Currently weight_quantize only support SM70/75/80/86/89/90/92/100. but got {arch} "
|
||||
)
|
||||
assert group_size == -1 or group_size == 64 or group_size == 128, (
|
||||
f"Currently weight_quantize only support group size of -1, 64 or 128. but got {group_size} "
|
||||
)
|
||||
|
||||
if in_dynamic_or_pir_mode():
|
||||
out = _C_ops.weight_only_linear(
|
||||
x, weight, bias, weight_scale, weight_dtype, arch, group_size
|
||||
)
|
||||
return out
|
||||
else:
|
||||
check_dtype(
|
||||
weight_dtype, 'weight_dtype', ['int8', 'int4'], 'weight_only_linear'
|
||||
)
|
||||
type = "weight_only_linear"
|
||||
helper = LayerHelper(type, **locals())
|
||||
dtype = x.dtype
|
||||
|
||||
inputs = {
|
||||
'x': [x],
|
||||
'weight': [weight],
|
||||
'weight_scale': [weight_scale],
|
||||
}
|
||||
if bias is not None:
|
||||
inputs["bias"] = [bias]
|
||||
attrs = {
|
||||
'weight_dtype': weight_dtype,
|
||||
'arch': arch,
|
||||
'group_size': group_size,
|
||||
}
|
||||
|
||||
out = helper.create_variable_for_type_inference(dtype)
|
||||
|
||||
helper.append_op(
|
||||
type=type,
|
||||
inputs=inputs,
|
||||
outputs={'out': out},
|
||||
attrs=attrs,
|
||||
)
|
||||
return out
|
||||
|
||||
|
||||
def llm_int8_linear(
|
||||
x: Tensor,
|
||||
weight: Tensor,
|
||||
bias: Tensor | None = None,
|
||||
weight_scale: Tensor | None = None,
|
||||
threshold: float = 6.0,
|
||||
) -> Tensor:
|
||||
"""
|
||||
Applies matrix multiplication of two tensors and then bias addition if provided.
|
||||
This method requires CUDA version >= 11.2.
|
||||
|
||||
Args:
|
||||
x (Tensor): the first input Tensor to be multiplied, the data type is float16 or bfloat16.
|
||||
weight (Tensor): the second input Tensor to be multiplied. Its rank must be 2.
|
||||
bias (Tensor|None): the input bias Tensor. If it is None, no bias addition would
|
||||
be performed. Otherwise, the bias is added to the matrix multiplication result.
|
||||
weight_scale (Tensor|None): the input scale Tensor Provided to weight for dequantization. Its rank must be 1.
|
||||
threshold(float): The min value of outlier in activation, outlier's channel will be apply multiply with x.dtype.
|
||||
|
||||
Returns:
|
||||
Tensor: the output Tensor, the data type is the same as that of x.
|
||||
|
||||
Examples:
|
||||
.. code-block:: pycon
|
||||
|
||||
>>> # doctest: +SKIP('No testing required')
|
||||
>>> import paddle
|
||||
>>> from paddle.nn.quant import llm_int8_linear
|
||||
|
||||
>>> x = paddle.cast(paddle.randn([1, 2, 64]), dtype='float16')
|
||||
>>> weight = paddle.cast(paddle.randint(0, 127, [32, 64]), dtype='int8')
|
||||
>>> scale = paddle.randn([32], dtype='float32')
|
||||
>>> bias = paddle.cast(paddle.randn([32]), dtype='float16')
|
||||
>>> if paddle.device.cuda.get_device_capability()[0] >= 8:
|
||||
... out = llm_int8_linear(x, weight, bias=bias, weight_scale=scale, threshold=6.0)
|
||||
... print(out.shape)
|
||||
paddle.Size([1, 2, 32])
|
||||
"""
|
||||
if in_dynamic_or_pir_mode():
|
||||
out = _C_ops.llm_int8_linear(x, weight, bias, weight_scale, threshold)
|
||||
return out
|
||||
else:
|
||||
type = "llm_int8_linear"
|
||||
helper = LayerHelper(type, **locals())
|
||||
dtype = x.dtype
|
||||
|
||||
inputs = {
|
||||
'x': [x],
|
||||
'weight': [weight],
|
||||
'weight_scale': [weight_scale],
|
||||
}
|
||||
if bias:
|
||||
inputs["bias"] = [bias]
|
||||
attrs = {'threshold': threshold}
|
||||
|
||||
out = helper.create_variable_for_type_inference(dtype)
|
||||
|
||||
helper.append_op(
|
||||
type=type,
|
||||
inputs=inputs,
|
||||
outputs={'out': out},
|
||||
attrs=attrs,
|
||||
)
|
||||
return out
|
||||
|
||||
|
||||
def apply_per_channel_scale(x: Tensor, scales: Tensor) -> Tensor:
|
||||
"""
|
||||
Apply pre-quant per channel scale on activations
|
||||
|
||||
Args:
|
||||
x (Tensor): Input tensor representing the activations, the data type can be float16 or bfloat16.
|
||||
scales(Tensor): Per-channel scale factors for pre-quantization. Data type should be compatible with x.
|
||||
|
||||
Returns:
|
||||
out (Tensor): The Tensor which is the pre-quant results, the data type is compatible with x.
|
||||
|
||||
Examples:
|
||||
.. code-block:: pycon
|
||||
|
||||
>>> # doctest: +SKIP('No testing required')
|
||||
>>> import paddle
|
||||
>>> from paddle.nn.quant import apply_per_channel_scale
|
||||
|
||||
>>> paddle.seed(2023)
|
||||
>>> x = paddle.rand(shape=[64, 32], dtype=paddle.float16)
|
||||
>>> scales = paddle.rand(shape=[32], dtype=paddle.float16)
|
||||
>>> out = apply_per_channel_scale(x, scales)
|
||||
"""
|
||||
|
||||
if in_dynamic_or_pir_mode():
|
||||
return _C_ops.apply_per_channel_scale(x, scales)
|
||||
else:
|
||||
type = "apply_per_channel_scale"
|
||||
helper = LayerHelper(type, **locals())
|
||||
out = helper.create_variable_for_type_inference(x.dtype)
|
||||
|
||||
helper.append_op(
|
||||
type=type,
|
||||
inputs={"x": [x], "scales": [scales]},
|
||||
outputs={"out": out},
|
||||
)
|
||||
return out
|
||||
@@ -0,0 +1,107 @@
|
||||
# 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 stub used in quantization."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from ..layer.layers import Layer
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from paddle import Tensor
|
||||
from paddle.quantization import QuantConfig
|
||||
from paddle.quantization.factory import QuanterFactory
|
||||
|
||||
|
||||
class Stub(Layer):
|
||||
r"""
|
||||
The stub is used as placeholders that will be replaced by observers before PTQ or QAT.
|
||||
It is hard to assign a quantization configuration to a functional API called in
|
||||
the forward of a layer. Instead, we can create a stub and add it to the sublayers of the layer.
|
||||
And call the stub before the functional API in the forward. The observer held by the
|
||||
stub will observe or quantize the inputs of the functional API.
|
||||
|
||||
Args:
|
||||
observer(QuanterFactory): The configured information of the observer to be inserted.
|
||||
It will use a global configuration to create the observers if the 'observer' is none.
|
||||
|
||||
Examples:
|
||||
.. code-block:: pycon
|
||||
|
||||
>>> import paddle
|
||||
>>> from paddle.nn.quant import Stub
|
||||
>>> from paddle.quantization.quanters import FakeQuanterWithAbsMaxObserver
|
||||
>>> from paddle.nn import Conv2D
|
||||
>>> from paddle.quantization import QAT, QuantConfig
|
||||
|
||||
>>> quanter = FakeQuanterWithAbsMaxObserver(moving_rate=0.9)
|
||||
>>> class Model(paddle.nn.Layer):
|
||||
... def __init__(self, num_classes=10):
|
||||
... super().__init__()
|
||||
... self.conv = Conv2D(3, 6, 3, stride=1, padding=1)
|
||||
... self.quant = Stub(quanter)
|
||||
...
|
||||
... def forward(self, inputs):
|
||||
... out = self.conv(inputs)
|
||||
... out = self.quant(out)
|
||||
... return paddle.nn.functional.relu(out)
|
||||
|
||||
>>> model = Model()
|
||||
>>> q_config = QuantConfig(activation=quanter, weight=quanter)
|
||||
>>> qat = QAT(q_config)
|
||||
>>> quant_model = qat.quantize(model)
|
||||
>>> print(quant_model)
|
||||
Model(
|
||||
(conv): QuantedConv2D(
|
||||
(weight_quanter): FakeQuanterWithAbsMaxObserverLayer()
|
||||
(activation_quanter): FakeQuanterWithAbsMaxObserverLayer()
|
||||
)
|
||||
(quant): QuanterStub(
|
||||
(_observer): FakeQuanterWithAbsMaxObserverLayer()
|
||||
)
|
||||
)
|
||||
"""
|
||||
|
||||
def __init__(self, observer: QuanterFactory | None = None) -> None:
|
||||
super().__init__()
|
||||
self._observer = observer
|
||||
|
||||
def forward(self, input: Tensor) -> Tensor:
|
||||
return input
|
||||
|
||||
|
||||
class QuanterStub(Layer):
|
||||
r"""
|
||||
It is an identity layer with an observer observing the input.
|
||||
Before QAT or PTQ, the stub in the model will be replaced with an instance of QuanterStub.
|
||||
The user should not use this class directly.
|
||||
|
||||
Args:
|
||||
layer(paddle.nn.Layer): The stub layer with an observer configure factory. If the observer
|
||||
of the stub layer is none, it will use 'q_config' to create an observer instance.
|
||||
q_config(QuantConfig): The quantization configuration for the current stub layer.
|
||||
"""
|
||||
|
||||
def __init__(self, layer: Stub, q_config: QuantConfig) -> None:
|
||||
super().__init__()
|
||||
self._observer = None
|
||||
if layer._observer is not None:
|
||||
self._observer = layer._observer._instance(layer)
|
||||
elif q_config.activation is not None:
|
||||
self._observer = q_config.activation._instance(layer)
|
||||
|
||||
def forward(self, input: Tensor) -> Tensor:
|
||||
return self._observer(input) if self._observer is not None else input
|
||||
Reference in New Issue
Block a user