wehub-resource-sync
567 lines
21 KiB
Python
567 lines
21 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 json
|
|
import os
|
|
|
|
import paddle
|
|
from paddle import nn
|
|
from paddle.distributed.fleet.meta_parallel import (
|
|
ColumnParallelLinear,
|
|
RowParallelLinear,
|
|
)
|
|
from paddle.quantization import PTQ, QAT, QuantConfig
|
|
from paddle.quantization.base_observer import BaseObserver
|
|
from paddleslim.common.wrapper_function import FuncWrapper
|
|
from paddleslim.quant.advanced import (
|
|
GPTQ,
|
|
AutoClip,
|
|
AWQSearch,
|
|
EMASampler,
|
|
MultiStepSampler,
|
|
PieceWiseSearch,
|
|
Shift,
|
|
Smooth,
|
|
)
|
|
from paddleslim.quant.advanced.utils import find_parent_layer_and_sub_name
|
|
from paddleslim.quant.layers import (
|
|
QuantizedColumnParallelLinear,
|
|
QuantizedRowParallelLinear,
|
|
)
|
|
from paddleslim.quant.layers.custom_attention import QuantizedCustomAttentionLayer
|
|
from paddleslim.quant.observers import (
|
|
AbsMaxChannelWiseWeightObserver,
|
|
GroupWiseWeightObserver,
|
|
)
|
|
from paddleslim.quant.observers.abs_max import AbsmaxObserver
|
|
from paddleslim.quant.observers.abs_max_headwise import AbsMaxHeadwiseObserver
|
|
from paddleslim.quant.observers.avg import AVGObserver
|
|
from paddleslim.quant.observers.avg_headwise import AvgHeadwiseObserver
|
|
from paddleslim.quant.observers.channel_wise import ChannelWiseObserver
|
|
|
|
from paddlenlp.peft import PrefixModelForCausalLM
|
|
from paddlenlp.peft.lora import (
|
|
ColumnParallelLoRALinear,
|
|
LoRALinear,
|
|
RowParallelLoRALinear,
|
|
)
|
|
from paddlenlp.peft.lora.lora_quant_layers import (
|
|
ColumnParallelQuantedLoRALinear,
|
|
QuantedLoRALinear,
|
|
RowParallelQuantedLoRALinear,
|
|
)
|
|
from paddlenlp.utils.log import logger
|
|
|
|
ACT_OBSERVER = dict(
|
|
abs_max=AbsmaxObserver,
|
|
avg=AVGObserver,
|
|
)
|
|
|
|
WEIGHT_OBSERVER = dict(
|
|
abs_max_channel_wise=AbsMaxChannelWiseWeightObserver,
|
|
groupwise=GroupWiseWeightObserver,
|
|
)
|
|
|
|
CACHEKV_OBSERVER = dict(
|
|
abs_max_headwise=AbsMaxHeadwiseObserver,
|
|
avg_headwise=AvgHeadwiseObserver,
|
|
)
|
|
|
|
FP8_OBSERVER = dict(
|
|
abs_max=AbsmaxObserver,
|
|
avg=AVGObserver,
|
|
)
|
|
|
|
|
|
def create_qat_model(quant_args, model, dtype):
|
|
from paddle.quantization.quanters import FakeQuanterWithAbsMaxObserver
|
|
from paddleslim.quant.quanters import (
|
|
FakeQuanterChannelWiseAbsMaxObserver,
|
|
PACTQuanter,
|
|
)
|
|
|
|
q_config = QuantConfig(activation=None, weight=None)
|
|
q_config.add_qat_layer_mapping(LoRALinear, QuantedLoRALinear)
|
|
q_config.add_qat_layer_mapping(RowParallelLoRALinear, RowParallelQuantedLoRALinear)
|
|
q_config.add_qat_layer_mapping(ColumnParallelLoRALinear, ColumnParallelQuantedLoRALinear)
|
|
if quant_args.quant_type == "a8w8":
|
|
activation = PACTQuanter(quanter=FakeQuanterWithAbsMaxObserver(), init_value=20.0, dtype=dtype)
|
|
weight = FakeQuanterChannelWiseAbsMaxObserver(bit_length=8, dtype="float32")
|
|
elif quant_args.quant_type == "weight_only_int4":
|
|
activation = None
|
|
weight = FakeQuanterChannelWiseAbsMaxObserver(bit_length=4, dtype="float32")
|
|
elif quant_args.quant_type == "weight_only_int8":
|
|
activation = None
|
|
weight = FakeQuanterChannelWiseAbsMaxObserver(bit_length=8, dtype="float32")
|
|
else:
|
|
raise ValueError("quant_type should be one of ['a8w8', 'weight_only_int4', 'weight_only_int8']")
|
|
|
|
q_config.add_type_config(RowParallelLoRALinear, weight=weight, activation=activation)
|
|
q_config.add_type_config(ColumnParallelLoRALinear, weight=weight, activation=activation)
|
|
q_config.add_type_config(LoRALinear, weight=weight, activation=activation)
|
|
q_config.add_type_config(nn.Linear, weight=weight, activation=activation)
|
|
|
|
qat = QAT(q_config)
|
|
model = qat.quantize(model, inplace=True)
|
|
return model
|
|
|
|
|
|
def apply_shift(quant_args, trainer, ptq_dataloader, ptq_model_config):
|
|
logger.info("***** Running Shift *****")
|
|
shift_sampler = EMASampler() if quant_args.shift_sampler == "ema" else None
|
|
shift = Shift(
|
|
model=trainer.model,
|
|
model_config=ptq_model_config,
|
|
sample_function=shift_sampler,
|
|
shift_all_linears=quant_args.shift_all_linears,
|
|
)
|
|
with paddle.no_grad():
|
|
trainer.ptq_loop(
|
|
ptq_dataloader,
|
|
description="Shift",
|
|
max_eval_iters=quant_args.shift_step,
|
|
)
|
|
shift.update_weight()
|
|
del shift, shift_sampler
|
|
logger.info("***** Shift done *****")
|
|
|
|
|
|
def apply_smooth(quant_args, trainer, ptq_dataloader, ptq_model_config):
|
|
|
|
if quant_args.do_awq:
|
|
logger.info("***** Running AWQ *****")
|
|
else:
|
|
logger.info("***** Running Smooth *****")
|
|
smooth_sampler = MultiStepSampler() if quant_args.smooth_sampler == "multi_step" else None
|
|
if quant_args.smooth_piecewise_search:
|
|
search_func = PieceWiseSearch(
|
|
k_piece=quant_args.smooth_k_piece,
|
|
bits_length=8,
|
|
search_piece=quant_args.smooth_search_piece,
|
|
search_alpha_min=quant_args.search_alpha_min,
|
|
search_alpha_max=quant_args.search_alpha_max,
|
|
search_scale_min=quant_args.search_scale_min,
|
|
search_scale_max=quant_args.search_scale_max,
|
|
weight_quant_method=quant_args.weight_quant_method,
|
|
act_quant_method=quant_args.act_quant_method,
|
|
)
|
|
elif quant_args.do_awq:
|
|
search_func = AWQSearch(
|
|
n_grid=20,
|
|
bits_length=4,
|
|
weight_quant_method=quant_args.weight_quant_method,
|
|
)
|
|
else:
|
|
search_func = None
|
|
smooth = Smooth(
|
|
trainer.model,
|
|
ptq_model_config,
|
|
alpha=0.5,
|
|
smooth_all_linears=quant_args.smooth_all_linears,
|
|
sample_function=smooth_sampler,
|
|
search_function=search_func,
|
|
smooth_method="awq" if quant_args.do_awq else "smoothquant",
|
|
)
|
|
with paddle.no_grad():
|
|
trainer.ptq_loop(
|
|
ptq_dataloader,
|
|
description="Smooth",
|
|
max_eval_iters=quant_args.smooth_step,
|
|
)
|
|
|
|
smooth.update_weight()
|
|
del smooth, smooth_sampler, search_func
|
|
logger.info("***** Smooth done *****")
|
|
|
|
|
|
def apply_autoclip(quant_args, trainer, ptq_dataloader):
|
|
"""
|
|
AutoClip
|
|
"""
|
|
print("-------------------Start AutoClip------------------")
|
|
sampler = MultiStepSampler()
|
|
auto_clip = AutoClip(
|
|
trainer.model,
|
|
weight_bits=4,
|
|
weight_quant_method=quant_args.weight_quant_method,
|
|
sample_function=sampler,
|
|
n_grid=20,
|
|
max_shrink=0.5,
|
|
)
|
|
with paddle.no_grad():
|
|
trainer.ptq_loop(
|
|
ptq_dataloader,
|
|
description="AutoClip",
|
|
max_eval_iters=quant_args.autoclip_step,
|
|
)
|
|
auto_clip.auto_clip()
|
|
del sampler, auto_clip
|
|
logger.info("***** AutoClip done *****")
|
|
|
|
|
|
def prepare_qconfig(args):
|
|
"""
|
|
Prepare qconfig
|
|
"""
|
|
args.quant_type = args.quant_type.lower()
|
|
|
|
if args.quant_type in ["a8w8_fp8"]:
|
|
use_fp8 = "aw"
|
|
args.quant_type = args.quant_type.replace("_fp8", "")
|
|
else:
|
|
use_fp8 = ""
|
|
|
|
weight_observer = (
|
|
WEIGHT_OBSERVER.get(args.weight_quant_method, None)
|
|
if "w" not in use_fp8
|
|
else FP8_OBSERVER.get(args.weight_quant_method, None)
|
|
)
|
|
act_observer = (
|
|
ACT_OBSERVER.get(args.act_quant_method, None)
|
|
if "a" not in use_fp8
|
|
else FP8_OBSERVER.get(args.act_quant_method, None)
|
|
)
|
|
cachekv_observer = CACHEKV_OBSERVER.get(args.cachekv_quant_method, None)
|
|
|
|
if "c8" in args.quant_type:
|
|
quant_type = args.quant_type.replace("c8", "")
|
|
cachekv_quant = True
|
|
cachekv_quant_bits = "int8"
|
|
else:
|
|
quant_type = args.quant_type.replace("c16", "")
|
|
cachekv_quant = False
|
|
|
|
q_config = QuantConfig(activation=None, weight=None)
|
|
if quant_type in ["a8w8", "w8a8"]:
|
|
|
|
if "w" in use_fp8:
|
|
w_quant_bit = (4, 3) if args.fp8_type[use_fp8.index("w")] == "e4m3" else (5, 2)
|
|
else:
|
|
w_quant_bit = 8
|
|
|
|
if "a" in use_fp8:
|
|
a_quant_bit = (4, 3) if args.fp8_type[use_fp8.index("a")] == "e4m3" else (5, 2)
|
|
else:
|
|
a_quant_bit = 8
|
|
activation = act_observer(quant_bits=a_quant_bit)
|
|
weight = weight_observer(quant_bits=w_quant_bit)
|
|
|
|
elif quant_type in ["wint4", "w4a16", "weight_only_int4"]:
|
|
activation = None
|
|
weight = weight_observer(quant_bits=4)
|
|
|
|
elif quant_type in ["wint8", "w8a16", "weight_only_int8"]:
|
|
activation = None
|
|
if "w" in use_fp8:
|
|
weight = weight_observer(quant_bits=(4, 3))
|
|
else:
|
|
weight = weight_observer(quant_bits=8)
|
|
else:
|
|
raise ValueError(
|
|
"quant_type should be in ['weight_only_int8/wint8', 'weight_only_int4/wint4', 'a8w8', 'a8w8c8', 'a8w8_fp8']"
|
|
)
|
|
|
|
q_config.add_qat_layer_mapping(ColumnParallelLinear, QuantizedColumnParallelLinear)
|
|
q_config.add_qat_layer_mapping(RowParallelLinear, QuantizedRowParallelLinear)
|
|
|
|
cachekv = None
|
|
if cachekv_quant:
|
|
if cachekv_quant_bits == "int8":
|
|
cachekv_quant_bit = 8
|
|
if "headwise" in args.cachekv_quant_method:
|
|
cachekv = [
|
|
cachekv_observer(quant_bits=cachekv_quant_bit, quant_axis=1),
|
|
cachekv_observer(quant_bits=cachekv_quant_bit, quant_axis=1),
|
|
]
|
|
else:
|
|
cachekv = [
|
|
cachekv_observer(quant_bits=cachekv_quant_bit),
|
|
cachekv_observer(quant_bits=cachekv_quant_bit),
|
|
]
|
|
q_config.add_qat_layer_mapping(FuncWrapper, QuantizedCustomAttentionLayer)
|
|
else:
|
|
raise ValueError("cachekv_quant_bits should be int8")
|
|
|
|
return activation, weight, cachekv, q_config
|
|
|
|
|
|
def load_quant_model(model, quant_args, load_quant_path, dtype="float32"):
|
|
"""
|
|
Load quantized model and its scales
|
|
"""
|
|
activation, weight, cachekv, q_config = prepare_qconfig(quant_args)
|
|
|
|
if cachekv is not None:
|
|
set_wrapper_for_attn(model)
|
|
|
|
skip_list_names = [] if quant_args.skip_list_names is None else quant_args.skip_list_names
|
|
for cur_name, cur_layer in model.named_sublayers():
|
|
skip = False
|
|
for k in skip_list_names:
|
|
if k in cur_name:
|
|
logger.info(f"Skip layer {cur_name}")
|
|
skip = True
|
|
if skip:
|
|
continue
|
|
|
|
if type(cur_layer) in [paddle.nn.Linear, ColumnParallelLinear, RowParallelLinear]:
|
|
logger.info(f"PTQ layer: {cur_name}")
|
|
q_config.add_name_config([cur_layer.full_name()], activation=activation, weight=weight)
|
|
|
|
if type(cur_layer) in [FuncWrapper] and cachekv is not None:
|
|
logger.info(f"PTQ layer: {cur_name}")
|
|
# set both act and weight for attention, actually act-k and act-v are quantized
|
|
q_config.add_name_config([cur_layer.full_name()], weight=cachekv[0], activation=cachekv[1])
|
|
|
|
ptq = PTQ(q_config)
|
|
model = ptq.quantize(model, inplace=True)
|
|
|
|
logger.info("Load quant model...")
|
|
if activation is not None:
|
|
with open(f"{load_quant_path}/act_scales.json") as outfile:
|
|
act_scales = json.load(outfile)
|
|
else:
|
|
act_scales = {}
|
|
|
|
if cachekv is not None:
|
|
with open(f"{load_quant_path}/cachekv_scales.json") as outfile:
|
|
cachekv_scales = json.load(outfile)
|
|
else:
|
|
cachekv_scales = {}
|
|
|
|
with open(f"{load_quant_path}/weight_scales.json") as outfile:
|
|
weight_scales = json.load(outfile)
|
|
|
|
for cur_name, cur_layer in model.named_sublayers():
|
|
if hasattr(cur_layer, "scales"):
|
|
|
|
if isinstance(cur_layer, ChannelWiseObserver) or isinstance(cur_layer, BaseObserver):
|
|
logger.info(f"Load scale for layer {cur_name}")
|
|
if "attn_func" in cur_name:
|
|
cur_name = cur_name.replace("attn_func.activation_quanter_v", "cachev_matmul.activation_quanter")
|
|
cur_name = cur_name.replace("attn_func.activation_quanter_k", "cachek_matmul.activation_quanter")
|
|
if cur_name in cachekv_scales:
|
|
cur_layer._scale = paddle.to_tensor(cachekv_scales[cur_name], dtype=dtype)
|
|
if cur_name + ".zero_point" in cachekv_scales:
|
|
cur_layer._zero_point = paddle.to_tensor(
|
|
cachekv_scales[cur_name + ".zero_point"], dtype=dtype
|
|
)
|
|
else:
|
|
cur_layer._zero_point = paddle.to_tensor(0.0, dtype=dtype)
|
|
else:
|
|
logger.info(f"No scale found for layer {cur_name}, remove it")
|
|
parent_layer, sub_name = find_parent_layer_and_sub_name(model, cur_name)
|
|
setattr(parent_layer, sub_name, None)
|
|
|
|
elif "activation_quanter" in cur_name:
|
|
if cur_name in act_scales:
|
|
cur_layer._scale = paddle.to_tensor(act_scales[cur_name], dtype=dtype)
|
|
cur_layer._zero_point = paddle.to_tensor(0.0, dtype=dtype)
|
|
else:
|
|
logger.info(f"No scale found for layer {cur_name}, remove it")
|
|
parent_layer, sub_name = find_parent_layer_and_sub_name(model, cur_name)
|
|
setattr(parent_layer, sub_name, None)
|
|
elif "weight_quanter" in cur_name:
|
|
if cur_name in weight_scales:
|
|
cur_layer._scale = paddle.to_tensor(weight_scales[cur_name], dtype=dtype)
|
|
else:
|
|
logger.info(f"No scale found for layer {cur_name}, remove it")
|
|
parent_layer, sub_name = find_parent_layer_and_sub_name(model, cur_name)
|
|
setattr(parent_layer, sub_name, None)
|
|
|
|
model = ptq.convert(model, inplace=True)
|
|
if os.path.exists(os.path.join(load_quant_path, "model_state.pdparams")):
|
|
logger.info(f"Load model checkpoint from {load_quant_path}")
|
|
model_path = os.path.join(load_quant_path, "model_state.pdparams")
|
|
model_dict = paddle.load(model_path, return_numpy=True)
|
|
model.set_dict(model_dict)
|
|
else:
|
|
raise Exception("Only support load model from pdparams now")
|
|
|
|
|
|
def apply_ptq(quant_args, trainer, ptq_dataloader):
|
|
logger.info("***** Running PTQ *****")
|
|
activation, weight, cachekv, q_config = prepare_qconfig(quant_args)
|
|
|
|
if cachekv is not None:
|
|
set_wrapper_for_attn(trainer.model)
|
|
|
|
skip_list_names = [] if quant_args.skip_list_names is None else quant_args.skip_list_names
|
|
|
|
for cur_name, cur_layer in trainer.model.named_sublayers():
|
|
skip = False
|
|
for k in skip_list_names:
|
|
if k in cur_name:
|
|
logger.info(f"Skip layer {cur_name}")
|
|
skip = True
|
|
if skip:
|
|
continue
|
|
|
|
if type(cur_layer) in [paddle.nn.Linear, ColumnParallelLinear, RowParallelLinear]:
|
|
logger.info(f"PTQ layer: {cur_name}")
|
|
q_config.add_name_config([cur_layer.full_name()], activation=activation, weight=weight)
|
|
|
|
if cachekv is not None and type(cur_layer) in [FuncWrapper]:
|
|
logger.info(f"PTQ layer: {cur_name}")
|
|
# set both act and weight for attention, actually act-k and act-v are quantized
|
|
q_config.add_name_config([cur_layer.full_name()], weight=cachekv[0], activation=cachekv[1])
|
|
|
|
ptq = PTQ(q_config)
|
|
trainer.model = ptq.quantize(trainer.model, inplace=True)
|
|
|
|
# enable observer
|
|
enable_observer(trainer.model)
|
|
logger.info("***** PTQ loop start *****")
|
|
trainer.ptq_loop(
|
|
ptq_dataloader,
|
|
description="PTQ",
|
|
max_eval_iters=quant_args.ptq_step,
|
|
)
|
|
# disable observer
|
|
disable_observer(trainer.model)
|
|
|
|
weight_scales = {}
|
|
act_scales = {}
|
|
cachekv_scales = {}
|
|
for cur_name, cur_layer in trainer.model.named_sublayers():
|
|
if isinstance(cur_layer, ChannelWiseObserver) or isinstance(cur_layer, BaseObserver):
|
|
if "_observer" not in cur_name:
|
|
if "attn_func" in cur_name:
|
|
cur_name = cur_name.replace("attn_func.activation_quanter_v", "cachev_matmul.activation_quanter")
|
|
cur_name = cur_name.replace("attn_func.activation_quanter_k", "cachek_matmul.activation_quanter")
|
|
cachekv_scales[cur_name] = cur_layer.scales().cast("float32").numpy().tolist()
|
|
elif "activation_quanter" in cur_name:
|
|
act_scales[cur_name] = cur_layer.scales().cast("float32").numpy().tolist()
|
|
elif "weight_quanter" in cur_name:
|
|
weight_scales[cur_name] = cur_layer.scales().cast("float32").numpy().tolist()
|
|
|
|
weight_scales_path = os.path.join(trainer.args.output_dir, "weight_scales.json")
|
|
with open(weight_scales_path, "w") as f:
|
|
json.dump(weight_scales, f)
|
|
logger.info(f"Weight scales saved in {weight_scales_path}.")
|
|
|
|
act_scales_path = os.path.join(trainer.args.output_dir, "act_scales.json")
|
|
with open(act_scales_path, "w") as f:
|
|
json.dump(act_scales, f)
|
|
logger.info(f"Activation scales saved in {act_scales_path}.")
|
|
|
|
cachekv_scales_path = os.path.join(trainer.args.output_dir, "cachekv_scales.json")
|
|
with open(cachekv_scales_path, "w") as f:
|
|
json.dump(cachekv_scales, f)
|
|
logger.info(f"CacheKV scales saved in {cachekv_scales_path}.")
|
|
|
|
trainer.model = ptq.convert(trainer.model, inplace=True)
|
|
logger.info("***** PTQ done *****")
|
|
|
|
|
|
def apply_gptq(quant_args, trainer, ptq_dataloader):
|
|
logger.info("***** Running GPTQ *****")
|
|
num_layer = 0
|
|
model = trainer.model
|
|
for cur_name, cur_layer in model.named_sublayers():
|
|
if type(cur_layer) in [paddle.nn.Linear, ColumnParallelLinear, RowParallelLinear]:
|
|
num_layer += 1
|
|
logger.info(f"GPTQ layer: {num_layer}, {cur_name}")
|
|
parent_layer, sub_name = find_parent_layer_and_sub_name(model, cur_name)
|
|
cur_quant_layer = GPTQ(cur_layer)
|
|
setattr(parent_layer, sub_name, cur_quant_layer)
|
|
with paddle.no_grad():
|
|
trainer.ptq_loop(
|
|
ptq_dataloader,
|
|
description="GPTQ",
|
|
max_eval_iters=quant_args.gptq_step,
|
|
)
|
|
cur_quant_layer.fasterquant(percdamp=0.1, groupsize=-1, actorder=True)
|
|
del cur_quant_layer
|
|
setattr(parent_layer, sub_name, cur_layer)
|
|
logger.info("***** GPTQ done *****")
|
|
|
|
|
|
def set_wrapper_for_attn(model: nn.Layer, attn_name="attn_func"):
|
|
for cur_name, cur_layer in model.named_sublayers():
|
|
if hasattr(cur_layer, attn_name):
|
|
logger.info(f"Set wrapper for {attn_name} in {cur_name}")
|
|
cur_layer.attn_func = FuncWrapper(cur_layer.attn_func)
|
|
|
|
|
|
def get_ptq_model_config(model):
|
|
if isinstance(model, PrefixModelForCausalLM):
|
|
base_model_prefix = model.model.base_model_prefix
|
|
else:
|
|
base_model_prefix = model.base_model_prefix
|
|
|
|
if base_model_prefix in ["chatglm"]:
|
|
raise NotImplementedError(f"{model} does not support Shift or Smooth.")
|
|
elif base_model_prefix == "chatglm_v2":
|
|
model_config = {"fused_qkv": False, "parallel_ffn": False, "skip_norm_list": ["rms_norm_56"]}
|
|
elif base_model_prefix == "bloom":
|
|
model_config = {"fused_qkv": True, "parallel_ffn": False}
|
|
elif base_model_prefix == "llama":
|
|
model_config = {"fused_qkv": False, "parallel_ffn": True}
|
|
elif base_model_prefix == "qwen2":
|
|
model_config = {"fused_qkv": False, "parallel_ffn": True}
|
|
else:
|
|
raise ValueError(
|
|
f"Unknown base_model_prefix: {model.base_model_prefix}. Supported base_model_prefix list: chatglm_V2, bloom, llama, qwen2."
|
|
)
|
|
return model_config
|
|
|
|
|
|
def enable_observer(model: nn.Layer):
|
|
# TODO maybe not support pp,tp etc.
|
|
for mod in model.sublayers():
|
|
if hasattr(mod, "observer_enabled"):
|
|
mod.observer_enabled = True
|
|
|
|
|
|
def disable_observer(model: nn.Layer):
|
|
# TODO maybe not support pp,tp etc.
|
|
for mod in model.sublayers():
|
|
if hasattr(mod, "observer_enabled"):
|
|
mod.observer_enabled = False
|
|
|
|
|
|
def add_quant_inp_out_hook(model: nn.Layer, tag_func):
|
|
def get_hook():
|
|
inp_ret = []
|
|
out_ret = []
|
|
|
|
def hook(layer, inp, out):
|
|
nonlocal inp_ret, out_ret
|
|
inp_ret.append(inp[0].flatten().numpy())
|
|
out_ret.append(out.flatten().numpy())
|
|
return out
|
|
|
|
return hook, inp_ret, out_ret
|
|
|
|
inp_dict = dict()
|
|
out_dict = dict()
|
|
|
|
handlers = []
|
|
for cur_name, cur_layer in model.named_sublayers():
|
|
if tag_func(cur_name):
|
|
hook, inp_ret, out_ret = get_hook()
|
|
handle = cur_layer.register_forward_post_hook(hook)
|
|
inp_dict[cur_name] = inp_ret
|
|
out_dict[cur_name] = out_ret
|
|
handlers.append(handle)
|
|
|
|
return inp_dict, out_dict
|
|
|
|
|
|
def save_dict(inp_dict, file_path):
|
|
import pickle
|
|
|
|
with open(file_path, "wb") as f:
|
|
pickle.dump(inp_dict, f)
|