wehub-resource-sync
279 lines
12 KiB
Python
279 lines
12 KiB
Python
# Copyright (c) 2025 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 re
|
|
|
|
import paddle
|
|
import paddle.nn as nn
|
|
from paddle.distributed.fleet.meta_parallel import (
|
|
ColumnParallelLinear,
|
|
RowParallelLinear,
|
|
)
|
|
from paddle.distributed.fleet.utils.sequence_parallel_utils import (
|
|
ColumnSequenceParallelLinear,
|
|
RowSequenceParallelLinear,
|
|
)
|
|
from paddle.incubate.nn.layer.fused_linear import FusedLinear
|
|
from paddle.nn.quant import weight_quantize
|
|
|
|
try:
|
|
from .qlora import qlora_weight_linear, qlora_weight_quantize
|
|
except:
|
|
qlora_weight_linear = None
|
|
qlora_weight_quantize = None
|
|
|
|
from ..utils.log import logger
|
|
from .qat_utils import quantize
|
|
from .quantization_linear import (
|
|
ColumnParallelQuantizationLinear,
|
|
QuantizationLinear,
|
|
RowParallelQuantizationLinear,
|
|
)
|
|
|
|
LINEAR_CLASSES = [
|
|
nn.Linear,
|
|
FusedLinear,
|
|
ColumnParallelLinear,
|
|
RowParallelLinear,
|
|
ColumnSequenceParallelLinear,
|
|
RowSequenceParallelLinear,
|
|
]
|
|
|
|
|
|
def parse_weight_quantize_algo(quantization_config, name):
|
|
if quantization_config.ignore_modules is not None and any(
|
|
re.fullmatch(ignore_module, name) for ignore_module in quantization_config.ignore_modules
|
|
):
|
|
weight_quantize_algo = None
|
|
elif isinstance(quantization_config.weight_quantize_algo, str):
|
|
weight_quantize_algo = quantization_config.weight_quantize_algo
|
|
else:
|
|
weight_quantize_algo = None
|
|
for algo in quantization_config.weight_quantize_algo:
|
|
if any(re.fullmatch(module, name) for module in quantization_config.weight_quantize_algo[algo]):
|
|
weight_quantize_algo = algo
|
|
return weight_quantize_algo
|
|
|
|
|
|
def replace_with_quantization_linear(model, quantization_config, llm_int8_threshold=6.0):
|
|
for name, child in model.named_sublayers():
|
|
weight_quantize_algo = parse_weight_quantize_algo(quantization_config, name)
|
|
if weight_quantize_algo is None:
|
|
continue
|
|
if any(isinstance(child, linear_class) for linear_class in LINEAR_CLASSES):
|
|
if child.bias is None:
|
|
bias_attr = False
|
|
else:
|
|
bias_attr = None
|
|
parent = model
|
|
*path, last = name.split(".")
|
|
for attr in path:
|
|
parent = getattr(parent, attr)
|
|
if isinstance(child, nn.Linear) or isinstance(child, FusedLinear):
|
|
if getattr(child.weight, "transpose_weight", False):
|
|
out_feature, in_features = child.weight.shape[0], child.weight.shape[1]
|
|
else:
|
|
in_features, out_feature = child.weight.shape[0], child.weight.shape[1]
|
|
quant_linear = QuantizationLinear(
|
|
in_features=in_features,
|
|
out_features=out_feature,
|
|
quantization_config=quantization_config,
|
|
weight_quantize_algo=weight_quantize_algo,
|
|
dtype=child._dtype,
|
|
bias_attr=bias_attr,
|
|
mp_moe=getattr(child.weight, "mp_moe", False),
|
|
is_distributed=getattr(child.weight, "is_distributed", False),
|
|
)
|
|
elif isinstance(child, ColumnParallelLinear):
|
|
quant_linear = ColumnParallelQuantizationLinear(
|
|
in_features=child.weight.shape[0],
|
|
output_size_per_partition=child.weight.shape[1],
|
|
quantization_config=quantization_config,
|
|
weight_quantize_algo=weight_quantize_algo,
|
|
dtype=child._dtype,
|
|
bias_attr=bias_attr,
|
|
gather_output=child.gather_output,
|
|
mp_skip_c_identity=child.mp_skip_c_identity,
|
|
)
|
|
elif isinstance(child, RowParallelLinear):
|
|
quant_linear = RowParallelQuantizationLinear(
|
|
input_size_per_partition=child.weight.shape[0],
|
|
out_features=child.weight.shape[1],
|
|
quantization_config=quantization_config,
|
|
weight_quantize_algo=weight_quantize_algo,
|
|
dtype=child._dtype,
|
|
bias_attr=bias_attr,
|
|
input_is_parallel=child.input_is_parallel,
|
|
mp_skip_c_identity=child.mp_skip_c_identity,
|
|
)
|
|
elif isinstance(child, ColumnSequenceParallelLinear):
|
|
quant_linear = ColumnParallelQuantizationLinear(
|
|
in_features=child.weight.shape[0],
|
|
output_size_per_partition=child.weight.shape[1],
|
|
quantization_config=quantization_config,
|
|
weight_quantize_algo=weight_quantize_algo,
|
|
dtype=child._dtype,
|
|
bias_attr=bias_attr,
|
|
gather_output=False,
|
|
sequence_parallel=True,
|
|
)
|
|
elif isinstance(child, RowSequenceParallelLinear):
|
|
quant_linear = RowParallelQuantizationLinear(
|
|
input_size_per_partition=child.weight.shape[0],
|
|
out_features=child.weight.shape[1],
|
|
quantization_config=quantization_config,
|
|
weight_quantize_algo=weight_quantize_algo,
|
|
dtype=child._dtype,
|
|
bias_attr=bias_attr,
|
|
input_is_parallel=True,
|
|
sequence_parallel=True,
|
|
)
|
|
setattr(parent, last, quant_linear)
|
|
del child
|
|
|
|
|
|
def convert_to_weight_quantize_state_dict(state_dict, name, quantization_config, dtype, weight_quantize_algo):
|
|
|
|
weight_name = name + ".weight"
|
|
quant_weight_name = name + ".quant_weight"
|
|
quant_scale_name = name + ".quant_scale"
|
|
act_scale_name = name + ".act_scale"
|
|
|
|
if quant_weight_name in state_dict and quant_scale_name in state_dict:
|
|
return state_dict
|
|
if weight_name in state_dict:
|
|
# gpu weight_quantize will fix in future
|
|
target_weight = state_dict.pop(weight_name).cast(dtype).cuda()
|
|
|
|
if weight_quantize_algo in ["a8w8linear", "a8w4linear", "fp8linear"]:
|
|
quant_weight, quant_scale = quantize(
|
|
target_weight,
|
|
weight_quantize_algo,
|
|
"weight",
|
|
quantization_config,
|
|
side="left",
|
|
apply_hadamard=quantization_config.apply_hadamard,
|
|
)
|
|
act_scale = paddle.ones([1], dtype=dtype).cuda()
|
|
act_scale.stop_gradient = True
|
|
state_dict[act_scale_name] = act_scale
|
|
else:
|
|
quant_weight, quant_scale = weight_quantize(
|
|
x=target_weight,
|
|
algo=weight_quantize_algo,
|
|
group_size=quantization_config.group_size,
|
|
)
|
|
state_dict[quant_weight_name] = quant_weight
|
|
state_dict[quant_scale_name] = quant_scale
|
|
del target_weight
|
|
return state_dict
|
|
|
|
|
|
def convert_to_qlora_state_dict(state_dict, name, quantization_config, dtype, weight_quantize_algo):
|
|
if qlora_weight_quantize is None:
|
|
raise ImportError(
|
|
"Please run the following commands to install qlora related package first: \n"
|
|
"1) git clone https://github.com/PaddlePaddle/PaddleSlim \n"
|
|
"2) cd PaddleSlim \n"
|
|
"3) python ./csrc/setup_cuda.py install"
|
|
)
|
|
weight_name = name + ".weight"
|
|
quant_weight_name = name + ".quant_weight"
|
|
quant_name_list = [quant_weight_name]
|
|
if not quantization_config.qlora_weight_double_quant:
|
|
quant_scale_name = name + ".quant_scale"
|
|
quant_name_list += [quant_scale_name]
|
|
else:
|
|
qquant_scale_name = name + ".qquant_scale"
|
|
double_quant_scale_name = name + ".double_quant_scale"
|
|
quant_sacle_offset_name = name + ".quant_sacle_offset"
|
|
quant_name_list += [qquant_scale_name, double_quant_scale_name, quant_sacle_offset_name]
|
|
|
|
if all(quant_name in state_dict for quant_name in quant_name_list):
|
|
return state_dict
|
|
elif weight_name in state_dict:
|
|
target_weight = state_dict.pop(weight_name).cast(dtype).cuda()
|
|
qlora_state_dict = qlora_weight_quantize(
|
|
weight=target_weight,
|
|
quant_algo=weight_quantize_algo,
|
|
double_quant=quantization_config.qlora_weight_double_quant,
|
|
block_size=quantization_config.qlora_weight_blocksize,
|
|
double_quant_block_size=quantization_config.qlora_weight_double_quant_block_size,
|
|
linear_name=name,
|
|
return_dict=True,
|
|
)
|
|
state_dict.update(qlora_state_dict)
|
|
del target_weight
|
|
|
|
return state_dict
|
|
|
|
|
|
def convert_to_quantize_state_dict(state_dict, quantization_linear_list, quantization_config, dtype):
|
|
for name in quantization_linear_list:
|
|
# Get quantization algorithm
|
|
weight_quantize_algo = parse_weight_quantize_algo(quantization_config, name)
|
|
if weight_quantize_algo is None:
|
|
continue
|
|
# Convert state dict
|
|
if weight_quantize_algo in [
|
|
"weight_only_int8",
|
|
"weight_only_int4",
|
|
"llm.int8",
|
|
"a8w8linear",
|
|
"a8w4linear",
|
|
"fp8linear",
|
|
]:
|
|
convert_to_weight_quantize_state_dict(state_dict, name, quantization_config, dtype, weight_quantize_algo)
|
|
elif weight_quantize_algo in ["fp4", "nf4"]:
|
|
convert_to_qlora_state_dict(state_dict, name, quantization_config, dtype, weight_quantize_algo)
|
|
else:
|
|
raise NotImplementedError(
|
|
f"Please check the quantization_config.weight_quantize_algo: {quantization_config.weight_quantize_algo}"
|
|
)
|
|
return state_dict
|
|
|
|
|
|
def update_loaded_state_dict_keys(state_dict, quantization_linear_list, quantization_config, ignore_warning=False):
|
|
for name in quantization_linear_list:
|
|
weight_name = name + ".weight"
|
|
quant_weight_name = name + ".quant_weight"
|
|
quant_scale_name = name + ".quant_scale"
|
|
act_scale_name = name + ".act_scale"
|
|
qquant_scale_name = name + ".qquant_scale"
|
|
double_quant_scale_name = name + ".double_quant_scale"
|
|
quant_sacle_offset_name = name + ".quant_sacle_offset"
|
|
|
|
if quant_weight_name in state_dict and quant_scale_name in state_dict:
|
|
continue
|
|
elif weight_name in state_dict:
|
|
state_dict.remove(weight_name)
|
|
state_dict.append(quant_weight_name)
|
|
if quantization_config.qlora_weight_double_quant:
|
|
state_dict.append(qquant_scale_name)
|
|
state_dict.append(double_quant_scale_name)
|
|
state_dict.append(quant_sacle_offset_name)
|
|
else:
|
|
state_dict.append(quant_scale_name)
|
|
weight_quantize_algo = parse_weight_quantize_algo(quantization_config, name)
|
|
if weight_quantize_algo in ["a8w8linear", "a8w4linear", "fp8linear"]:
|
|
state_dict.append(act_scale_name)
|
|
|
|
else:
|
|
if not ignore_warning:
|
|
logger.warning(
|
|
f"Cannot find {weight_name} in state_dict or {quant_weight_name} and {quant_scale_name} in state_dict"
|
|
)
|
|
|
|
return state_dict
|