chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -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
|
||||
Reference in New Issue
Block a user