Model Compression: 14GB to 450MB While Keeping 90% Quality

# Core dependencies
pip install torch>=2.0.0 torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118
 
# HuggingFace ecosystem
pip install transformers>=4.36.0 datasets>=2.14.0 accelerate>=0.25.0
 
# Quantization libraries (CUDA required)
pip install bitsandbytes>=0.44.0  # INT8 quantization
pip install auto-gptq[triton]>=0.7.0  # INT4 quantization with GPTQ
 
# Parameter-efficient fine-tuning
pip install peft>=0.5.0  # LoRA implementation
 
# Evaluation and utilities
pip install sentencepiece protobuf

# Test all dependencies
import torch
import transformers
from datasets import load_dataset
from peft import LoraConfig
import bitsandbytes as bnb
 
print(f"PyTorch: {torch.__version__}")
print(f"CUDA available: {torch.cuda.is_available()}")
print(f"Transformers: {transformers.__version__}")
 
# Test GPTQ availability
try:
    from auto_gptq import AutoGPTQForCausalLM
    print("✓ GPTQ available")
except ImportError:
    print("✗ GPTQ not available (install auto-gptq[triton])")
 
# Expected output:
# PyTorch: 2.0.0+cu118
# CUDA available: True
# Transformers: 4.36.0
# ✓ GPTQ available

from transformers import LlamaConfig, LlamaForCausalLM
 
# Teacher: Llama-7B (already trained)
teacher = LlamaForCausalLM.from_pretrained("meta-llama/Llama-2-7b-hf")
teacher.eval()
for param in teacher.parameters():
    param.requires_grad = False
 
# Student: Scaled-down architecture
student_config = LlamaConfig(
    hidden_size=1536,        # vs 4096 in teacher
    num_hidden_layers=16,    # vs 32 in teacher
    num_attention_heads=12,  # vs 32 in teacher
    intermediate_size=4096,  # vs 11008 in teacher
)
 
student = LlamaForCausalLM(student_config)
print(f"Teacher params: {teacher.num_parameters() / 1e9:.2f}B")
print(f"Student params: {student.num_parameters() / 1e9:.2f}B")
# Teacher params: 6.74B
# Student params: 1.52B

import torch
import torch.nn.functional as F
 
def distillation_loss(
    student_logits,
    teacher_logits,
    labels,
    temperature=2.0,
    alpha=0.7
):
    """
    Combined distillation + task loss.
    
    Args:
        student_logits: [batch, seq_len, vocab_size]
        teacher_logits: [batch, seq_len, vocab_size]
        labels: [batch, seq_len]
        temperature: Softening parameter (higher = softer distribution)
        alpha: Weight for distillation vs task loss (0=task only, 1=distill only)
    
    Returns:
        loss: Scalar tensor
    """
    # Reshape for cross-entropy
    batch_size, seq_len, vocab_size = student_logits.shape
    student_logits_flat = student_logits.view(-1, vocab_size)
    teacher_logits_flat = teacher_logits.view(-1, vocab_size)
    labels_flat = labels.view(-1)
    
    # Task loss (standard cross-entropy with true labels)
    loss_ce = F.cross_entropy(
        student_logits_flat,
        labels_flat,
        ignore_index=-100  # Ignore padding tokens
    )
    
    # Distillation loss (KL divergence with temperature scaling)
    student_soft = F.log_softmax(student_logits_flat / temperature, dim=-1)
    teacher_soft = F.softmax(teacher_logits_flat / temperature, dim=-1)
    
    loss_kd = F.kl_div(
        student_soft,
        teacher_soft,
        reduction='batchmean'
    ) * (temperature ** 2)  # Compensate for temperature
    
    # Combined loss
    loss = alpha * loss_kd + (1 - alpha) * loss_ce
    
    return loss, {"loss_kd": loss_kd.item(), "loss_ce": loss_ce.item()}

from torch.utils.data import DataLoader
from transformers import get_cosine_schedule_with_warmup
 
# Setup
device = "cuda"
teacher = teacher.to(device)
student = student.to(device)
 
optimizer = torch.optim.AdamW(student.parameters(), lr=2e-4, weight_decay=0.01)
scheduler = get_cosine_schedule_with_warmup(
    optimizer,
    num_warmup_steps=1000,
    num_training_steps=100000
)
 
