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