Initial add of unet/int8 model.

.gitattributes

@@ -33,3 +33,6 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+quant_param.json filter=lfs diff=lfs merge=lfs -text
+quant_params.json filter=lfs diff=lfs merge=lfs -text
+unet/int8/quant_params.json filter=lfs diff=lfs merge=lfs -text

unet/int8/config.json

@@ -0,0 +1,69 @@
+{
+  "_class_name": "UNet2DConditionModel",
+  "_diffusers_version": "0.19.0.dev0",
+  "act_fn": "silu",
+  "addition_embed_type": "text_time",
+  "addition_embed_type_num_heads": 64,
+  "addition_time_embed_dim": 256,
+  "attention_head_dim": [
+    5,
+    10,
+    20
+  ],
+  "block_out_channels": [
+    320,
+    640,
+    1280
+  ],
+  "center_input_sample": false,
+  "class_embed_type": null,
+  "class_embeddings_concat": false,
+  "conv_in_kernel": 3,
+  "conv_out_kernel": 3,
+  "cross_attention_dim": 2048,
+  "cross_attention_norm": null,
+  "down_block_types": [
+    "DownBlock2D",
+    "CrossAttnDownBlock2D",
+    "CrossAttnDownBlock2D"
+  ],
+  "downsample_padding": 1,
+  "dual_cross_attention": false,
+  "encoder_hid_dim": null,
+  "encoder_hid_dim_type": null,
+  "flip_sin_to_cos": true,
+  "freq_shift": 0,
+  "in_channels": 4,
+  "layers_per_block": 2,
+  "mid_block_only_cross_attention": null,
+  "mid_block_scale_factor": 1,
+  "mid_block_type": "UNetMidBlock2DCrossAttn",
+  "norm_eps": 1e-05,
+  "norm_num_groups": 32,
+  "num_attention_heads": null,
+  "num_class_embeds": null,
+  "only_cross_attention": false,
+  "out_channels": 4,
+  "projection_class_embeddings_input_dim": 2816,
+  "resnet_out_scale_factor": 1.0,
+  "resnet_skip_time_act": false,
+  "resnet_time_scale_shift": "default",
+  "sample_size": 128,
+  "time_cond_proj_dim": null,
+  "time_embedding_act_fn": null,
+  "time_embedding_dim": null,
+  "time_embedding_type": "positional",
+  "timestep_post_act": null,
+  "transformer_layers_per_block": [
+    1,
+    2,
+    10
+  ],
+  "up_block_types": [
+    "CrossAttnUpBlock2D",
+    "CrossAttnUpBlock2D",
+    "UpBlock2D"
+  ],
+  "upcast_attention": null,
+  "use_linear_projection": true
+}

unet/int8/params.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1047f8e694b0ce7d2fb0754b519b1d3aa7c316bfe74900474f69a261d0077fc7
+size 5136204272

unet/int8/quant_params.json

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:fbc5f010261abb44cf5c149fc1471cb6cfc272d4c5e94cfdbfd80c9f9d52eb39
+size 85103981

unet/int8/reference/math_model.py

