wehub-resource-sync
874 lines
32 KiB
Python
874 lines
32 KiB
Python
import gc
|
|
import torch
|
|
import logging
|
|
import inspect
|
|
import functools
|
|
|
|
from tqdm import tqdm
|
|
from collections import defaultdict
|
|
from typing import Tuple, List, Union, Dict
|
|
|
|
logging.basicConfig(level=logging.ERROR)
|
|
|
|
class AwqQuantizer:
|
|
def __init__(
|
|
self,
|
|
model,
|
|
modules_to_not_convert=None,
|
|
apply_clip=True,
|
|
n_parallel_calib_samples=None,
|
|
max_calib_samples=128,
|
|
max_calib_seq_len=512,
|
|
max_chunk_memory=1024 * 1024 * 1024,
|
|
) -> None:
|
|
self.awq_model = model
|
|
self.model = model
|
|
self.tokenizer = model.tokenizer
|
|
self.w_bit = model.args.quant_bit
|
|
self.group_size = model.args.quant_block
|
|
self.zeropoint = not model.args.sym
|
|
self.calib_data = 'wikitext' if model.args.calib_data is None else model.args.calib_data
|
|
self.split = 'test'
|
|
self.duo_scaling = True
|
|
self.apply_clip = apply_clip
|
|
self.n_parallel_calib_samples = n_parallel_calib_samples
|
|
self.max_calib_samples = max_calib_samples
|
|
self.max_calib_seq_len = max_calib_seq_len
|
|
self.max_chunk_memory = max_chunk_memory
|
|
self.modules_to_not_convert = (
|
|
modules_to_not_convert if modules_to_not_convert is not None else []
|
|
)
|
|
self.modules, self.module_kwargs, self.inps = self.init_quant(
|
|
n_samples=self.max_calib_samples, max_seq_len=self.max_calib_seq_len
|
|
)
|
|
|
|
def pseudo_quantize_tensor(self, w: torch.Tensor):
|
|
org_w_shape = w.shape
|
|
if self.group_size > 0:
|
|
assert org_w_shape[-1] % self.group_size == 0
|
|
w = w.reshape(-1, self.group_size)
|
|
assert w.dim() == 2
|
|
assert torch.isnan(w).sum() == 0
|
|
# zero point quantization
|
|
if self.zeropoint:
|
|
max_val = w.amax(dim=1, keepdim=True)
|
|
min_val = w.amin(dim=1, keepdim=True)
|
|
offset = 1 << (self.w_bit - 1)
|
|
clip_max = offset - 1
|
|
clip_min = -offset
|
|
scales = (max_val - min_val) / (clip_max - clip_min)
|
|
zeros = - torch.round(min_val / scales) + clip_min
|
|
qw = torch.round(w / scales) + zeros
|
|
qw = torch.clamp(qw, clip_min, clip_max)
|
|
w = (qw - zeros) * scales
|
|
zeros = min_val.view(org_w_shape[0], -1)
|
|
else:
|
|
abs_max = w.abs().amax(dim=1, keepdim=True)
|
|
offset = 1 << (self.w_bit - 1)
|
|
clip_max = offset - 1
|
|
clip_min = -clip_max
|
|
scales = abs_max / clip_max
|
|
w = torch.clamp(torch.round(w / scales), clip_min, clip_max) * scales
|
|
zeros = None
|
|
|
|
assert torch.isnan(scales).sum() == 0
|
|
assert torch.isnan(w).sum() == 0
|
|
|
|
scales = scales.view(org_w_shape[0], -1)
|
|
w = w.reshape(org_w_shape)
|
|
|
|
return w, scales, zeros
|
|
|
|
def quantize(self):
|
|
for i in tqdm(range(len(self.modules)), desc="AWQ"):
|
|
# Move module and inputs to correct device
|
|
common_device = next(self.modules[i].parameters()).device
|
|
if common_device is None or str(common_device) == "cpu":
|
|
best_device = AwqQuantizer.get_best_device()
|
|
AwqQuantizer.to_device(self.modules[i], best_device)
|
|
common_device = best_device
|
|
|
|
if self.module_kwargs.get("position_ids") is not None:
|
|
self.module_kwargs["position_ids"] = self.module_kwargs[
|
|
"position_ids"
|
|
].to(common_device)
|
|
|
|
if self.module_kwargs.get("attention_mask") is not None:
|
|
self.module_kwargs["attention_mask"] = self.module_kwargs[
|
|
"attention_mask"
|
|
].to(common_device)
|
|
|
|
self.inps = self.inps.to(common_device)
|
|
|
|
# [STEP 1]: Get layer, extract linear modules, extract input features
|
|
named_linears = AwqQuantizer.get_named_linears(self.modules[i])
|
|
|
|
# Filter out the linear layers we don't want to exclude
|
|
named_linears = AwqQuantizer.exclude_layers_to_not_quantize(
|
|
named_linears, self.modules_to_not_convert
|
|
)
|
|
|
|
input_feat = self._get_input_feat(self.modules[i], named_linears)
|
|
AwqQuantizer.clear_memory()
|
|
|
|
# [STEP 2]: Compute and apply scale list
|
|
module_config = []
|
|
# q, k, v proj
|
|
module_config.append(
|
|
dict(
|
|
prev_op=self.modules[i].input_layernorm,
|
|
layers=[
|
|
self.modules[i].self_attn.q_proj,
|
|
self.modules[i].self_attn.k_proj,
|
|
self.modules[i].self_attn.v_proj,
|
|
],
|
|
inp=input_feat["self_attn.q_proj"],
|
|
module2inspect=self.modules[i].self_attn,
|
|
kwargs=self.module_kwargs,
|
|
)
|
|
)
|
|
# o_proj
|
|
if self.modules[i].self_attn.v_proj.weight.shape == self.modules[i].self_attn.o_proj.weight.shape:
|
|
module_config.append(
|
|
dict(
|
|
prev_op=self.modules[i].self_attn.v_proj,
|
|
layers=[self.modules[i].self_attn.o_proj],
|
|
inp=input_feat["self_attn.o_proj"],
|
|
)
|
|
)
|
|
# mlp gate
|
|
module_config.append(
|
|
dict(
|
|
prev_op=self.modules[i].post_attention_layernorm,
|
|
layers=[self.modules[i].mlp.gate_proj, self.modules[i].mlp.up_proj],
|
|
inp=input_feat["mlp.gate_proj"],
|
|
module2inspect=self.modules[i].mlp,
|
|
)
|
|
)
|
|
# mlp down
|
|
module_config.append(
|
|
dict(
|
|
prev_op=self.modules[i].mlp.up_proj,
|
|
layers=[self.modules[i].mlp.down_proj],
|
|
inp=input_feat["mlp.down_proj"],
|
|
)
|
|
)
|
|
scales_list = [
|
|
self._search_best_scale(self.modules[i], **layer)
|
|
for layer in module_config
|
|
]
|
|
AwqQuantizer.apply_scale(self.modules[i], scales_list, input_feat_dict=input_feat)
|
|
# [STEP 3]: Compute and apply clipping list
|
|
if self.apply_clip:
|
|
clip_list = self._search_best_clip(
|
|
self.modules[i], named_linears, input_feat
|
|
)
|
|
AwqQuantizer.apply_clip(self.modules[i], clip_list)
|
|
|
|
AwqQuantizer.clear_memory()
|
|
AwqQuantizer.to_device(self.modules[i], torch.device('cpu'))
|
|
|
|
|
|
@torch.no_grad()
|
|
def _module_forward(
|
|
self, x: torch.Tensor, module: torch.nn.Module, module_kwargs: Dict
|
|
) -> torch.Tensor:
|
|
|
|
if self.n_parallel_calib_samples is None:
|
|
# runs through all samples at once
|
|
module_output = module(x, **module_kwargs)
|
|
if isinstance(module_output, tuple):
|
|
module_output = module_output[0]
|
|
else:
|
|
# memory efficiently runs through all calibration samples
|
|
# but only n_parallel_calib_samples at a time
|
|
module_output = []
|
|
partitioned_inputs = torch.split(x, self.n_parallel_calib_samples)
|
|
for x_partial in partitioned_inputs:
|
|
partial_output = module(x_partial, **module_kwargs)
|
|
|
|
if isinstance(partial_output, tuple):
|
|
partial_output = partial_output[0]
|
|
|
|
module_output.append(partial_output.cpu())
|
|
|
|
module_output = torch.cat(module_output, dim=0)
|
|
|
|
return module_output
|
|
|
|
@torch.no_grad()
|
|
def _search_best_scale(
|
|
self,
|
|
module,
|
|
prev_op,
|
|
layers: List[torch.nn.Linear],
|
|
inp: torch.Tensor,
|
|
module2inspect=None,
|
|
kwargs={},
|
|
):
|
|
if module2inspect is None:
|
|
assert len(layers) == 1
|
|
module2inspect = layers[0]
|
|
|
|
if "use_cache" in kwargs:
|
|
kwargs.pop("use_cache")
|
|
|
|
# Put x on the right device
|
|
inp = inp.to(next(layers[0].parameters()).device)
|
|
|
|
# [STEP 1]: Compute per-channel mean of normalised weights
|
|
# All layer weights are concatted together
|
|
weight = torch.cat([_m.weight for _m in layers], dim=0)
|
|
org_shape = weight.shape
|
|
# The weights are reshaped to be organised by quantization group
|
|
weight = weight.view(-1, self.group_size)
|
|
# Calculates the relative magnitude of the weights within each of the quantization groups,
|
|
# and rescales each group individually so that each group has weights on a 0-1 scale.
|
|
w_scale = weight.abs() / (weight.abs().amax(dim=1, keepdim=True) + 1e-6)
|
|
# Resizes the rescaled weight matrix back up to its original dimensions
|
|
w_scale = w_scale.view(org_shape)
|
|
# Gets the average rescaled magnitude for each output channel
|
|
w_mean = w_scale.mean(0)
|
|
AwqQuantizer.clear_memory(weight)
|
|
|
|
# [STEP 2]: Compute per-channel mean of the input activation with chunking
|
|
# move inp to cpu to avoid memory leak
|
|
inp_flat = inp.cpu().abs().view(-1, inp.shape[-1])
|
|
num_elements = inp_flat.size(0)
|
|
num_channels = inp_flat.size(1)
|
|
element_size_bytes = inp_flat.element_size() * 2 # multiplied by 2 for FP32
|
|
|
|
# Calculate chunk size dynamically based on max_chunk_memory
|
|
chunk_size = int(self.max_chunk_memory // (element_size_bytes * num_channels))
|
|
chunk_size = min(chunk_size, num_elements)
|
|
|
|
# Use float32 for sum calculation
|
|
x_sum = torch.zeros(num_channels, dtype=torch.float32, device=inp.device)
|
|
|
|
for i in range(0, num_elements, chunk_size):
|
|
end = min(i + chunk_size, num_elements)
|
|
chunk_sum = inp_flat[i:end].to(torch.float32).sum(dim=0)
|
|
x_sum += chunk_sum.to(inp.device)
|
|
|
|
x_mean = (x_sum / num_elements).to(inp.dtype)
|
|
AwqQuantizer.clear_memory(x_sum)
|
|
|
|
inp = inp.to(next(layers[0].parameters()).device)
|
|
# [STEP 3]: Compute output of module
|
|
with torch.no_grad():
|
|
module_kwargs = self._sanitize_kwargs(kwargs, module2inspect)
|
|
fp16_output = self._module_forward(inp, module2inspect, module_kwargs)
|
|
|
|
# [STEP 4]: Compute loss
|
|
best_scales = self._compute_best_scale(
|
|
inp, w_mean, x_mean, module2inspect, layers, fp16_output, module_kwargs
|
|
)
|
|
|
|
return (
|
|
AwqQuantizer.get_op_name(module, prev_op),
|
|
tuple([AwqQuantizer.get_op_name(module, m) for m in layers]),
|
|
best_scales,
|
|
)
|
|
|
|
def _compute_best_scale(
|
|
self,
|
|
x: torch.Tensor,
|
|
w_mean: torch.Tensor,
|
|
x_mean: torch.Tensor,
|
|
module2inspect: torch.nn.Module,
|
|
linears2scale: List[torch.nn.Linear],
|
|
fp16_output: torch.Tensor,
|
|
kwargs: Dict={},
|
|
):
|
|
"""
|
|
Compute loss and select best scales
|
|
|
|
L(s) = || Q(W * s) (s^-1 * X) - W * X ||
|
|
Q: weight quantization function | pseudo_quantize_tensor(W * s)
|
|
X: inputs from calib dataset | X
|
|
W: original weights in FP16 | layer
|
|
s: per channel scaling factor | s^-1 * X
|
|
"""
|
|
n_grid = 20
|
|
history = []
|
|
best_ratio = -1
|
|
best_scales = None
|
|
best_error = float("inf")
|
|
|
|
device = x.device
|
|
x_mean = x_mean.view(-1).to(device)
|
|
w_mean = w_mean.view(-1).to(device)
|
|
|
|
ord_weights = []
|
|
for fc in linears2scale:
|
|
ord_weights.append(fc.weight.data.clone())
|
|
|
|
for ratio in range(n_grid):
|
|
# create new scales
|
|
ratio = ratio / n_grid
|
|
|
|
# NOTE: s^-1 * x is fused here, according to paper
|
|
if self.duo_scaling:
|
|
scales = (x_mean.pow(ratio) / (w_mean.pow(1 - ratio) + 1e-4)).clamp(min=1e-4)
|
|
else:
|
|
scales = x_mean.pow(ratio).clamp(min=1e-4).view(-1)
|
|
scales = scales / (scales.max() * scales.min()).sqrt()
|
|
scales_view = scales.view(1, -1).to(device)
|
|
|
|
# avoid scaling values that overflow
|
|
scales[torch.isinf(scales)] = 1
|
|
scales[torch.isnan(scales)] = 1
|
|
|
|
# Q(W * s)
|
|
for fc in linears2scale:
|
|
fc.weight.mul_(scales_view)
|
|
fc.weight.data = (
|
|
self.pseudo_quantize_tensor(fc.weight.data)[0] / scales_view
|
|
)
|
|
|
|
# W * X
|
|
int_w_output = self._module_forward(x, module2inspect, kwargs)
|
|
|
|
# compute mean squared error (L2 norm)
|
|
loss = self._compute_loss(fp16_output, int_w_output, device)
|
|
|
|
history.append(loss)
|
|
if loss < best_error:
|
|
best_error = loss
|
|
best_ratio = ratio
|
|
best_scales = scales.clone()
|
|
|
|
for fc, ord_weight in zip(linears2scale, ord_weights):
|
|
fc.weight.data = ord_weight.clone()
|
|
|
|
del ord_weights
|
|
|
|
if best_ratio == -1:
|
|
logging.debug(history)
|
|
raise Exception
|
|
|
|
assert torch.isnan(best_scales).sum() == 0, best_scales
|
|
|
|
return best_scales.detach().cpu()
|
|
|
|
@torch.no_grad()
|
|
def _compute_loss(
|
|
self,
|
|
fp16_output: torch.Tensor,
|
|
int_w_output: torch.Tensor,
|
|
device: torch.device,
|
|
):
|
|
loss = 0.0
|
|
fp16_output_flat = fp16_output.view(-1)
|
|
int_w_output_flat = int_w_output.view(-1)
|
|
num_elements = fp16_output_flat.size(0)
|
|
element_size_bytes = fp16_output.element_size()
|
|
|
|
# Calculate chunk size dynamically based on max_chunk_memory
|
|
# Divide the max_chunk_memory by twice the element size
|
|
chunk_size = self.max_chunk_memory // (element_size_bytes * 2)
|
|
chunk_size = min(chunk_size, num_elements)
|
|
|
|
# Split the computation into chunks
|
|
fp16_chunks = torch.split(fp16_output_flat, chunk_size)
|
|
int_w_chunks = torch.split(int_w_output_flat, chunk_size)
|
|
|
|
# Compute the loss for each chunk
|
|
for fp16_chunk, int_w_chunk in zip(fp16_chunks, int_w_chunks):
|
|
chunk_loss = (fp16_chunk.to(device) - int_w_chunk.to(device)).float().pow(2).sum().item()
|
|
loss += chunk_loss
|
|
|
|
# Normalize the loss by the total number of elements
|
|
loss /= num_elements
|
|
|
|
return loss
|
|
|
|
@torch.no_grad()
|
|
def _search_best_clip(self, layer, named_linears, input_feat):
|
|
clip_list = []
|
|
avoid_clipping = ["q_", "k_", "query", "key", "Wqkv"]
|
|
|
|
for name in named_linears:
|
|
# due to qk bmm, it is hard to clip precisely
|
|
if any([_ in name for _ in avoid_clipping]):
|
|
continue
|
|
|
|
named_linears[name].to(AwqQuantizer.get_best_device())
|
|
max_val = self._compute_best_clip(
|
|
named_linears[name].weight, input_feat[name]
|
|
)
|
|
clip_list.append((name, max_val))
|
|
named_linears[name].cpu()
|
|
|
|
return clip_list
|
|
|
|
@torch.no_grad()
|
|
def _compute_best_clip(
|
|
self,
|
|
w: torch.Tensor,
|
|
input_feat: torch.Tensor,
|
|
n_grid=20,
|
|
max_shrink=0.5,
|
|
n_sample_token=512,
|
|
):
|
|
assert w.dim() == 2
|
|
org_w_shape = w.shape
|
|
# w [co, ci] -> [co, 1, n_group, group size]
|
|
# input_feat [n_token, ci] -> [1, n_token, n_group, group size]
|
|
group_size = self.group_size if self.group_size > 0 else org_w_shape[1]
|
|
input_feat = input_feat.view(-1, input_feat.shape[-1])
|
|
input_feat = input_feat.reshape(1, input_feat.shape[0], -1, group_size)
|
|
|
|
# Compute input feature step size (minimum 1)
|
|
step_size = max(1, input_feat.shape[1] // n_sample_token)
|
|
input_feat = input_feat[:, ::step_size]
|
|
|
|
w = w.reshape(org_w_shape[0], 1, -1, group_size)
|
|
|
|
oc_batch_size = 256 if org_w_shape[0] % 256 == 0 else 64 # prevent OOM
|
|
assert org_w_shape[0] % oc_batch_size == 0
|
|
w_all = w
|
|
best_max_val_all = []
|
|
|
|
for i_b in range(org_w_shape[0] // oc_batch_size):
|
|
w = w_all[i_b * oc_batch_size : (i_b + 1) * oc_batch_size]
|
|
|
|
org_max_val = w.abs().amax(dim=-1, keepdim=True) # co, 1, n_group, 1
|
|
|
|
best_max_val = org_max_val.clone()
|
|
min_errs = torch.ones_like(org_max_val) * 1e9
|
|
input_feat = input_feat.to(w.device)
|
|
org_out = (input_feat * w).sum(dim=-1) # co, n_token, n_group
|
|
|
|
for i_s in range(int(max_shrink * n_grid)):
|
|
max_val = org_max_val * (1 - i_s / n_grid)
|
|
min_val = -max_val
|
|
cur_w = torch.clamp(w, min_val, max_val)
|
|
q_w = self.pseudo_quantize_tensor(cur_w)[0]
|
|
cur_out = (input_feat * q_w).sum(dim=-1)
|
|
|
|
# co, 1, n_group, 1
|
|
err = (cur_out - org_out).pow(2).mean(dim=1).view(min_errs.shape)
|
|
del cur_w
|
|
del cur_out
|
|
cur_best_idx = err < min_errs
|
|
min_errs[cur_best_idx] = err[cur_best_idx]
|
|
best_max_val[cur_best_idx] = max_val[cur_best_idx]
|
|
best_max_val_all.append(best_max_val)
|
|
|
|
best_max_val = torch.cat(best_max_val_all, dim=0)
|
|
|
|
AwqQuantizer.clear_memory(input_feat)
|
|
AwqQuantizer.clear_memory(org_out)
|
|
|
|
return best_max_val.squeeze(1)
|
|
|
|
@staticmethod
|
|
@torch.no_grad()
|
|
def apply_clip(module, clip_list: Tuple[str, torch.Tensor]):
|
|
for name, max_val in clip_list:
|
|
layer: torch.nn.Linear = AwqQuantizer.get_op_by_name(module, name)
|
|
layer.to(AwqQuantizer.get_best_device())
|
|
max_val = max_val.to(layer.weight.device)
|
|
org_shape = layer.weight.shape
|
|
layer.weight.data = layer.weight.data.reshape(*max_val.shape[:2], -1)
|
|
layer.weight.data = torch.clamp(layer.weight.data, -max_val, max_val)
|
|
layer.weight.data = layer.weight.data.reshape(org_shape)
|
|
layer.cpu()
|
|
|
|
@staticmethod
|
|
@torch.no_grad()
|
|
def scale_fc_fcs(fc1: torch.nn.Linear, fcs: List[torch.nn.Linear], scales: torch.Tensor):
|
|
if not isinstance(fcs, list):
|
|
fcs = [fcs]
|
|
|
|
scales = scales.to(fc1.weight.device)
|
|
|
|
fc1.weight[-scales.size(0) :].div_(scales.view(-1, 1))
|
|
if fc1.bias is not None:
|
|
fc1.bias.div_(scales.view(-1))
|
|
|
|
for fc in fcs:
|
|
fc.weight.mul_(scales.view(1, -1))
|
|
|
|
for p in fc1.parameters():
|
|
assert torch.isnan(p).sum() == 0
|
|
for fc in fcs:
|
|
for p in fc.parameters():
|
|
assert torch.isnan(p).sum() == 0
|
|
|
|
@staticmethod
|
|
def is_allowed_act_fns(op):
|
|
from transformers.activations import NewGELUActivation, PytorchGELUTanh, GELUActivation
|
|
allowed_act_fns = [
|
|
torch.nn.GELU,
|
|
NewGELUActivation,
|
|
PytorchGELUTanh,
|
|
GELUActivation,
|
|
]
|
|
return (op in allowed_act_fns)
|
|
|
|
@staticmethod
|
|
def is_allowed_norms(op):
|
|
if isinstance(op, torch.nn.LayerNorm):
|
|
return True
|
|
if any(t in str(type(op)) for t in ['LlamaRMSNorm', 'GemmaRMSNorm', 'CohereLayerNorm']):
|
|
return True
|
|
return False
|
|
|
|
@staticmethod
|
|
@torch.no_grad()
|
|
def scale_fc_fc(fc1: torch.nn.Linear, fc2: torch.nn.Linear, scales: torch.Tensor):
|
|
assert isinstance(fc1, torch.nn.Linear)
|
|
assert isinstance(fc2, torch.nn.Linear)
|
|
|
|
scales = scales.to(fc1.weight.device)
|
|
fc1.weight[-scales.size(0) :].div_(scales.view(-1, 1))
|
|
if fc1.bias is not None:
|
|
fc1.bias.div_(scales.view(-1))
|
|
|
|
fc2.weight.mul_(scales.view(1, -1))
|
|
|
|
for p in fc1.parameters():
|
|
assert torch.isnan(p).sum() == 0
|
|
for p in fc2.parameters():
|
|
assert torch.isnan(p).sum() == 0
|
|
|
|
@staticmethod
|
|
@torch.no_grad()
|
|
def scale_ln_fcs(ln: torch.nn.Linear, fcs: List[torch.nn.Linear], scales: torch.Tensor):
|
|
if not isinstance(fcs, list):
|
|
fcs = [fcs]
|
|
|
|
scales = scales.to(ln.weight.device)
|
|
|
|
# GemmaRMSNorm is different from Llama's in that it multiplies
|
|
# (1 + weight) to the output, instead of just weight.
|
|
if 'GemmaRMSNorm' in str(type(ln)):
|
|
ln.weight += 1
|
|
ln.weight.div_(scales)
|
|
ln.weight -= 1
|
|
else:
|
|
ln.weight.div_(scales)
|
|
|
|
if hasattr(ln, "bias") and ln.bias is not None:
|
|
ln.bias.div_(scales)
|
|
|
|
for fc in fcs:
|
|
fc.weight.mul_(scales.view(1, -1))
|
|
|
|
for p in ln.parameters():
|
|
assert torch.isnan(p).sum() == 0
|
|
for fc in fcs:
|
|
for p in fc.parameters():
|
|
assert torch.isnan(p).sum() == 0
|
|
|
|
@staticmethod
|
|
@torch.no_grad()
|
|
def scale_gelu_fc(gelu, fc: torch.nn.Linear, scales: torch.Tensor):
|
|
assert AwqQuantizer.is_allowed_act_fns(gelu)
|
|
assert isinstance(fc, torch.nn.Linear)
|
|
|
|
fc.weight.mul_(scales.view(1, -1).to(fc.weight.device))
|
|
|
|
for p in fc.parameters():
|
|
assert torch.isnan(p).sum() == 0
|
|
|
|
@staticmethod
|
|
def apply_scale(module, scales_list, input_feat_dict=None):
|
|
for prev_op_name, layer_names, scales in scales_list:
|
|
prev_op = AwqQuantizer.get_op_by_name(module, prev_op_name)
|
|
layers = [AwqQuantizer.get_op_by_name(module, name) for name in layer_names]
|
|
|
|
best_device = AwqQuantizer.get_best_device()
|
|
prev_op.to(best_device)
|
|
for layer in layers:
|
|
layer.to(best_device)
|
|
scales.to(best_device)
|
|
if (
|
|
isinstance(prev_op, torch.nn.Linear)
|
|
and type(layers) == list
|
|
and isinstance(layers[0], torch.nn.Linear)
|
|
):
|
|
if len(layers) == 1:
|
|
AwqQuantizer.scale_fc_fc(prev_op, layers[0], scales)
|
|
else:
|
|
AwqQuantizer.scale_fc_fcs(prev_op, layers, scales)
|
|
elif (
|
|
AwqQuantizer.is_allowed_norms(prev_op)
|
|
or "rmsnorm" in str(prev_op.__class__).lower()
|
|
):
|
|
AwqQuantizer.scale_ln_fcs(prev_op, layers, scales)
|
|
|
|
elif AwqQuantizer.is_allowed_act_fns(prev_op):
|
|
AwqQuantizer.scale_gelu_fc(prev_op, layers[0], scales)
|
|
else:
|
|
raise NotImplementedError(f"prev_op {type(prev_op)} not supported yet!")
|
|
|
|
# apply the scaling to input feat if given; prepare it for clipping
|
|
if input_feat_dict is not None:
|
|
for layer_name in layer_names:
|
|
# Skip the modules that are not quantized
|
|
if layer_name in input_feat_dict:
|
|
inp = input_feat_dict[layer_name]
|
|
inp.div_(scales.view(1, -1).to(inp.device))
|
|
|
|
prev_op.cpu()
|
|
for layer in layers:
|
|
layer.cpu()
|
|
scales.cpu()
|
|
|
|
@staticmethod
|
|
def exclude_layers_to_not_quantize(linear_layers, modules_to_not_convert):
|
|
if modules_to_not_convert is None:
|
|
return linear_layers
|
|
|
|
filtered_layers = {}
|
|
for name, linear_layer in linear_layers.items():
|
|
if not any(key in name for key in modules_to_not_convert):
|
|
filtered_layers[name] = linear_layer
|
|
return filtered_layers
|
|
|
|
@staticmethod
|
|
def to_device(module, device):
|
|
for child_name, child_module in module.named_children():
|
|
if child_name == 'self_attn':
|
|
for sub_name, sub_child in child_module.named_children():
|
|
if sub_name != 'config':
|
|
sub_child.to(device)
|
|
else:
|
|
child_module.to(device)
|
|
|
|
|
|
@staticmethod
|
|
def get_named_linears(module):
|
|
linears = {}
|
|
for child_name, child_module in module.named_children():
|
|
if child_name == 'self_attn':
|
|
for name, mod in child_module.named_children():
|
|
if name != 'config':
|
|
if isinstance(mod, torch.nn.Linear):
|
|
linears[f"{child_name}.{name}"] = mod
|
|
else:
|
|
for name, mod in child_module.named_modules():
|
|
if isinstance(mod, torch.nn.Linear):
|
|
full_name = f"{child_name}.{name}" if name else child_name
|
|
linears[full_name] = mod
|
|
|
|
return linears
|
|
|
|
|
|
@staticmethod
|
|
def get_op_by_name(module, op_name):
|
|
for child_name, child_module in module.named_children():
|
|
if child_name == op_name:
|
|
return child_module
|
|
if child_name == 'self_attn':
|
|
for name, mod in child_module.named_children():
|
|
if name != 'config':
|
|
full_name = f"{child_name}.{name}"
|
|
if full_name == op_name:
|
|
return mod
|
|
else:
|
|
for name, mod in child_module.named_modules():
|
|
full_name = f"{child_name}.{name}" if name else child_name
|
|
if full_name == op_name:
|
|
return mod
|
|
|
|
if op_name == "":
|
|
return module
|
|
|
|
raise ValueError(f"Cannot find op {op_name} in module {module}")
|
|
|
|
@staticmethod
|
|
def get_calib_dataset(
|
|
data: Union[str, List[str], List[List[int]]] = "pileval",
|
|
tokenizer=None,
|
|
n_samples=128,
|
|
max_seq_len=512,
|
|
split="train",
|
|
text_column="text",
|
|
):
|
|
if isinstance(data, str):
|
|
from datasets import load_dataset
|
|
if data == "pileval":
|
|
dataset = load_dataset("mit-han-lab/pile-val-backup", split="validation")
|
|
elif data == "wikitext":
|
|
|
|
dataset = load_dataset("Salesforce/wikitext", "wikitext-2-raw-v1", split=split)
|
|
else:
|
|
dataset = load_dataset(data, split=split)
|
|
|
|
elif isinstance(data, list):
|
|
if isinstance(data[0], str):
|
|
dataset = [{text_column: text} for text in data]
|
|
elif isinstance(data[0][0], int):
|
|
dataset = data
|
|
else:
|
|
raise NotImplementedError(
|
|
"Either pass a string to a huggingface dataset or a list"
|
|
"that is preprocessed with one sample of text per element"
|
|
" or a list of list of int for tokenized words."
|
|
)
|
|
else:
|
|
raise NotImplementedError(
|
|
"Either pass a string to a huggingface dataset or a list"
|
|
"that is preprocessed with one sample of text per element"
|
|
" or a list of list of int for tokenized words."
|
|
)
|
|
|
|
samples = []
|
|
n_run = 0
|
|
for data in dataset:
|
|
if isinstance(data, list):
|
|
line_encoded = data
|
|
else:
|
|
line = data[text_column]
|
|
line = line.strip()
|
|
line_encoded = tokenizer.encode(line)
|
|
if len(line_encoded) > max_seq_len:
|
|
continue
|
|
sample = torch.tensor([line_encoded])
|
|
if sample.numel() == 0:
|
|
continue
|
|
samples.append(sample)
|
|
n_run += 1
|
|
if n_run == n_samples:
|
|
break
|
|
# now concatenate all samples and split according to max sequence length
|
|
cat_samples = torch.cat(samples, dim=1)
|
|
n_split = cat_samples.shape[1] // max_seq_len
|
|
logging.debug(f" * Split into {n_split} blocks")
|
|
return [
|
|
cat_samples[:, i * max_seq_len : (i + 1) * max_seq_len] for i in range(n_split)
|
|
]
|
|
|
|
@staticmethod
|
|
def get_best_device():
|
|
if torch.backends.mps.is_available():
|
|
return "mps"
|
|
elif torch.cuda.is_available():
|
|
return "cuda:0"
|
|
else:
|
|
return "cpu"
|
|
|
|
@staticmethod
|
|
def clear_memory(weight=None):
|
|
if weight is not None:
|
|
del weight
|
|
gc.collect()
|
|
torch.cuda.empty_cache()
|
|
|
|
|
|
@staticmethod
|
|
def get_op_name(module, op):
|
|
if module is op:
|
|
return ""
|
|
for child_name, child_module in module.named_children():
|
|
if child_name == 'self_attn':
|
|
if child_module is op:
|
|
return child_name
|
|
for name, mod in child_module.named_children():
|
|
if name != 'config':
|
|
if mod is op:
|
|
return f"{child_name}.{name}"
|
|
for sub_name, sub_mod in mod.named_modules():
|
|
if sub_mod is op:
|
|
full_name = f"{child_name}.{name}.{sub_name}" if sub_name else f"{child_name}.{name}"
|
|
return full_name
|
|
else:
|
|
if child_module is op:
|
|
return child_name
|
|
|
|
for name, mod in child_module.named_modules():
|
|
if mod is op:
|
|
full_name = f"{child_name}.{name}" if name else child_name
|
|
return full_name
|
|
|
|
raise ValueError(f"Cannot find op {op} in module {module}")
|
|
|
|
@staticmethod
|
|
def append_str_prefix(x, prefix):
|
|
if isinstance(x, str):
|
|
return prefix + x
|
|
elif isinstance(x, tuple):
|
|
return tuple([AwqQuantizer.append_str_prefix(y, prefix) for y in x])
|
|
elif isinstance(x, list):
|
|
return [AwqQuantizer.append_str_prefix(y, prefix) for y in x]
|
|
else:
|
|
return x
|
|
|
|
def init_quant(self, n_samples=128, max_seq_len=512):
|
|
modules = self.awq_model.blocks
|
|
samples = AwqQuantizer.get_calib_dataset(
|
|
data=self.calib_data,
|
|
tokenizer=self.tokenizer,
|
|
n_samples=n_samples,
|
|
max_seq_len=max_seq_len,
|
|
split=self.split
|
|
)
|
|
samples = torch.cat(samples[:1], dim=0) # just using 1 batch
|
|
inps = []
|
|
layer_kwargs = {}
|
|
# build inps
|
|
seq_len = samples.numel()
|
|
new_tokens = 0
|
|
best_device = AwqQuantizer.get_best_device()
|
|
inps = self.model.embedding(samples).to(best_device)
|
|
position_ids = self.model.get_position_ids(seq_len, new_tokens)
|
|
rotary_pos_emb = self.model.rotary(position_ids)
|
|
attention_mask = self.model.get_attention_mask(seq_len, new_tokens)
|
|
layer_kwargs["rotary_pos_emb"] = rotary_pos_emb.to(best_device)
|
|
layer_kwargs["attention_mask"] = attention_mask.to(best_device)
|
|
del samples
|
|
AwqQuantizer.clear_memory()
|
|
return modules, layer_kwargs, inps
|
|
|
|
def _get_input_feat(self, layer, named_linears):
|
|
# firstly, get input features of all linear layers
|
|
def cache_input_hook(m, x, y, name, feat_dict):
|
|
x = x[0]
|
|
x = x.detach().cpu()
|
|
feat_dict[name].append(x)
|
|
input_feat = defaultdict(list)
|
|
handles = []
|
|
for name in named_linears:
|
|
handles.append(
|
|
named_linears[name].register_forward_hook(
|
|
functools.partial(cache_input_hook, name=name, feat_dict=input_feat)
|
|
)
|
|
)
|
|
# get output as next layer's input
|
|
|
|
# Sanitize the kwargs in case we use transformers version that contains
|
|
# kwargs that are not handled by the module.
|
|
# Useful for trust_remote_code models.
|
|
module_kwargs = self._sanitize_kwargs(self.module_kwargs, layer)
|
|
|
|
self.inps = self._module_forward(self.inps, layer, module_kwargs)
|
|
for h in handles:
|
|
h.remove()
|
|
# now solve for scaling and clipping
|
|
input_feat = {k: torch.cat(v, dim=0) for k, v in input_feat.items()}
|
|
|
|
return input_feat
|
|
|
|
def _sanitize_kwargs(self, inputs_kwargs, module):
|
|
"""
|
|
Remove the arguments that are not supported in the module's
|
|
forward pass to avoid breaking behaviour between different versions
|
|
of transformers.
|
|
|
|
Args:
|
|
inputs_kwargs (`dict`):
|
|
The input dictionary to pass to the model layer
|
|
module (`torch.nn.Module`):
|
|
Target module to quantize.
|
|
"""
|
|
module_signature = inspect.signature(module.forward).parameters
|
|
sanitized_kwargs = {}
|
|
for k, v in inputs_kwargs.items():
|
|
if k in module_signature:
|
|
sanitized_kwargs[k] = v
|
|
return sanitized_kwargs
|