# Training
student.train()
for step, batch in enumerate(dataloader):
    input_ids = batch["input_ids"].to(device)
    labels = batch["labels"].to(device)
    
    try:
        # Teacher forward (no gradients)
        with torch.no_grad():
            teacher_outputs = teacher(input_ids)
            teacher_logits = teacher_outputs.logits
        
        # Student forward
        student_outputs = student(input_ids)
        student_logits = student_outputs.logits
        
        # Check for shape mismatch (common distillation error)
        if teacher_logits.shape != student_logits.shape:
            print(f"Shape mismatch at step {step}: teacher {teacher_logits.shape} vs student {student_logits.shape}")
            optimizer.zero_grad()
            continue
        
        # Compute distillation loss
        loss, metrics = distillation_loss(
            student_logits,
            teacher_logits,
            labels,
            temperature=2.0,
            alpha=0.7
        )
        
        # Check for NaN loss (indicates numerical instability)
        if torch.isnan(loss):
            print(f"NaN loss detected at step {step}. Skipping batch.")
            optimizer.zero_grad()
            continue
        
        # Backward pass
        optimizer.zero_grad()
        loss.backward()
        torch.nn.utils.clip_grad_norm_(student.parameters(), 1.0)
        optimizer.step()
        scheduler.step()
        
        if step % 100 == 0:
            print(f"Step {step}: loss={loss.item():.4f}, "
                  f"kd={metrics['loss_kd']:.4f}, ce={metrics['loss_ce']:.4f}")
    
    except RuntimeError as e:
        if "out of memory" in str(e):
            print(f"OOM at step {step}. Clearing cache and skipping batch.")
            torch.cuda.empty_cache()
            optimizer.zero_grad()
            continue
        else:
            raise e

class FeatureDistillationLoss(torch.nn.Module):
    def __init__(self, student_dim, teacher_dim):
        super().__init__()
        # Project teacher features to student dimension
        self.projection = torch.nn.Linear(teacher_dim, student_dim)
        
    def forward(self, student_hidden, teacher_hidden):
        """
        Args:
            student_hidden: [batch, seq_len, student_dim]
            teacher_hidden: [batch, seq_len, teacher_dim]
        
        Returns:
            MSE loss between projected features
        """
        teacher_proj = self.projection(teacher_hidden)
        return F.mse_loss(student_hidden, teacher_proj)
 
# Usage: Add to distillation training
feature_loss_fn = FeatureDistillationLoss(
    student_dim=1536,
    teacher_dim=4096
).to(device)
 
# In training loop, add feature matching
student_hidden = student_outputs.hidden_states[8]  # Middle layer
teacher_hidden = teacher_outputs.hidden_states[16]  # Corresponding teacher layer
 
loss_features = feature_loss_fn(student_hidden, teacher_hidden)
total_loss = loss + 0.1 * loss_features  # Add feature loss with small weight

import torch.quantization as quantization
 
def quantize_model_int8(model, calibration_dataloader):
    """
    Quantize model to INT8 using dynamic quantization.
    
    Args:
        model: PyTorch model to quantize
        calibration_dataloader: Small dataset for calibration
    
    Returns:
        Quantized model
    """
    # Prepare model for quantization
    model.eval()
    model.qconfig = quantization.get_default_qconfig('fbgemm')  # x86 backend
    
    # Fuse operations (e.g., Conv+ReLU) for better performance
    model_fused = quantization.fuse_modules(model, [['conv', 'relu']])
    
    # Prepare for quantization (insert observers)
    model_prepared = quantization.prepare(model_fused)
    
    # Calibrate on representative data
    with torch.no_grad():
        for batch in calibration_dataloader:
            model_prepared(batch)
    
    # Convert to quantized model
    model_quantized = quantization.convert(model_prepared)
    
    return model_quantized
 
# Usage
model_int8 = quantize_model_int8(model, calibration_loader)
 
# Compare sizes
print(f"FP16 size: {get_model_size(model) / 1e9:.2f} GB")
print(f"INT8 size: {get_model_size(model_int8) / 1e9:.2f} GB")
# FP16 size: 13.5 GB
# INT8 size: 6.8 GB (2× reduction)

from auto_gptq import AutoGPTQForCausalLM, BaseQuantizeConfig
 
