activation_quantize.py

"""
Activation quantization pass: quantize intermediate activations in addition to weights.
This improves memory efficiency and enables full int8 inference.
"""
import torch
from typing import Dict, Any, Tuple

def compute_activation_scales(model: torch.nn.Module, sample_inputs: torch.Tensor, 
                            num_bits: int = 8) -> Dict[str, float]:
    """
    Run a forward pass to collect activation ranges and compute quantization scales.
    """
    model.eval()
    activation_scales = {}
    
    # Hook to capture intermediate activations
    activations = {}
    
    def hook_fn(name):
        def hook(module, input, output):
            if isinstance(output, torch.Tensor):
                activations[name] = output.detach()
        return hook
    
    # Register hooks
    hooks = []
    for name, module in model.named_modules():
        if name:  # Skip root module
            hook = module.register_forward_hook(hook_fn(name))
            hooks.append(hook)
    
    # Forward pass
    with torch.no_grad():
        output = model(sample_inputs)
        activations['output'] = output.detach()
    
    # Remove hooks
    for hook in hooks:
        hook.remove()
    
    # Compute scales
    for name, activation in activations.items():
        max_val = activation.abs().max().item()
        if max_val == 0:
            scale = 1.0
        else:
            scale = max_val / (2**(num_bits - 1) - 1)
        activation_scales[name] = scale
    
    return activation_scales

def activation_quantization_pass(mir: Dict[str, Any], model: torch.nn.Module, 
                               sample_input: torch.Tensor, num_bits: int = 8) -> Dict[str, Any]:
    """
    Add activation quantization metadata to MIR nodes.
    """
    # Compute activation scales
    act_scales = compute_activation_scales(model, sample_input, num_bits)
    
    # Add activation quantization to each node
    for node in mir["graph"]["nodes"]:
        origin_name = node.get("metadata", {}).get("origin_name")
        if origin_name and origin_name in act_scales:
            if "compression" not in node:
                node["compression"] = {}
            
            node["compression"]["activation_quantized"] = True
            node["compression"]["activation_bits"] = num_bits
            node["compression"]["activation_scale"] = act_scales[origin_name]
            node["compression"]["activation_dtype"] = "int8"
    
    # Store activation scales in global metadata
    if "metadata" not in mir:
        mir["metadata"] = {}
    mir["metadata"]["activation_quantization"] = act_scales
    mir["metadata"]["activation_num_bits"] = num_bits
    
    return mir