app.py

import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
import gradio as gr
import matplotlib.pyplot as plt
from matplotlib.colors import TwoSlopeNorm
import io
from PIL import Image

# Implementation of the W8A16LinearLayer
class W8A16LinearLayer(nn.Module):
    def __init__(self, in_features, out_features, bias=True, dtype=torch.float32):
        super().__init__()
        
        self.register_buffer(
            "int8_weights",
            torch.randint(
                -128, 127, (out_features, in_features), dtype=torch.int8
            )
        )
        
        self.register_buffer("scales", 
                         torch.randn((out_features), dtype=dtype))
        
        if bias:
            self.register_buffer("bias", 
                             torch.randn((1, out_features), 
                                       dtype=dtype))
        else:
            self.bias = None
            
    def quantize(self, weights):
        """
        Quantize floating point weights to int8 precision
        
        Args:
            weights: Tensor of weights to quantize (shape: out_features x in_features)
            
        Returns:
            None (updates the int8_weights and scales directly)
        """
        w_fp32 = weights.clone().to(torch.float32)

        # Calculate scales as the max absolute value for each output row
        # divided by 127 (max value for int8)
        scales = w_fp32.abs().max(dim=-1).values / 127
        scales = scales.to(weights.dtype)

        # Quantize by dividing by scales and rounding to nearest integer
        int8_weights = torch.round(weights / scales.unsqueeze(1)).to(torch.int8)

        # Update the model parameters
        self.int8_weights = int8_weights
        self.scales = scales
        
        return int8_weights, scales
    
    def forward(self, input):
        """
        Forward pass through the quantized linear layer
        
        Args:
            input: Input tensor (shape: batch_size x seq_len x in_features)
            
        Returns:
            output: Output tensor after the linear transformation
        """
        # Cast int8 weights to input dtype while preserving the values
        casted_weights = self.int8_weights.to(input.dtype)
        
        # Perform the linear multiplication and apply the scaling factor
        output = F.linear(input, casted_weights) * self.scales
        
        # Add bias if present
        if self.bias is not None:
            output = output + self.bias
            
        return output

# Helper functions for visualization

def plot_weight_matrix(weights, title="Weight Matrix"):
    """Create a heatmap visualization of weight matrices"""
    plt.figure(figsize=(10, 8))
    
    # Create a centered colormap
    vmax = max(abs(weights.min().item()), abs(weights.max().item()))
    vmin = -vmax
    norm = TwoSlopeNorm(vmin=vmin, vcenter=0, vmax=vmax)
    
    plt.imshow(weights.detach().numpy(), cmap='RdBu_r', norm=norm)
    plt.colorbar(label='Weight Value')
    plt.title(title)
    
    # Save the plot to a bytes buffer
    buf = io.BytesIO()
    plt.savefig(buf, format='png')
    plt.close()
    buf.seek(0)
    
    return Image.open(buf)

def plot_weight_distribution(weights, title="Weight Distribution"):
    """Create a histogram visualization of weight distributions"""
    plt.figure(figsize=(10, 6))
    
    # Flatten the weights to 1D for histogram
    flat_weights = weights.flatten().detach().numpy()
    
    plt.hist(flat_weights, bins=50, alpha=0.7, color='blue')
    plt.xlabel('Weight Value')
    plt.ylabel('Frequency')
    plt.title(title)
    plt.grid(alpha=0.3)
    
    # Save the plot to a bytes buffer
    buf = io.BytesIO()
    plt.savefig(buf, format='png')
    plt.close()
    buf.seek(0)
    
    return Image.open(buf)

def calculate_quantization_error(original, quantized, scales):
    """Calculate error metrics between original and dequantized weights"""
    # Dequantize the weights
    dequantized = quantized.float() * scales.unsqueeze(1)
    
    # Calculate error metrics
    abs_error = (original - dequantized).abs()
    max_error = abs_error.max().item()
    mean_error = abs_error.mean().item()
    
    return max_error, mean_error, dequantized