def quantize_model_int4_gptq(model_name, calibration_dataset, bits=4):
    """
    Quantize to INT4 using GPTQ algorithm.
    
    Args:
        model_name: HuggingFace model identifier
        calibration_dataset: Dataset for calibration (e.g., C4, WikiText)
        bits: Target bit-width (4 or 8)
    
    Returns:
        Quantized model
    """
    # Configure quantization
    quantize_config = BaseQuantizeConfig(
        bits=bits,                    # 4-bit quantization
        group_size=128,               # Quantize in groups of 128
        desc_act=False,               # Disable activation order
        damp_percent=0.01,            # Dampening for numerical stability
    )
    
    # Load model and quantize
    model = AutoGPTQForCausalLM.from_pretrained(
        model_name,
        quantize_config=quantize_config
    )
    
    model.quantize(calibration_dataset)
    
    # Save quantized model
    model.save_quantized("./model-gptq-4bit")
    
    return model
 
# Usage
from datasets import load_dataset
 
calibration_data = load_dataset("allenai/c4", "en", split="train[:1000]")
model_int4 = quantize_model_int4_gptq(
    "meta-llama/Llama-2-7b-hf",
    calibration_data,
    bits=4
)
 
# Size comparison
# FP16: 13.5 GB
# INT4: 3.5 GB (4× reduction!)

class QuantizedLinear(torch.nn.Module):
    """
    Linear layer with quantization-aware training.
    Uses fake quantization during training, real quantization during inference.
    """
    def __init__(self, in_features, out_features, bits=8):
        super().__init__()
        self.weight = torch.nn.Parameter(torch.randn(out_features, in_features))
        self.bias = torch.nn.Parameter(torch.zeros(out_features))
        
        self.bits = bits
        self.register_buffer('scale', torch.tensor(1.0))
        self.register_buffer('zero_point', torch.tensor(0))
        
    def forward(self, x):
        if self.training:
            # Fake quantization (differentiable)
            w_quant = self.fake_quantize(self.weight)
        else:
            # Real quantization (inference)
            w_quant = self.quantize(self.weight)
        
        return F.linear(x, w_quant, self.bias)
    
    def fake_quantize(self, x):
        """Simulate quantization during training."""
        qmin = -(2 ** (self.bits - 1))
        qmax = 2 ** (self.bits - 1) - 1
        
        # Update scale
        self.scale = x.abs().max() / (2 ** (self.bits - 1) - 1)
        
        # Quantize and dequantize
        x_quant = torch.round(x / self.scale).clamp(qmin, qmax)
        return x_quant * self.scale
    
    def quantize(self, x):
        """Real quantization for inference."""
        return torch.round(x / self.scale).clamp(
            -(2 ** (self.bits - 1)),
            2 ** (self.bits - 1) - 1
        ) * self.scale
 
# Replace all Linear layers with QuantizedLinear
def convert_to_qat(model, bits=8):
    for name, module in model.named_children():
        if isinstance(module, torch.nn.Linear):
            setattr(model, name, QuantizedLinear(
                module.in_features,
                module.out_features,
                bits=bits
            ))
        else:
            convert_to_qat(module, bits)
    return model
 
# Usage
model_qat = convert_to_qat(model, bits=8)
# Train normally - quantization is baked in!

def prune_magnitude(model, sparsity=0.5, structured=False):
    """
    Prune model by magnitude.
    
    Args:
        model: PyTorch model
        sparsity: Fraction of parameters to remove (0.5 = 50%)
        structured: If True, prune entire neurons; if False, individual weights
    
    Returns:
        Model with pruning masks applied
    """
    import torch.nn.utils.prune as prune
    
    for name, module in model.named_modules():
        if isinstance(module, torch.nn.Linear):
            if structured:
                # Prune entire output neurons
                prune.ln_structured(
                    module,
                    name='weight',
                    amount=sparsity,
                    n=2,         # L2 norm
                    dim=0        # Prune rows (output neurons)
                )
            else:
                # Prune individual weights
                prune.l1_unstructured(
                    module,
                    name='weight',
                    amount=sparsity
                )
    
    return model
 
# Usage
model_pruned = prune_magnitude(model, sparsity=0.5, structured=True)
 
# Make pruning permanent
for module in model_pruned.modules():
    if isinstance(module, torch.nn.Linear):
        prune.remove(module, 'weight')

