wehub-resource-sync
222 lines
8.6 KiB
Python
222 lines
8.6 KiB
Python
# Copyright (c) 2024 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 numpy as np
|
|
import paddle
|
|
from paddle.distributed import fleet
|
|
|
|
from paddlenlp.quantization.checkpoint_quantization_utils import (
|
|
asymmetry_qdq_weight,
|
|
cal_ratio,
|
|
group_wise_quant_dequant,
|
|
merge_int4,
|
|
qdq_weight,
|
|
split_int8,
|
|
)
|
|
from paddlenlp.utils.env import (
|
|
ASYMMETRY_QUANT_SCALE_MAX,
|
|
ASYMMETRY_QUANT_SCALE_MIN,
|
|
MOMENT1_KEYNAME,
|
|
MOMENT2_KEYNAME,
|
|
SYMMETRY_QUANT_SCALE,
|
|
)
|
|
from paddlenlp.utils.log import logger
|
|
|
|
|
|
def dequant_unified_optimizer(state_dict, ckpt_quant_stage, scale_dict, use_pd=False):
|
|
"""
|
|
dequantize unified optimizer state dict.
|
|
Args:
|
|
state_dict (`dict`):
|
|
unified checkpoint optimizer state dict.
|
|
ckpt_quant_stage (`str`):
|
|
checkpoint quantization stage, chosen in ["O0", "O1", "O2"].
|
|
scale_dict (`int`):
|
|
compression checkpoint scale dict.
|
|
"""
|
|
logger.info(f"Start unified checkpoint dequantization, stage {ckpt_quant_stage}.")
|
|
tp_rank, tp_degree = -1, 1
|
|
if paddle.distributed.get_world_size() > 1:
|
|
hcg = fleet.get_hybrid_communicate_group()
|
|
tp_group = hcg.get_model_parallel_group()
|
|
tp_rank, tp_degree = tp_group.rank, tp_group.nranks
|
|
|
|
if ckpt_quant_stage == "O1":
|
|
# set eps
|
|
eps = 1e-8
|
|
for quant_key in state_dict.keys():
|
|
is_moment1 = MOMENT1_KEYNAME in quant_key
|
|
is_moment2 = MOMENT2_KEYNAME in quant_key
|
|
if is_moment1:
|
|
# dequant m1
|
|
scale_key = quant_key + SYMMETRY_QUANT_SCALE
|
|
weight = state_dict[quant_key]
|
|
scales = scale_dict[scale_key]
|
|
weight, _ = qdq_weight(
|
|
weight,
|
|
scales=scales,
|
|
quant_bit=8,
|
|
dequant=True,
|
|
tp_rank=tp_rank,
|
|
tp_degree=tp_degree,
|
|
use_pd=use_pd,
|
|
)
|
|
state_dict[quant_key] = weight
|
|
elif is_moment2:
|
|
# dequant ratio
|
|
weight = state_dict[quant_key]
|
|
min_scale_key = quant_key + ASYMMETRY_QUANT_SCALE_MIN
|
|
max_scale_key = quant_key + ASYMMETRY_QUANT_SCALE_MAX
|
|
mins, maxs = scale_dict[min_scale_key], scale_dict[max_scale_key]
|
|
weight, _ = asymmetry_qdq_weight(
|
|
weight,
|
|
mins=mins,
|
|
maxs=maxs,
|
|
quant_bit=8,
|
|
dequant=True,
|
|
tp_rank=tp_rank,
|
|
tp_degree=tp_degree,
|
|
use_pd=use_pd,
|
|
)
|
|
# cal m2
|
|
if use_pd:
|
|
weight = paddle.square(1.0 / weight - eps)
|
|
else:
|
|
weight = np.square(1.0 / weight - eps)
|
|
state_dict[quant_key] = weight
|
|
elif ckpt_quant_stage == "O2":
|
|
# set eps
|
|
eps = 1e-8
|
|
m1_state_dict = {}
|
|
for quant_key in state_dict.keys():
|
|
# not all optimizer weights in O2 stage were quantized to int8,
|
|
# the norm-like weights were still remain in float32.
|
|
if state_dict[quant_key].dtype != paddle.int8:
|
|
logger.info(f"{quant_key} skip.")
|
|
continue
|
|
# split int8
|
|
weight = state_dict[quant_key]
|
|
m1_quant, ratio_quant = split_int8(weight.numpy())
|
|
# dequant ratio
|
|
ratio_min_scale_key = quant_key + ASYMMETRY_QUANT_SCALE_MIN
|
|
ratio_max_scale_key = quant_key + ASYMMETRY_QUANT_SCALE_MAX
|
|
m1_scale_key = quant_key[: -len(MOMENT2_KEYNAME)] + MOMENT1_KEYNAME + SYMMETRY_QUANT_SCALE
|
|
m1_scales = scale_dict[m1_scale_key]
|
|
ratio_mins, ratio_maxs = scale_dict[ratio_min_scale_key], scale_dict[ratio_max_scale_key]
|
|
m1_weight = group_wise_quant_dequant(
|
|
m1_quant,
|
|
mins=m1_scales,
|
|
maxs=None,
|
|
quant_bits=4,
|
|
quant=False,
|
|
tp_rank=tp_rank,
|
|
tp_degree=tp_degree,
|
|
use_pd=use_pd,
|
|
symmetry=True,
|
|
)
|
|
ratio_weight = group_wise_quant_dequant(
|
|
ratio_quant,
|
|
mins=ratio_mins,
|
|
maxs=ratio_maxs,
|
|
quant_bits=4,
|
|
quant=False,
|
|
tp_rank=tp_rank,
|
|
tp_degree=tp_degree,
|
|
use_pd=use_pd,
|
|
)
|
|
|
|
if use_pd:
|
|
ratio_weight = paddle.square(1.0 / ratio_weight - eps)
|
|
else:
|
|
ratio_weight = np.square(1.0 / ratio_weight - eps)
|
|
state_dict[quant_key] = ratio_weight
|
|
m1_state_dict[quant_key[: -len(MOMENT2_KEYNAME)] + MOMENT1_KEYNAME] = m1_weight
|
|
state_dict.update(m1_state_dict)
|
|
|
|
logger.info(f"Unified checkpoint dequantization done, stage {ckpt_quant_stage}.")
|
|
|
|
return state_dict
|
|
|
|
|
|
def quant_unified_optimizer(state_dict, state_dict_type, ckpt_quant_stage, async_save=False):
|
|
"""
|
|
quantize unified optimizer state dict.
|
|
Args:
|
|
state_dict (`dict`):
|
|
unified checkpoint optimizer state dict.
|
|
state_dict_type (`str`):
|
|
state_dict type, chosen in ["model_weight", "master_weight", "optimizer_weight"].
|
|
ckpt_quant_stage (`str`):
|
|
checkpoint quantization stage, chosen in ["O0", "O1", "O2"].
|
|
async_save (`bool`):
|
|
whether use async_save.
|
|
"""
|
|
logger.info(f"Start unified checkpoint quantization, stage {ckpt_quant_stage}.")
|
|
|
|
quant = False
|
|
if ckpt_quant_stage != "O0":
|
|
quant = True
|
|
del_key = []
|
|
if quant and state_dict_type == "optimizer_weight":
|
|
scales_dict = {}
|
|
for k in state_dict.keys():
|
|
momentum1 = k.endswith(MOMENT1_KEYNAME)
|
|
momentum2 = k.endswith(MOMENT2_KEYNAME)
|
|
|
|
quant_weight = None
|
|
|
|
if ckpt_quant_stage == "O1":
|
|
# m1: wint8, 1/(sqrt(m2)+eps): wint8
|
|
if momentum2:
|
|
# m1: m1_quant_weight, m2: ratio
|
|
m1_key = k.split("/")[0] + "/" + MOMENT1_KEYNAME
|
|
ratio = cal_ratio(state_dict[m1_key], state_dict[k])
|
|
m1_quant, scales = qdq_weight(state_dict[m1_key], quant_bit=8)
|
|
quant_weight, mins, maxs = asymmetry_qdq_weight(ratio, quant_bit=8)
|
|
state_dict[m1_key] = m1_quant
|
|
scales_dict[m1_key + SYMMETRY_QUANT_SCALE] = scales
|
|
scales_dict[k + ASYMMETRY_QUANT_SCALE_MIN] = mins
|
|
scales_dict[k + ASYMMETRY_QUANT_SCALE_MAX] = maxs
|
|
elif not momentum1:
|
|
quant_weight = state_dict[k]
|
|
elif ckpt_quant_stage == "O2":
|
|
# m1: bw-wint4, 1/(sqrt(m2)+eps): bw-wint4
|
|
if momentum2:
|
|
# skip norm-like parameters
|
|
if len(state_dict[k].shape) < 2:
|
|
continue
|
|
# m1: m1_quant_weight, m2: ratio
|
|
m1_key = k.split("/")[0] + "/" + MOMENT1_KEYNAME
|
|
ratio = cal_ratio(state_dict[m1_key], state_dict[k])
|
|
m1_quant, m1_scales = group_wise_quant_dequant(state_dict[m1_key], quant_bits=4, symmetry=True)
|
|
quant_weight, r_mins, r_maxs = group_wise_quant_dequant(ratio, quant_bits=4)
|
|
quant_weight = merge_int4(m1_quant, quant_weight)
|
|
scales_dict[m1_key + SYMMETRY_QUANT_SCALE] = m1_scales
|
|
scales_dict[k + ASYMMETRY_QUANT_SCALE_MIN] = r_mins
|
|
scales_dict[k + ASYMMETRY_QUANT_SCALE_MAX] = r_maxs
|
|
del_key.append(m1_key)
|
|
elif not momentum1:
|
|
quant_weight = state_dict[k]
|
|
|
|
if quant_weight is not None:
|
|
state_dict[k] = quant_weight
|
|
|
|
for k in del_key:
|
|
state_dict.pop(k, None)
|
|
|
|
state_dict.update(scales_dict)
|
|
logger.info(f"Unified checkpoint quantization done, stage {ckpt_quant_stage}.")
|
|
|
|
return state_dict
|