File size: 3,808 Bytes
11690a5 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 | # File authors: Haotian Tang, Shang Yang, Yujun Lin, Song Han
# @article{lin2024awq,
# title={AWQ: Activation-aware Weight Quantization for On-Device LLM Compression and Acceleration},
# author={Lin, Ji and Tang, Jiaming and Tang, Haotian and Yang, Shang and Chen, Wei-Ming and Wang, Wei-Chen and Xiao, Guangxuan and Dang, Xingyu and Gan, Chuang and Han, Song},
# journal={Proceedings of Machine Learning and Systems},
# volume={6},
# pages={87--100},
# year={2024}
# }
import torch
class ActivationBuffer:
"""
Pre-allocated Buffer for activation in the siglip model.
Args:
model: The input model
batched_seq_len: The batched sequence length. Sum of all the sequence lengths in the batch.
"""
def __init__(self, model):
self.model_class = model.__class__.__name__
if self.model_class == "SiglipEncoder":
self.model_dtype = model.layers[0].self_attn.k_proj.weight.dtype
self.device = "cuda"
assert self.model_class in [
"SiglipEncoder",
], f"model_class: {self.model_class} is currently not supported."
assert (
self.model_dtype == torch.float16
), f"model_dtype is expected to be fp16. Current: {self.model_dtype}."
self.intermediate_size = model.config.intermediate_size
self.hidden_size = model.config.hidden_size
def allocate_activation_buffer(self, batched_seq_len):
if self.model_class == "SiglipEncoder":
self.__allocate_activation_buffer_siglip(batched_seq_len)
else:
raise NotImplementedError(
f"model_class: {self.model_class} is currently not supported."
)
def __allocate_activation_buffer_siglip(self, batched_seq_len):
# Allocate fp16 activation buffer.
self.act_buffer = torch.empty(
(batched_seq_len * max(self.hidden_size * 3, 2 * self.intermediate_size)),
device=self.device,
dtype=torch.float16,
)
self.qkv_proj_act_buffer = self.act_buffer[
: batched_seq_len * self.hidden_size * 3
].view(
batched_seq_len, self.hidden_size * 3
) # qkv
self.in_out_fc2_act_buffer = self.act_buffer[
: batched_seq_len * self.hidden_size
].view(
batched_seq_len, self.hidden_size
) # LN1, Wo_out, LN2, all_out
self.fc1_buffer = self.act_buffer[
: batched_seq_len * self.intermediate_size
].view(batched_seq_len, self.intermediate_size)
self.actfn_buffer = self.act_buffer[
batched_seq_len
* self.intermediate_size : 2
* batched_seq_len
* self.intermediate_size
].view(batched_seq_len, self.intermediate_size)
# Allocate quantized activation buffer.
self.quantized_act_buffer = torch.empty(
(batched_seq_len * max(self.hidden_size, self.intermediate_size)),
device=self.device,
dtype=torch.int8,
)
self.quantized_hidden_states_buffer = self.quantized_act_buffer[
: batched_seq_len * self.hidden_size
].view(
batched_seq_len, self.hidden_size
) # Wo_in,
self.quantized_mlp_act_buffer = self.quantized_act_buffer[
: batched_seq_len * self.intermediate_size
].view(batched_seq_len, self.intermediate_size)
# per token
self.quantized_scale_buffer = torch.empty(
(batched_seq_len), device=self.device, dtype=torch.float16
)
# For faster act-quant implementation
self.tmp = torch.empty(
(batched_seq_len * self.intermediate_size),
device=self.device,
dtype=torch.float16,
)
|