# Gradio UI components

def initialize_model(in_features, out_features, with_bias, dtype_str):
    """Initialize a new quantized linear layer model"""
    # Map dtype string to torch dtype
    dtype_map = {
        "float32": torch.float32,
        "float16": torch.float16,
        "bfloat16": torch.bfloat16
    }
    dtype = dtype_map[dtype_str]
    
    # Create the model
    model = W8A16LinearLayer(in_features, out_features, bias=with_bias, dtype=dtype)
    
    # Generate random weights for visualization
    random_weights = torch.randn((out_features, in_features), dtype=dtype)
    
    # Original weights visualization
    weights_vis = plot_weight_matrix(random_weights, "Original Weights")
    dist_vis = plot_weight_distribution(random_weights, "Original Weight Distribution")
    
    # Quantize the weights
    int8_weights, scales = model.quantize(random_weights)
    
    # Quantized weights visualization
    q_weights_vis = plot_weight_matrix(int8_weights, "Quantized Weights (INT8)")
    q_dist_vis = plot_weight_distribution(int8_weights, "Quantized Weight Distribution")
    
    # Calculate quantization error
    max_error, mean_error, dequantized = calculate_quantization_error(
        random_weights, int8_weights, scales
    )
    
    # Dequantized weights visualization
    deq_weights_vis = plot_weight_matrix(dequantized, "Dequantized Weights")
    
    # Error visualization
    error = (random_weights - dequantized).abs()
    error_vis = plot_weight_matrix(error, "Quantization Error (Absolute)")
    
    # Create model summary
    model_info = f"""
    ## Model Configuration
    - Input Features: {in_features}
    - Output Features: {out_features}
    - Bias: {"Yes" if with_bias else "No"}
    - Data Type: {dtype_str}
    
    ## Quantization Stats
    - Original Weights Shape: {random_weights.shape}
    - Quantized Weights Shape: {int8_weights.shape}
    - Scales Shape: {scales.shape}
    - Maximum Quantization Error: {max_error:.6f}
    - Mean Quantization Error: {mean_error:.6f}
    - Memory Savings: {100 * (1 - (int8_weights.element_size() + scales.element_size() * scales.numel()/int8_weights.numel()) / random_weights.element_size()):.2f}%
    """
    
    # Create sample input/output for the model
    sample_input = torch.randn(1, in_features, dtype=dtype)
    sample_output = model(sample_input)
    
    io_info = f"""
    ## Sample Input/Output
    - Input Shape: {sample_input.shape}
    - Output Shape: {sample_output.shape}
    - Output Range: [{sample_output.min().item():.4f}, {sample_output.max().item():.4f}]
    """
    
    return model_info, io_info, weights_vis, q_weights_vis, deq_weights_vis, dist_vis, q_dist_vis, error_vis