def iterative_magnitude_pruning(
    model,
    train_fn,
    target_sparsity=0.9,
    num_iterations=5
):
    """
    Iterative magnitude pruning with retraining.
    
    Based on "The Lottery Ticket Hypothesis" (Frankle & Carbin, 2019)
    
    Args:
        model: Model to prune
        train_fn: Function that trains model for one cycle
        target_sparsity: Final sparsity target
        num_iterations: Number of prune-retrain cycles
    
    Returns:
        Pruned model
    """
    # Save initial weights
    initial_state = {name: param.clone() for name, param in model.named_parameters()}
    
    # Current sparsity starts at 0
    current_sparsity = 0
    
    for iteration in range(num_iterations):
        print(f"\n=== Iteration {iteration + 1}/{num_iterations} ===")
        
        # Calculate sparsity for this iteration
        # Use exponential schedule: gradually increase sparsity
        current_sparsity = target_sparsity * (1 - (1 - (iteration + 1) / num_iterations) ** 3)
        
        # Prune by current sparsity
        model = prune_magnitude(model, sparsity=current_sparsity, structured=False)
        
        # Reset to initial weights (lottery ticket insight!)
        for name, param in model.named_parameters():
            if 'weight' in name:
                # Get mask
                mask = (param.data != 0).float()
                # Reset to initial weights with mask
                param.data = initial_state[name] * mask
        
        # Train for one cycle
        print(f"Training with sparsity={current_sparsity:.1%}")
        model = train_fn(model)
        
        # Evaluate
        eval_results = evaluate(model)
        print(f"Sparsity: {current_sparsity:.1%}, Accuracy: {eval_results['accuracy']:.2%}")
    
    return model

def prune_attention_heads(model, num_heads_to_prune=4):
    """
    Prune least important attention heads.
    
    Importance measured by average attention weight magnitude.
    """
    # Collect attention head importance scores
    head_importance = {}
    
    model.eval()
    with torch.no_grad():
        for name, module in model.named_modules():
            if 'self_attn' in name:
                # Run forward pass and collect attention weights
                # (Implementation depends on model architecture)
                importance = module.weight.abs().mean(dim=(1, 2))
                head_importance[name] = importance
    
    # Find least important heads
    all_importance = torch.cat(list(head_importance.values()))
    threshold = torch.kthvalue(all_importance, num_heads_to_prune).values
    
    # Prune heads below threshold
    for name, module in model.named_modules():
        if name in head_importance:
            heads_to_keep = head_importance[name] > threshold
            # Modify attention module to keep only important heads
            # (Architecture-specific implementation)
    
    return model

class LoRALinear(torch.nn.Module):
    """
    Linear layer with LoRA adaptation.
    
    Replaces: y = W x
    With: y = W_0 x + (B A) x
    
    Where W_0 is frozen, B and A are trainable.
    """
    def __init__(
        self,
        in_features,
        out_features,
        rank=16,
        alpha=16,
        dropout=0.1
    ):
        super().__init__()
        
        # Frozen base weights (will be loaded from pretrained)
        self.base_layer = torch.nn.Linear(in_features, out_features, bias=True)
        self.base_layer.weight.requires_grad = False
        
        # LoRA parameters
        self.lora_A = torch.nn.Parameter(torch.randn(rank, in_features) * 0.01)
        self.lora_B = torch.nn.Parameter(torch.zeros(out_features, rank))
        
        self.rank = rank
        self.alpha = alpha
        self.scaling = alpha / rank
        
        self.dropout = torch.nn.Dropout(dropout)
        
    def forward(self, x):
        # Base output (frozen)
        base_out = self.base_layer(x)
        
        # LoRA adaptation: x @ A.T @ B.T = (x @ A.T) @ B.T
        lora_out = self.dropout(x) @ self.lora_A.T @ self.lora_B.T
        
        return base_out + self.scaling * lora_out
    
    def merge_weights(self):
        """Merge LoRA weights into base layer for inference."""
        if self.rank > 0:
            self.base_layer.weight.data += (
                self.scaling * (self.lora_B @ self.lora_A)
            )
            # Zero out LoRA params
            self.lora_A = None
            self.lora_B = None
 
