add quantization

quantization.py

@@ -0,0 +1,313 @@
+"""
+Mixed-Precision Quantization Script for Small Language Models
+Supports selective quantization of different model components with configurable bitwidths.
+"""
+
+import torch
+import torch.nn as nn
+from transformers import AutoModelForCausalLM, AutoTokenizer, AutoConfig
+import argparse
+import os
+import json
+from pathlib import Path
+from typing import Dict, Optional, Tuple
+import time
+
+class MixedPrecisionQuantizer:
+    """
+    Quantizes model components with different precision levels.
+    Supports more aggressive quantization for attention layers while
+    preserving higher precision for FFN layers.
+    """
+    
+    def __init__(
+        self,
+        model_name: str,
+        attention_bits: int = 4,
+        ffn_bits: int = 8,
+        embedding_bits: int = 8,
+        output_dir: str = "./quantized_models",
+        device: str = "cuda" if torch.cuda.is_available() else "cpu"
+    ):
+        self.model_name = model_name
+        self.attention_bits = attention_bits
+        self.ffn_bits = ffn_bits
+        self.embedding_bits = embedding_bits
+        self.output_dir = Path(output_dir)
+        self.device = device
+        
+        # Create output directory
+        self.output_dir.mkdir(parents=True, exist_ok=True)
+        
+        print(f"Initializing quantizer for {model_name}")
+        print(f"Attention layers: {attention_bits}-bit")
+        print(f"FFN layers: {ffn_bits}-bit")
+        print(f"Embeddings: {embedding_bits}-bit")
+        print(f"Device: {device}")
+    
+    def load_model(self) -> Tuple[nn.Module, AutoTokenizer]:
+        """Load the pretrained model and tokenizer."""
+        print(f"\nLoading model: {self.model_name}")
+        start_time = time.time()
+        
+        # Load with low_cpu_mem_usage for large models
+        model = AutoModelForCausalLM.from_pretrained(
+            self.model_name,
+            torch_dtype=torch.float32,
+            low_cpu_mem_usage=True,
+            trust_remote_code=True
+        )
+        
+        tokenizer = AutoTokenizer.from_pretrained(
+            self.model_name,
+            trust_remote_code=True
+        )
+        
+        load_time = time.time() - start_time
+        print(f"Model loaded in {load_time:.2f} seconds")
+        
+        # Calculate original model size
+        param_count = sum(p.numel() for p in model.parameters())
+        param_size_mb = sum(p.numel() * p.element_size() for p in model.parameters()) / (1024 ** 2)
+        print(f"Parameters: {param_count:,} ({param_size_mb:.2f} MB)")
+        
+        return model, tokenizer
+    
+    def quantize_linear_layer(self, layer: nn.Linear, bits: int) -> nn.Linear:
+        """
+        Quantize a linear layer to specified bit width using symmetric quantization.
+        """
+        if bits == 32:
+            return layer
+        
+        weight = layer.weight.data
+        bias = layer.bias.data if layer.bias is not None else None
+        
+        # Symmetric quantization
+        qmin = -(2 ** (bits - 1))
+        qmax = 2 ** (bits - 1) - 1
+        
+        # Calculate scale
+        max_val = torch.max(torch.abs(weight))
+        scale = max_val / qmax
+        
+        # Quantize
+        weight_q = torch.clamp(torch.round(weight / scale), qmin, qmax)
+        
+        # Store quantized weights and scale
+        layer.weight.data = weight_q.to(torch.int8 if bits <= 8 else torch.int16)
+        layer.weight_scale = scale
+        layer.quantized = True
+        layer.bits = bits
+        
+        return layer
+    
+    def identify_layer_type(self, name: str, module: nn.Module) -> str:
+        """
+        Identify if a layer is part of attention, FFN, embedding, or other components.
+        """
+        name_lower = name.lower()
+        
+        # Attention-related patterns
+        attention_patterns = [
+            'attn', 'attention', 'q_proj', 'k_proj', 'v_proj', 
+            'qkv', 'query', 'key', 'value', 'o_proj', 'out_proj',
+            'c_attn', 'c_proj'
+        ]
+        
+        # FFN-related patterns
+        ffn_patterns = [
+            'mlp', 'ffn', 'fc', 'dense', 'intermediate',
+            'gate_proj', 'up_proj', 'down_proj', 'w1', 'w2', 'w3'
+        ]
+        
+        # Embedding patterns
+        embedding_patterns = ['embed', 'wte', 'wpe', 'lm_head']
+        
+        if any(pattern in name_lower for pattern in attention_patterns):
+            return 'attention'
+        elif any(pattern in name_lower for pattern in ffn_patterns):
+            return 'ffn'
+        elif any(pattern in name_lower for pattern in embedding_patterns):
+            return 'embedding'
+        else:
+            return 'other'
+    
+    def quantize_model(self, model: nn.Module) -> Tuple[nn.Module, Dict]:
+        """
+        Apply mixed-precision quantization to the model.
+        """
+        print("\nApplying mixed-precision quantization...")
+        start_time = time.time()
+        
+        stats = {
+            'attention_layers': 0,
+            'ffn_layers': 0,
+            'embedding_layers': 0,
+            'other_layers': 0,
+            'total_quantized': 0
+        }
+        
+        # Iterate through all modules
+        for name, module in model.named_modules():
+            if isinstance(module, nn.Linear):
+                layer_type = self.identify_layer_type(name, module)
+                
+                # Select quantization bitwidth based on layer type
+                if layer_type == 'attention':
+                    bits = self.attention_bits
+                    stats['attention_layers'] += 1
+                elif layer_type == 'ffn':
+                    bits = self.ffn_bits
+                    stats['ffn_layers'] += 1
+                elif layer_type == 'embedding':
+                    bits = self.embedding_bits
+                    stats['embedding_layers'] += 1
+                else:
+                    bits = self.ffn_bits  # Default to FFN bitwidth
+                    stats['other_layers'] += 1
+                
+                # Quantize the layer
+                self.quantize_linear_layer(module, bits)
+                stats['total_quantized'] += 1
+        
+        quant_time = time.time() - start_time
+        print(f"\nQuantization completed in {quant_time:.2f} seconds")
+        print(f"Quantized layers breakdown:")
+        print(f"  - Attention: {stats['attention_layers']} layers ({self.attention_bits}-bit)")
+        print(f"  - FFN: {stats['ffn_layers']} layers ({self.ffn_bits}-bit)")
+        print(f"  - Embedding: {stats['embedding_layers']} layers ({self.embedding_bits}-bit)")
+        print(f"  - Other: {stats['other_layers']} layers ({self.ffn_bits}-bit)")
+        print(f"  - Total quantized: {stats['total_quantized']} layers")
+        
+        return model, stats
+    
+    def save_quantized_model(
+        self,
+        model: nn.Module,
+        tokenizer: AutoTokenizer,
+        stats: Dict
+    ) -> str:
+        """Save the quantized model, tokenizer, and metadata."""
+        # Create model-specific output directory
+        model_short_name = self.model_name.split('/')[-1]
+        quant_config = f"attn{self.attention_bits}_ffn{self.ffn_bits}_emb{self.embedding_bits}"
+        save_dir = self.output_dir / f"{model_short_name}_{quant_config}"
+        save_dir.mkdir(parents=True, exist_ok=True)
+        
+        print(f"\nSaving quantized model to: {save_dir}")
+        
+        # Save model
+        model.save_pretrained(save_dir)
+        
+        # Save tokenizer
+        tokenizer.save_pretrained(save_dir)
+        
+        # Calculate quantized model size
+        quantized_size_mb = sum(
+            p.numel() * p.element_size() for p in model.parameters()
+        ) / (1024 ** 2)
+        
+        # Save metadata
+        metadata = {
+            'original_model': self.model_name,
+            'quantization_config': {
+                'attention_bits': self.attention_bits,
+                'ffn_bits': self.ffn_bits,
+                'embedding_bits': self.embedding_bits
+            },
+            'layer_stats': stats,
+            'model_size_mb': quantized_size_mb,
+            'quantization_timestamp': time.strftime('%Y-%m-%d %H:%M:%S')
+        }
+        
+        with open(save_dir / 'quantization_metadata.json', 'w') as f:
+            json.dump(metadata, f, indent=2)
+        
+        print(f"Quantized model size: {quantized_size_mb:.2f} MB")
+        print(f"Metadata saved to: {save_dir / 'quantization_metadata.json'}")
+        
+        return str(save_dir)
+    
+    def run(self) -> str:
+        """Execute the full quantization pipeline."""
+        print("=" * 80)
+        print("MIXED-PRECISION QUANTIZATION PIPELINE")
+        print("=" * 80)
+        
+        # Load model
+        model, tokenizer = self.load_model()
+        
+        # Quantize model
+        quantized_model, stats = self.quantize_model(model)
+        
+        # Save quantized model
+        save_path = self.save_quantized_model(quantized_model, tokenizer, stats)
+        
+        print("\n" + "=" * 80)
+        print("QUANTIZATION COMPLETE")
+        print("=" * 80)
+        print(f"Saved to: {save_path}")
+        
+        return save_path
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="Mixed-Precision Quantization for Small Language Models"
+    )
+    parser.add_argument(
+        '--model_name',
+        type=str,
+        required=True,
+        help='HuggingFace model name or path'
+    )
+    parser.add_argument(
+        '--attention_bits',
+        type=int,
+        default=4,
+        help='Bit width for attention layers (default: 4)'
+    )
+    parser.add_argument(
+        '--ffn_bits',
+        type=int,
+        default=8,
+        help='Bit width for FFN layers (default: 8)'
+    )
+    parser.add_argument(
+        '--embedding_bits',
+        type=int,
+        default=8,
+        help='Bit width for embedding layers (default: 8)'
+    )
+    parser.add_argument(
+        '--output_dir',
+        type=str,
+        default='./quantized_models',
+        help='Output directory for quantized models'
+    )
+    parser.add_argument(
+        '--device',
+        type=str,
+        default='cuda' if torch.cuda.is_available() else 'cpu',
+        help='Device to use (cuda/cpu)'
+    )
+    
+    args = parser.parse_args()
+    
+    # Initialize quantizer
+    quantizer = MixedPrecisionQuantizer(
+        model_name=args.model_name,
+        attention_bits=args.attention_bits,
+        ffn_bits=args.ffn_bits,
+        embedding_bits=args.embedding_bits,
+        output_dir=args.output_dir,
+        device=args.device
+    )
+    
+    # Run quantization
+    quantizer.run()
+
+
+if __name__ == "__main__":
+    main()