wehub-resource-sync
116 lines
4.1 KiB
Python
116 lines
4.1 KiB
Python
# Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
|
|
import paddle
|
|
from paddleslim.lc.quantizers.quant_func import dequantize_8bit, quantize_8bit
|
|
from paddleslim_ops import dequant_blockwise, quant_blockwise
|
|
|
|
|
|
def qlora_weight_quantize(
|
|
weight,
|
|
quant_algo="nf4",
|
|
double_quant=False,
|
|
block_size=64,
|
|
double_quant_block_size=256,
|
|
linear_name=None,
|
|
return_dict=True,
|
|
):
|
|
quant_weight, quant_scale = quant_blockwise(weight, None, blocksize=block_size, quant_type=quant_algo)
|
|
if double_quant:
|
|
quant_sacle_offset = quant_scale.mean()
|
|
quant_scale -= quant_sacle_offset
|
|
qquant_scale, double_quant_scale = quantize_8bit(
|
|
quant_scale, None, double_quant_block_size, quant_type="dynamic_fp8"
|
|
)
|
|
if not return_dict:
|
|
return quant_weight, (qquant_scale, double_quant_scale, quant_sacle_offset)
|
|
qquant_scale_name = f"{linear_name}.qquant_scale" if linear_name else "qquant_scale"
|
|
double_quant_scale_name = f"{linear_name}.double_quant_scale" if linear_name else "double_quant_scale"
|
|
quant_sacle_offset_name = f"{linear_name}.quant_sacle_offset" if linear_name else "quant_sacle_offset"
|
|
qlora_state_dict = {
|
|
qquant_scale_name: qquant_scale,
|
|
double_quant_scale_name: double_quant_scale,
|
|
quant_sacle_offset_name: quant_sacle_offset,
|
|
}
|
|
else:
|
|
quant_scale_name = f"{linear_name}.quant_scale" if linear_name else "quant_scale"
|
|
qlora_state_dict = {quant_scale_name: quant_scale}
|
|
if not return_dict:
|
|
return quant_weight, (quant_scale)
|
|
quant_weight_name = f"{linear_name}.quant_weight" if linear_name else "quant_weight"
|
|
qlora_state_dict[quant_weight_name] = quant_weight
|
|
return qlora_state_dict
|
|
|
|
|
|
def qlora_weight_dequantize(
|
|
quant_weight, quant_algo, state, double_quant=False, block_size=64, double_quant_block_size=256
|
|
):
|
|
if double_quant:
|
|
qquant_scale, double_quant_scale, quant_sacle_offset = state
|
|
quant_scale = dequantize_8bit(
|
|
qquant_scale, None, double_quant_scale, double_quant_block_size, quant_type="dynamic_fp8"
|
|
)
|
|
quant_scale += quant_sacle_offset
|
|
else:
|
|
quant_scale = state
|
|
out = dequant_blockwise(quant_weight, None, quant_scale, blocksize=block_size, quant_type=quant_algo)
|
|
return out
|
|
|
|
|
|
def qlora_weight_quantize_dequantize(
|
|
weight, quant_algo="nf4", double_quant=False, block_size=64, double_quant_block_size=256
|
|
):
|
|
dtype = weight.dtype
|
|
quant_weight, state = qlora_weight_quantize(
|
|
weight=weight,
|
|
quant_algo=quant_algo,
|
|
double_quant=double_quant,
|
|
block_size=block_size,
|
|
double_quant_block_size=double_quant_block_size,
|
|
return_dict=False,
|
|
)
|
|
quant_dequant_weight = (
|
|
qlora_weight_dequantize(
|
|
quant_weight=quant_weight,
|
|
quant_algo=quant_algo,
|
|
state=state,
|
|
double_quant=double_quant,
|
|
block_size=block_size,
|
|
double_quant_block_size=double_quant_block_size,
|
|
)
|
|
.reshape(weight.shape)
|
|
.cast(dtype)
|
|
)
|
|
return quant_dequant_weight
|
|
|
|
|
|
def qlora_weight_linear(
|
|
x,
|
|
quant_weight,
|
|
dtype,
|
|
state,
|
|
quant_algo="nf4",
|
|
double_quant=False,
|
|
block_size=64,
|
|
double_quant_block_size=256,
|
|
bias=None,
|
|
):
|
|
weight = (
|
|
qlora_weight_dequantize(quant_weight, quant_algo, state, double_quant, block_size, double_quant_block_size)
|
|
.cast(dtype)
|
|
.reshape([x.shape[-1], -1])
|
|
)
|
|
out = paddle.nn.functional.linear(x, weight, bias)
|
|
return out
|