There’s a typo, it should be
custom_module_config = {
“float_to_observed_custom_module_class”: {
torch.nn.MultiheadAttention: torch.nn.quantizable.MultiheadAttention
},
“observed_to_quantized_custom_module_class”: {
torch.nn.quantizable.MultiheadAttention: torch.nn.quantized.MultiheadAttention
},
import os
import torch
from torch import nn, Tensor
from torch.ao.quantization import QuantStub, DeQuantStub
class m_AttentionPool2d(nn.Module):
def __init__(self):
super(m_AttentionPool2d, self).__init__()
self.m_atten = nn.MultiheadAttention(embed_dim=2048, num_heads=32, bias=True, add_bias_kv=False)
self.quant = QuantStub()
self.dequant = DeQuantStub()
def forward(self, x):
x = self.quant(x)
(op , _) = self.m_atten(query=x[:1], key=x, value=x, need_weights=False, average_attn_weights=False, is_causal=False)
op = self.dequant(op)
return op
model = m_AttentionPool2d()
example_inputs = torch.rand(50,20,2048)
example_outputs = model(example_inputs)
print(example_outputs.shape) # op - torch.Size([1, 20, 2048])
model.eval()
custom_module_config = {
"float_to_observed_custom_module_class": {torch.nn.MultiheadAttention: torch.nn.quantizable.MultiheadAttention},
"observed_to_quantized_custom_module_class": {torch.nn.quantizable.MultiheadAttention: torch.nn.quantized.MultiheadAttention},
}
print("before", model)
model.qconfig = torch.ao.quantization.get_default_qconfig()
model_prepared = torch.ao.quantization.prepare(model, prepare_custom_config_dict=custom_module_config)
print("prepared", model_prepared)
model_prepared(example_inputs)
model_int8 = torch.ao.quantization.convert(model_prepared,convert_custom_config_dict=custom_module_config,)
print("converted", model_int8)
Result
torch.Size([1, 20, 2048])
before m_AttentionPool2d(
(m_atten): MultiheadAttention(
(out_proj): NonDynamicallyQuantizableLinear(in_features=2048, out_features=2048, bias=True)
)
(quant): QuantStub()
(dequant): DeQuantStub()
)
prepared m_AttentionPool2d(
(m_atten): QuantizableMultiheadAttention(
(out_proj): Linear(
in_features=2048, out_features=2048, bias=True
(activation_post_process): HistogramObserver(min_val=inf, max_val=-inf)
)
(linear_Q): Linear(
in_features=2048, out_features=2048, bias=True
(activation_post_process): HistogramObserver(min_val=inf, max_val=-inf)
)
(linear_K): Linear(
in_features=2048, out_features=2048, bias=True
(activation_post_process): HistogramObserver(min_val=inf, max_val=-inf)
)
(linear_V): Linear(
in_features=2048, out_features=2048, bias=True
(activation_post_process): HistogramObserver(min_val=inf, max_val=-inf)
)
(q_scaling_product): FloatFunctional(
(activation_post_process): HistogramObserver(min_val=inf, max_val=-inf)
)
(quant_attn_output): QuantStub(
(activation_post_process): HistogramObserver(min_val=inf, max_val=-inf)
)
(quant_attn_output_weights): QuantStub(
(activation_post_process): HistogramObserver(min_val=inf, max_val=-inf)
)
(dequant_q): DeQuantStub()
(dequant_k): DeQuantStub()
(dequant_v): DeQuantStub()
)
(quant): QuantStub(
(activation_post_process): HistogramObserver(min_val=inf, max_val=-inf)
)
(dequant): DeQuantStub()
)
converted m_AttentionPool2d(
(m_atten): QuantizedMultiheadAttention(
(out_proj): QuantizedLinear(in_features=2048, out_features=2048, scale=0.011808564886450768, zero_point=59, qscheme=torch.per_channel_affine)
(linear_Q): QuantizedLinear(in_features=2048, out_features=2048, scale=0.025253426283597946, zero_point=65, qscheme=torch.per_channel_affine)
(linear_K): QuantizedLinear(in_features=2048, out_features=2048, scale=0.025217663496732712, zero_point=66, qscheme=torch.per_channel_affine)
(linear_V): QuantizedLinear(in_features=2048, out_features=2048, scale=0.0245734341442585, zero_point=64, qscheme=torch.per_channel_affine)
(q_scaling_product): QFunctional(
scale=1.0, zero_point=0
(activation_post_process): Identity()
)
(quant_attn_output): Quantize(scale=tensor([0.0197]), zero_point=tensor([65]), dtype=torch.quint8)
(quant_attn_output_weights): Quantize(scale=tensor([0.0002]), zero_point=tensor([0]), dtype=torch.quint8)
(dequant_q): DeQuantize()
(dequant_k): DeQuantize()
(dequant_v): DeQuantize()
)
(quant): Quantize(scale=tensor([0.0079]), zero_point=tensor([0]), dtype=torch.quint8)
(dequant): DeQuantize()
)