@@ -0,0 +1,126 @@
+import torch.nn as nn
+import torch
+
+def quantize(tensor, scale, zero_point, is_asym=False):
+    if is_asym:
+        clamp_min, clamp_max = torch.tensor(0.), torch.tensor(255.)
+    else:
+        clamp_min, clamp_max = torch.tensor(-128.), torch.tensor(127.)
+    quant_tensor = torch.clamp(torch.round(tensor/scale + zero_point), clamp_min, clamp_max) 
+    return quant_tensor
+
+def dequantize(tensor, scale, zero_point):
+    return (tensor - zero_point) * scale
+
+
+class QuantLinear(nn.Module):
+    def __init__(self, in_ch, out_ch, quant_param):
+        super().__init__()
+        mul_factor = torch.tensor(quant_param['smoothquant_mul']).view(quant_param['smoothquant_mul_shape'])
+        self.register_buffer('mul_factor', mul_factor)
+        self.linear = nn.Linear(in_ch, out_ch)
+        weight_scale = torch.tensor(quant_param['weight_scale']).view(quant_param['weight_scale_shape'])
+        weight_zp = torch.tensor(quant_param['weight_zp']).view(quant_param['weight_zp_shape'])
+        input_scale = torch.tensor(quant_param['input_scale']).view(quant_param['input_scale_shape'])
+        input_zp = torch.tensor(quant_param['input_zp']).view(quant_param['input_zp_shape'])
+        self.register_buffer('weight_scale', weight_scale)
+        self.register_buffer('weight_zp', weight_zp)
+        self.register_buffer('input_scale', input_scale)
+        self.register_buffer('input_zp', input_zp)
+
+    # I.e., "fake quantization"
+    def qdq_forward(self, x):
+        scaled_x = x * self.mul_factor
+        quant_weight = quantize(self.linear.weight, self.weight_scale, self.weight_zp, is_asym=True)
+        quant_input = quantize(scaled_x, self.input_scale, self.input_zp, is_asym=False)
+        dequantized_weight = dequantize(quant_weight, self.weight_scale, self.weight_zp)
+        dequantized_input = dequantize(quant_input, self.input_scale, self.input_zp)
+        out = torch.nn.functional.linear(dequantized_input, dequantized_weight, self.linear.bias)
+        return out
+
+    # Accelerated version
+    def qop_forward(self, x):
+        # With an integer linear kernel, if the weight zero point is not zero,
+        # A correction term must be calculated to correct the output.
+        # The correction term calculated as follows:
+        #  - sum the input tensor across the dot-product dimentions: (e.g., `torch.sum(quant_input, dim=-1)`)
+        #  - multiply this sum with every weight zero-point (e.g., `torch.sum(quant_input, dim=-1) * self.weight_zp`
+        #  - Subtract from previous output (e.g., `quant_output -= torch.sum(quant_input, dim=-1) * self.weight_zp`)
+        #  - All other code is just to make sure the broadcasting semantics work correctly
+        weight_zp_int8 = (self.weight_zp - 128).to(torch.int8).to(torch.float32) # Conversion from uint8 -> int8, can be computed offline
+        quant_weight = quantize(self.linear.weight, self.weight_scale, weight_zp_int8, is_asym=False).to(torch.int8)
+        fused_input_scale = self.input_scale / self.mul_factor # Fuse SmoothQuant and input scales, can be computed offline
+        quant_input = quantize(x, fused_input_scale, self.input_zp, is_asym=False).to(torch.int8)
+        quant_output = torch.nn.functional.linear(quant_input.to(torch.float32), quant_weight.to(torch.float32), None).to(torch.int32) # Convert inputs to FP32 to avoid F.linear quantizing the output to int8
+        correction = torch.sum(quant_input, dim=-1, keepdim=True).to(torch.int32) * weight_zp_int8.to(torch.int8).view([1]*(quant_input.ndim-1) + [self.weight_zp.nelement()]) # Correct for weight zero-point
+        quant_output = quant_output - correction
+        output = dequantize(quant_output, (self.weight_scale * self.input_scale).view([1]*(quant_output.ndim-1) + [(self.weight_scale * self.input_scale).nelement()]), 0.0)
+        output += self.linear.bias
+        return output
+
+    def forward(self, x, qop=False):
+        if qop:
+            return self.qop_forward(x)
+        else:
+            return self.qdq_forward(x)
+
+class QuantConv2d(nn.Module):
+    def __init__(self, in_ch, out_ch, kernel_size, quant_param):
+        super().__init__()
+        mul_factor = torch.tensor(quant_param['smoothquant_mul']).view(quant_param['smoothquant_mul_shape'])
+        self.register_buffer('mul_factor', mul_factor)
+        self.conv2d = nn.Conv2d(in_ch, out_ch, kernel_size)
+        weight_scale = torch.tensor(quant_param['weight_scale']).view(quant_param['weight_scale_shape'])
+        weight_zp = torch.tensor(quant_param['weight_zp']).view(quant_param['weight_zp_shape'])
+        input_scale = torch.tensor(quant_param['input_scale']).view(quant_param['input_scale_shape'])
+        input_zp = torch.tensor(quant_param['input_zp']).view(quant_param['input_zp_shape'])
+        self.register_buffer('weight_scale', weight_scale)
+        self.register_buffer('weight_zp', weight_zp)
+        self.register_buffer('input_scale', input_scale)
+        self.register_buffer('input_zp', input_zp)
+
+    # I.e., "fake quantization"
+    def qdq_forward(self, x):
+        scaled_x = x * self.mul_factor
+        quant_weight = quantize(self.conv2d.weight, self.weight_scale, self.weight_zp, is_asym=True)
+        quant_input = quantize(scaled_x, self.input_scale, self.input_zp, is_asym=False)
+        dequantized_weight = dequantize(quant_weight, self.weight_scale, self.weight_zp)
+        dequantized_input = dequantize(quant_input, self.input_scale, self.input_zp)
+        out = torch.nn.functional.conv2d(dequantized_input, dequantized_weight, self.conv2d.bias)
+        return out
+
+    # Accelerated version
+    def qop_forward(self, x):
+        # With an integer conv2d kernel, if the weight zero point is not zero,
+        # A correction term must be calculated to correct the output.
+        # Conceptually, it's identical to the linear case except that it's difficult
+        # to reduce the input across the dot-product dimension. This leaves us with two obvious options:
+        #  1. Manually compute the reduction via Im2Col -> `torch.sum`
+        #  2. Add an extra _output channel_ to the convolution with a kernel made from all ones (e.g., `torch.ones()`)
+        # In this example, I've used option #2.
+        # The correction term is then calculated as follows:
+        #  - Add an extra output channel to the weight tensor with all values equal to 1 to calculate the sum (e.g., `torch.cat((quant_weight, torch.ones(shape)), dim=0)`)
+        #  - Extract the sum from the output tensor (e.g., `sum = quant_output[:,-1,:,:]`)
+        #  - multiply this sum with every weight zero-point (e.g., `sum * self.weight_zp`
+        #  - Subtract from previous output (e.g., `quant_output -= sum * self.weight_zp`)
+        #  - All other code is just to make sure the broadcasting semantics work correctly
+        weight_zp_int8 = (self.weight_zp - 128).to(torch.int8).to(torch.float32) # Conversion from uint8 -> int8, can be computed offline
+        quant_weight = quantize(self.conv2d.weight, self.weight_scale, weight_zp_int8, is_asym=False).to(torch.int8)
+        b_shape = list(quant_weight.shape) # Used for weight zero-point correction
+        b_shape[0] = 1 # Used for weight zero-point correction
+        weight_cat = torch.ones((1,1,1,1)).broadcast_to(b_shape).to(torch.int8) # Used for weight zero-point correction
+        quant_weight = torch.cat((quant_weight,weight_cat),dim=0).to(torch.int8) # Create extra output channel, used for weight zero-point correction
+        fused_input_scale = self.input_scale / self.mul_factor # Fuse SmoothQuant and input scales, can be computed offline
+        quant_input = quantize(x, fused_input_scale, self.input_zp, is_asym=False).to(torch.int8)
+        quant_output = torch.nn.functional.conv2d(quant_input.to(torch.float32), quant_weight.to(torch.float32), None).to(torch.int32) # Convert inputs to FP32 to avoid F.conv2d quantizing the output to int8
+        correction = quant_output[:,-1,:,:] * weight_zp_int8.to(torch.int8).view([1, self.weight_zp.nelement()] +  [1]*(quant_output.ndim-2)) # Correct zero-point for weight
+        quant_output = quant_output[:,:-1,:,:] - correction
+        output = dequantize(quant_output, (self.weight_scale * self.input_scale).view([1, (self.weight_scale * self.input_scale).nelement()] + [1]*(quant_output.ndim-2)), 0.0)
+        output += self.conv2d.bias.view([1, self.conv2d.bias.nelement()] + [1]*(quant_output.ndim-2))
+        return output
+
+    def forward(self, x, qop=False):
+        if qop:
+            return self.qop_forward(x)
+        else:
+            return self.qdq_forward(x)