# Replace Linear layers with LoRA versions
def add_lora_to_model(model, rank=16, alpha=16, target_modules=None):
    """
    Add LoRA adapters to specified modules.
    
    Args:
        model: Base model
        rank: LoRA rank
        alpha: LoRA alpha (scaling factor)
        target_modules: List of module names to add LoRA to
                       (e.g., ['q_proj', 'v_proj', 'k_proj', 'o_proj'])
    
    Returns:
        Model with LoRA adapters
    """
    if target_modules is None:
        target_modules = ['q_proj', 'v_proj']  # Attention projections
    
    for name, module in model.named_modules():
        # Check if this module should get LoRA
        should_add_lora = any(target in name for target in target_modules)
        
        if should_add_lora and isinstance(module, torch.nn.Linear):
            # Replace with LoRA version
            parent_name = '.'.join(name.split('.')[:-1])
            child_name = name.split('.')[-1]
            
            parent = model.get_submodule(parent_name) if parent_name else model
            
            lora_layer = LoRALinear(
                module.in_features,
                module.out_features,
                rank=rank,
                alpha=alpha
            )
            
            # Copy base weights
            lora_layer.base_layer.weight.data = module.weight.data.clone()
            if module.bias is not None:
                lora_layer.base_layer.bias.data = module.bias.data.clone()
            
            setattr(parent, child_name, lora_layer)
    
    return model
 
# Usage
from transformers import AutoModelForCausalLM
 
model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-2-7b-hf")
model_lora = add_lora_to_model(
    model,
    rank=16,
    alpha=16,
    target_modules=['q_proj', 'v_proj', 'k_proj', 'o_proj']
)
 
# Count trainable parameters
total_params = sum(p.numel() for p in model_lora.parameters())
trainable_params = sum(p.numel() for p in model_lora.parameters() if p.requires_grad)
 
print(f"Total parameters: {total_params / 1e9:.2f}B")
print(f"Trainable parameters: {trainable_params / 1e6:.2f}M")
print(f"Trainable %: {100 * trainable_params / total_params:.4f}%")
# Total parameters: 6.74B
# Trainable parameters: 4.19M
# Trainable %: 0.0622%

from transformers import AutoModelForCausalLM, BitsAndBytesConfig
from peft import LoraConfig, get_peft_model
 
def create_qlora_model(model_name, rank=64, alpha=16):
    """
    Create model with 4-bit quantization + LoRA.
    
    Enables fine-tuning 65B models on single 48GB GPU!
    """
    # 4-bit quantization config
    bnb_config = BitsAndBytesConfig(
        load_in_4bit=True,
        bnb_4bit_quant_type="nf4",        # NormalFloat4
        bnb_4bit_compute_dtype=torch.bfloat16,
        bnb_4bit_use_double_quant=True,   # Nested quantization
    )
    
    # Load base model in 4-bit
    model = AutoModelForCausalLM.from_pretrained(
        model_name,
        quantization_config=bnb_config,
        device_map="auto",
        trust_remote_code=True
    )
    
    # Add LoRA adapters
    lora_config = LoraConfig(
        r=rank,
        lora_alpha=alpha,
        target_modules=["q_proj", "k_proj", "v_proj", "o_proj"],
        lora_dropout=0.05,
        bias="none",
        task_type="CAUSAL_LM"
    )
    
    model = get_peft_model(model, lora_config)
    model.print_trainable_parameters()
    
    return model
 
# Usage
model_qlora = create_qlora_model("meta-llama/Llama-2-7b-hf", rank=64)
# trainable params: 4,194,304 || all params: 6,738,415,616 || trainable%: 0.0622
 
# Train normally - base model stays in INT4, adapters in BF16!

# Pseudo-code
teacher = load_model("Llama-13B")  # 26 GB
student = distill(teacher, target_size=7B)  # → 14 GB
student_int8 = quantize(student, bits=8)  # → 7 GB
student_finetuned = finetune_lora(student_int8, task_data)  # Same size
# Final: 7 GB, 94% of original quality

teacher = load_model("Llama-7B")  # 14 GB
student = distill(teacher, target_size=3B)  # → 6 GB
student_int4 = quantize_gptq(student, bits=4)  # → 1.5 GB
# Final: 1.5 GB, runs on consumer GPUs at 40+ tok/s

teacher = load_model("Llama-7B")  # 14 GB
student = distill(teacher, target_size=1.5B)  # → 3 GB
student_pruned = prune_structured(student, sparsity=0.33)  # → 2 GB
student_int4 = quantize_gptq(student_pruned, bits=4)  # → 500 MB
# Final: 500 MB, runs on phones!