def quantize_custom_weights(in_features, out_features, with_bias, dtype_str, weight_pattern):
    """Quantize custom weights based on the selected pattern"""
    # Map dtype string to torch dtype
    dtype_map = {
        "float32": torch.float32,
        "float16": torch.float16,
        "bfloat16": torch.bfloat16
    }
    dtype = dtype_map[dtype_str]
    
    # Create the model
    model = W8A16LinearLayer(in_features, out_features, bias=with_bias, dtype=dtype)
    
    # Generate weights based on pattern
    if weight_pattern == "random":
        custom_weights = torch.randn((out_features, in_features), dtype=dtype)
    elif weight_pattern == "eye":
        # Identity matrix (or closest approximation if dimensions don't match)
        custom_weights = torch.zeros((out_features, in_features), dtype=dtype)
        min_dim = min(out_features, in_features)
        custom_weights[:min_dim, :min_dim] = torch.eye(min_dim, dtype=dtype)
    elif weight_pattern == "ones":
        custom_weights = torch.ones((out_features, in_features), dtype=dtype)
    elif weight_pattern == "alternating":
        custom_weights = torch.ones((out_features, in_features), dtype=dtype)
        # Create a checkerboard pattern
        for i in range(out_features):
            for j in range(in_features):
                if (i + j) % 2 == 1:
                    custom_weights[i, j] = -1.0
    elif weight_pattern == "gradient":
        # Linear gradient from -1 to 1
        x = torch.linspace(-1, 1, in_features)
        y = torch.linspace(-1, 1, out_features)
        xx, yy = torch.meshgrid(x, y, indexing='ij')
        custom_weights = (xx + yy).t().to(dtype)
    
    # Original weights visualization
    weights_vis = plot_weight_matrix(custom_weights, f"Original Weights ({weight_pattern})")
    dist_vis = plot_weight_distribution(custom_weights, "Original Weight Distribution")
    
    # Quantize the weights
    int8_weights, scales = model.quantize(custom_weights)
    
    # Quantized weights visualization
    q_weights_vis = plot_weight_matrix(int8_weights, "Quantized Weights (INT8)")
    q_dist_vis = plot_weight_distribution(int8_weights, "Quantized Weight Distribution")
    
    # Calculate quantization error
    max_error, mean_error, dequantized = calculate_quantization_error(
        custom_weights, int8_weights, scales
    )
    
    # Dequantized weights visualization
    deq_weights_vis = plot_weight_matrix(dequantized, "Dequantized Weights")
    
    # Error visualization
    error = (custom_weights - dequantized).abs()
    error_vis = plot_weight_matrix(error, "Quantization Error (Absolute)")
    
    # Quantization details
    quant_info = f"""
    ## Quantization Details
    - Original Data Type: {dtype_str}
    - Quantized Data Type: int8 (8-bit)
    - Weight Pattern: {weight_pattern}
    
    ## Error Analysis
    - Maximum Quantization Error: {max_error:.6f}
    - Mean Quantization Error: {mean_error:.6f}
    - Memory Savings: {100 * (1 - (int8_weights.element_size() + scales.element_size() * scales.numel()/int8_weights.numel()) / custom_weights.element_size()):.2f}%
    
    ## Tensor Shapes
    - Original Weights: {custom_weights.shape}
    - Quantized Weights: {int8_weights.shape}
    - Quantization Scales: {scales.shape}
    """
    
    # Create row histograms for quantization scales
    plt.figure(figsize=(10, 6))
    plt.hist(scales.detach().cpu().numpy(), bins=30, alpha=0.7, color='green')
    plt.xlabel('Scale Value')
    plt.ylabel('Frequency')
    plt.title('Distribution of Quantization Scales')
    plt.grid(alpha=0.3)
    
    # Save the plot to a bytes buffer
    buf = io.BytesIO()
    plt.savefig(buf, format='png')
    plt.close()
    buf.seek(0)
    scales_vis = Image.open(buf)
    
    return quant_info, weights_vis, q_weights_vis, deq_weights_vis, dist_vis, q_dist_vis, error_vis, scales_vis