unet/int8/reference/test_quant_conv2d.py

@@ -0,0 +1,39 @@
+import torch
+import torch.nn as nn
+from math_model import QuantConv2d
+
+torch.manual_seed(0)
+
+batch_size = 1
+out_ch = 128
+in_ch = 64
+k = 3
+h = 5
+w = 5
+
+i = 2*torch.rand((batch_size,in_ch,h,w)) - 1.
+l = nn.Conv2d(in_ch, out_ch, k, bias=True)
+
+quant_params = {
+    'smoothquant_mul': torch.rand((in_ch,)),
+    'smoothquant_mul_shape': (1,in_ch,1,1),
+    'weight_scale': torch.rand((out_ch,)),
+    'weight_scale': torch.max(torch.abs(torch.flatten(l.weight, start_dim=1)), dim=1).values / 128.,
+    'weight_scale_shape': (out_ch,1,1,1),
+    'weight_zp': torch.clamp(torch.round((torch.mean((l.weight), dim=(1,2,3))) * (128 / torch.max(torch.abs(torch.flatten(l.weight, start_dim=1)), dim=1).values)) + 128, 0, 255),
+    'weight_zp_shape': (out_ch,1,1,1),
+    'input_scale': torch.max(torch.abs(i)) / 128.,
+    'input_scale_shape': tuple(),
+    'input_zp': torch.zeros((1,)),
+    'input_zp_shape': tuple(),
+}
+
+print(quant_params)
+
+ql = QuantConv2d(in_ch, out_ch, k, quant_params)
+ql.conv2d.load_state_dict(l.state_dict())
+o_qdq = ql(i)
+o_qop = ql(i, qop=True)
+print(o_qdq.shape)
+print(o_qop.shape)
+print(o_qdq - o_qop)

