quantize.py

"""
Simple per-tensor quantization pass. For MVP we only compute scales and record
quantized metadata. We do NOT yet replace model weights in the running model.
"""
import torch
from typing import Dict, Any

def quantize_state_dict(state_dict: Dict[str, torch.Tensor], num_bits: int = 8):
    quant_meta = {}
    for k, t in state_dict.items():
        if not isinstance(t, torch.Tensor):
            continue
        max_val = t.abs().max().item()  # Ensure this is a scalar
        
        # avoid division by zero
        if max_val == 0:
            scale = 1.0
        else:
            scale = max_val / (2**(num_bits - 1) - 1)  # This should be scalar / scalar
            
        quant_meta[k] = {"quantized": True, "bits": num_bits, "scale": float(scale), "dtype": "int8"}
    return quant_meta

def quantization_pass(mir: Dict[str, Any], model, num_bits: int = 8) -> Dict[str, Any]:
    """
    mir: the MIR dict (will be annotated)
    model: the original torch model (to access state_dict)
    returns updated mir with quantization metadata and embedded weights
    """
    sd = model.state_dict()
    meta = quantize_state_dict(sd, num_bits=num_bits)
    
    # Store quantization metadata
    if "metadata" not in mir:
        mir["metadata"] = {}
    mir["metadata"]["quantization"] = meta
    mir["metadata"]["quantization_num_bits"] = num_bits
    
    # Embed weights into MIR nodes
    _embed_weights_in_mir_nodes(mir, model)
    
    return mir

def _embed_weights_in_mir_nodes(mir: Dict[str, Any], model):
    """Embed PyTorch model weights into MIR nodes"""
    state_dict = model.state_dict()
    
    # Create mapping from module names to parameters
    param_map = {}
    for name, param in state_dict.items():
        param_map[name] = param
    
    # Embed weights into nodes
    for node in mir['graph']['nodes']:
        node_id = node['id']
        op_type = node['op_type']
        
        # Try to find corresponding module weights
        # Handle both "conv1" and "conv_1" naming conventions
        possible_names = [node_id, node_id.replace('_', '.')]
        
        if op_type in ['conv2d', 'linear']:
            for base_name in possible_names:
                weight_key = f"{base_name}.weight"
                bias_key = f"{base_name}.bias"
                
                if weight_key in param_map:
                    if 'weights' not in node:
                        node['weights'] = {}
                    
                    # Convert tensors to lists for JSON serialization
                    node['weights']['weight'] = param_map[weight_key].detach().cpu().numpy().tolist()
                    
                    # Add bias if it exists
                    if bias_key in param_map:
                        node['weights']['bias'] = param_map[bias_key].detach().cpu().numpy().tolist()
                    else:
                        # Create zero bias for layers without bias
                        if op_type == 'conv2d':
                            out_channels = param_map[weight_key].shape[0]
                            node['weights']['bias'] = torch.zeros(out_channels).numpy().tolist()
                        elif op_type == 'linear':
                            out_features = param_map[weight_key].shape[0]
                            node['weights']['bias'] = torch.zeros(out_features).numpy().tolist()
                    
                    break
        
        elif op_type == 'batchnorm2d':
            for base_name in possible_names:
                weight_key = f"{base_name}.weight"
                bias_key = f"{base_name}.bias"
                mean_key = f"{base_name}.running_mean"
                var_key = f"{base_name}.running_var"
                
                if weight_key in param_map:
                    if 'weights' not in node:
                        node['weights'] = {}
                    
                    # Convert tensors to lists for JSON serialization
                    node['weights']['weight'] = param_map[weight_key].detach().cpu().numpy().tolist()
                    node['weights']['bias'] = param_map[bias_key].detach().cpu().numpy().tolist()
                    node['weights']['running_mean'] = param_map[mean_key].detach().cpu().numpy().tolist()
                    node['weights']['running_var'] = param_map[var_key].detach().cpu().numpy().tolist()
                    
                    break


def activation_quantization_pass(mir: Dict[str, Any], model, input_tensor: torch.Tensor, 
                               num_bits: int = 8) -> Dict[str, Any]:
    """
    Compute activation quantization scales using a safe approach that handles 
    complex model architectures like CNNs.
    """
    model.eval()
    activation_scales = {}
    
    with torch.no_grad():
        try:
            # Try to run the full model to get a reference output scale
            output = model(input_tensor)
            output_max = output.abs().max().item()
            output_scale = output_max / (2**(num_bits - 1) - 1) if output_max > 0 else 1.0
            
            print(f"Model output range: [{output.min().item():.4f}, {output.max().item():.4f}]")
            print(f"Output quantization scale: {output_scale:.6f}")
            
        except Exception as e:
            print(f"Warning: Could not run full model: {e}")
            output_scale = 0.1 / (2**(num_bits - 1) - 1)  # Default scale
        
        # Use heuristic activation scales based on operation types
        # In production, you'd use activation hooks for precise measurements
        nodes = mir["graph"]["nodes"]
        
        for node in nodes:
            node_id = node["id"]
            op_type = node.get("op_type", "")
            
            if op_type == "conv2d":
                # Conv layers typically produce moderate activations
                activation_scales[node_id] = output_scale * 2.0
            elif op_type == "relu":
                # ReLU clips negative values, can use tighter scale
                activation_scales[node_id] = output_scale * 1.5
            elif op_type == "linear":
                # Linear layers can have wider ranges
                activation_scales[node_id] = output_scale
            else:
                # Default for unknown ops
                activation_scales[node_id] = output_scale * 1.2
        
        # Set output scale
        activation_scales["output"] = output_scale
    
    # Add to MIR metadata
    if "metadata" not in mir:
        mir["metadata"] = {}
    mir["metadata"]["activation_quantization"] = activation_scales
    mir["metadata"]["activation_num_bits"] = num_bits
    
    return mir