# Create Gradio interface
with gr.Blocks(title="8-Bit Weight Quantizer") as demo:
    gr.Markdown("# PyTorch 8-Bit Weight Quantizer")
    gr.Markdown("""
    This tool demonstrates quantization of neural network weights to INT8 precision.
    It implements a custom `W8A16LinearLayer` that uses 8-bit weights with 16-bit activations.
    """)
    
    with gr.Tabs():
        with gr.TabItem("Initialize Model"):
            with gr.Row():
                with gr.Column():
                    in_feat = gr.Slider(minimum=1, maximum=512, value=16, step=1, label="Input Features")
                    out_feat = gr.Slider(minimum=1, maximum=512, value=32, step=1, label="Output Features")
                    with_bias = gr.Checkbox(value=True, label="Include Bias")
                    dtype = gr.Dropdown(choices=["float32", "float16", "bfloat16"], value="float32", label="Data Type")
                    init_btn = gr.Button("Initialize Model")
                
                with gr.Column():
                    model_info = gr.Markdown()
                    io_info = gr.Markdown()
            
            with gr.Row():
                orig_weights = gr.Image(label="Original Weights")
                quant_weights = gr.Image(label="Quantized Weights (INT8)")
                dequant_weights = gr.Image(label="Dequantized Weights")
            
            with gr.Row():
                orig_dist = gr.Image(label="Original Weight Distribution")
                quant_dist = gr.Image(label="Quantized Weight Distribution")
                error_vis = gr.Image(label="Quantization Error")
                
        with gr.TabItem("Custom Quantization"):
            with gr.Row():
                with gr.Column():
                    c_in_feat = gr.Slider(minimum=1, maximum=512, value=16, step=1, label="Input Features")
                    c_out_feat = gr.Slider(minimum=1, maximum=512, value=32, step=1, label="Output Features")
                    c_with_bias = gr.Checkbox(value=True, label="Include Bias")
                    c_dtype = gr.Dropdown(choices=["float32", "float16", "bfloat16"], value="float32", label="Data Type")
                    weight_pattern = gr.Dropdown(
                        choices=["random", "eye", "ones", "alternating", "gradient"], 
                        value="random", 
                        label="Weight Pattern"
                    )
                    quantize_btn = gr.Button("Quantize Weights")
                
                with gr.Column():
                    quant_details = gr.Markdown()
            
            with gr.Row():
                c_orig_weights = gr.Image(label="Original Weights")
                c_quant_weights = gr.Image(label="Quantized Weights (INT8)")
                c_dequant_weights = gr.Image(label="Dequantized Weights")
            
            with gr.Row():
                c_orig_dist = gr.Image(label="Original Weight Distribution")
                c_quant_dist = gr.Image(label="Quantized Weight Distribution")
                c_error_vis = gr.Image(label="Quantization Error")
            
            with gr.Row():
                scales_dist = gr.Image(label="Quantization Scales Distribution")
                
        with gr.TabItem("About"):
            gr.Markdown("""
            ## 8-bit Quantizer Implementation
            
            This implementation includes:
            
            1. **W8A16LinearLayer** - A PyTorch module that uses INT8 weights and FP16/BF16/FP32 activations
            2. **Quantization** - Converts FP32/FP16/BF16 weights to INT8 using per-output-channel scaling
            3. **Visualization** - Shows the impact of quantization on weight distributions and errors
            
            ### How It Works:
            
            1. For each output channel, find the maximum absolute weight value
            2. Scale all weights in that channel so the maximum fits in INT8 range (-128 to 127)
            3. Round scaled weights to integers and store as INT8
            4. During inference, multiply INT8 weights by scaling factors to recover approximate FP values
            
            The quantization process reduces memory usage by up to 75% compared to FP32 weights.
            
            ### References:
            
            - This implementation is based on modern techniques used in LLM quantization
            - Similar methods are used in libraries like bitsandbytes, AutoGPTQ, and GPTQ-for-LLaMa
            """)
    
    # Connect buttons to functions
    init_btn.click(
        initialize_model,
        inputs=[in_feat, out_feat, with_bias, dtype],
        outputs=[model_info, io_info, orig_weights, quant_weights, dequant_weights, orig_dist, quant_dist, error_vis]
    )
    
    quantize_btn.click(
        quantize_custom_weights,
        inputs=[c_in_feat, c_out_feat, c_with_bias, c_dtype, weight_pattern],
        outputs=[quant_details, c_orig_weights, c_quant_weights, c_dequant_weights, c_orig_dist, c_quant_dist, c_error_vis, scales_dist]
    )

# Launch the app
if __name__ == "__main__":
    demo.launch()