unet/int8/reference/test_quant_linear.py

@@ -0,0 +1,35 @@
+import torch
+import torch.nn as nn
+from math_model import QuantLinear
+
+torch.manual_seed(0)
+
+batch_size = 1
+out_ch = 128
+in_ch = 64
+
+i = 2*torch.rand((batch_size,in_ch)) - 1.
+l = nn.Linear(in_ch, out_ch, bias=True)
+
+quant_params = {
+    'smoothquant_mul': torch.rand((in_ch,)),
+    'smoothquant_mul_shape': (1,in_ch),
+    'weight_scale': torch.max(torch.abs(l.weight), dim=1).values / 128.,
+    'weight_scale_shape': (out_ch,1),
+    'weight_zp': torch.clamp(torch.round((torch.mean((l.weight), dim=1)) * (128 / torch.max(torch.abs(l.weight), dim=1).values)) + 128, 0, 255),
+    'weight_zp_shape': (out_ch,1),
+    'input_scale': torch.max(torch.abs(i)) / 128.,
+    'input_scale_shape': tuple(),
+    'input_zp': torch.zeros((1,)),
+    'input_zp_shape': tuple(),
+}
+
+print(quant_params)
+
+ql = QuantLinear(in_ch, out_ch, quant_params)
+ql.linear.load_state_dict(l.state_dict())
+o_qdq = ql(i)
+o_qop = ql(i, qop=True)
+print(o_qdq.shape)
+print(o_qop.shape)
+print(o_qdq - o_qop)