# 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