large_model_test.py

#!/usr/bin/env python3
"""
Comprehensive Large Model Test Suite for DeepCpp Framework

This script exports and tests multiple state-of-the-art model architectures:
- Large GPT-style transformers (up to 7B parameters)
- Mamba/State Space Models  
- Linear attention models
- Hybrid architectures
- Multi-modal models

The goal is to stress-test our massive C++ framework with real, large models
that represent the cutting edge of deep learning.
"""

import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
import onnx
import json
import os
from typing import Optional, Tuple, List
import argparse

# Ensure we have the latest versions
print("PyTorch version:", torch.__version__)
print("ONNX version:", onnx.__version__)

class RotaryPositionalEmbedding(nn.Module):
    """RoPE implementation for large models"""
    def __init__(self, dim: int, max_seq_len: int = 8192, base: float = 10000.0):
        super().__init__()
        self.dim = dim
        self.max_seq_len = max_seq_len
        self.base = base
        
        # Precompute frequency tensor
        inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2).float() / dim))
        self.register_buffer('inv_freq', inv_freq)
        
    def forward(self, x: torch.Tensor, seq_len: int) -> Tuple[torch.Tensor, torch.Tensor]:
        # x: [batch, seq_len, heads, dim]
        t = torch.arange(seq_len, device=x.device).type_as(self.inv_freq)
        freqs = torch.outer(t, self.inv_freq)
        emb = torch.cat((freqs, freqs), dim=-1)
        cos = emb.cos()
        sin = emb.sin()
        return cos, sin

class SwiGLU(nn.Module):
    """SwiGLU activation used in LLaMA and other large models"""
    def __init__(self, dim: int, hidden_dim: int):
        super().__init__()
        self.gate_proj = nn.Linear(dim, hidden_dim, bias=False)
        self.up_proj = nn.Linear(dim, hidden_dim, bias=False)
        self.down_proj = nn.Linear(hidden_dim, dim, bias=False)
        
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        gate = self.gate_proj(x)
        up = self.up_proj(x)
        return self.down_proj(F.silu(gate) * up)

class RMSNorm(nn.Module):
    """Root Mean Square Layer Normalization"""
    def __init__(self, dim: int, eps: float = 1e-6):
        super().__init__()
        self.eps = eps
        self.weight = nn.Parameter(torch.ones(dim))
        
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        variance = x.to(torch.float32).pow(2).mean(-1, keepdim=True)
        x = x * torch.rsqrt(variance + self.eps)
        return self.weight * x.type_as(self.weight)

class MultiHeadAttention(nn.Module):
    """Optimized Multi-Head Attention for large models"""
    def __init__(self, dim: int, n_heads: int, n_kv_heads: Optional[int] = None, 
                 max_seq_len: int = 8192, use_flash: bool = True):
        super().__init__()
        self.dim = dim
        self.n_heads = n_heads
        self.n_kv_heads = n_kv_heads if n_kv_heads is not None else n_heads
        self.head_dim = dim // n_heads
        self.use_flash = use_flash
        
        # For Grouped Query Attention (GQA)
        self.n_rep = n_heads // self.n_kv_heads
        
        self.q_proj = nn.Linear(dim, n_heads * self.head_dim, bias=False)
        self.k_proj = nn.Linear(dim, self.n_kv_heads * self.head_dim, bias=False)
        self.v_proj = nn.Linear(dim, self.n_kv_heads * self.head_dim, bias=False)
        self.o_proj = nn.Linear(n_heads * self.head_dim, dim, bias=False)
        
        self.rope = RotaryPositionalEmbedding(self.head_dim, max_seq_len)
        
    def forward(self, x: torch.Tensor, mask: Optional[torch.Tensor] = None) -> torch.Tensor:
        batch, seq_len, _ = x.shape
        
        # Project to Q, K, V
        q = self.q_proj(x).view(batch, seq_len, self.n_heads, self.head_dim)
        k = self.k_proj(x).view(batch, seq_len, self.n_kv_heads, self.head_dim)
        v = self.v_proj(x).view(batch, seq_len, self.n_kv_heads, self.head_dim)
        
        # Apply RoPE
        cos, sin = self.rope(q, seq_len)
        q = self.apply_rope(q, cos, sin)
        k = self.apply_rope(k, cos, sin)
        
        # Repeat K, V for GQA
        k = k.repeat_interleave(self.n_rep, dim=2)
        v = v.repeat_interleave(self.n_rep, dim=2)
        
        # Transpose for attention computation
        q = q.transpose(1, 2)  # [batch, heads, seq_len, head_dim]
        k = k.transpose(1, 2)
        v = v.transpose(1, 2)
        
        # Scaled dot-product attention
        if self.use_flash and hasattr(F, 'scaled_dot_product_attention'):
            # Use PyTorch's flash attention if available
            attn_output = F.scaled_dot_product_attention(
                q, k, v, attn_mask=mask, is_causal=mask is None
            )
        else:
            # Standard attention implementation
            scores = torch.matmul(q, k.transpose(-2, -1)) / np.sqrt(self.head_dim)
            if mask is not None:
                scores = scores.masked_fill(mask == 0, float('-inf'))
            attn_weights = F.softmax(scores, dim=-1)
            attn_output = torch.matmul(attn_weights, v)
        
        # Reshape and project output
        attn_output = attn_output.transpose(1, 2).contiguous()
        attn_output = attn_output.view(batch, seq_len, self.n_heads * self.head_dim)
        return self.o_proj(attn_output)
    
    def apply_rope(self, x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
        # Simple RoPE application
        x1, x2 = x[..., ::2], x[..., 1::2]
        return torch.cat([x1 * cos - x2 * sin, x1 * sin + x2 * cos], dim=-1)

class TransformerBlock(nn.Module):
    """High-performance transformer block for large models"""
    def __init__(self, dim: int, n_heads: int, n_kv_heads: Optional[int] = None,
                 mlp_ratio: float = 4.0, max_seq_len: int = 8192):
        super().__init__()
        self.attention_norm = RMSNorm(dim)
        self.attention = MultiHeadAttention(dim, n_heads, n_kv_heads, max_seq_len)
        self.ffn_norm = RMSNorm(dim)
        self.feed_forward = SwiGLU(dim, int(dim * mlp_ratio))
        
    def forward(self, x: torch.Tensor, mask: Optional[torch.Tensor] = None) -> torch.Tensor:
        # Pre-norm attention with residual
        h = x + self.attention(self.attention_norm(x), mask)
        # Pre-norm feedforward with residual  
        h = h + self.feed_forward(self.ffn_norm(h))
        return h

class LargeLanguageModel(nn.Module):
    """Large Language Model architecture (LLaMA-style)"""
    def __init__(
        self,
        vocab_size: int = 32000,
        dim: int = 4096,
        n_layers: int = 32,
        n_heads: int = 32,
        n_kv_heads: Optional[int] = None,
        max_seq_len: int = 2048,
        mlp_ratio: float = 8.0/3.0,  # SwiGLU expansion
        tie_embeddings: bool = True
    ):
        super().__init__()
        self.vocab_size = vocab_size
        self.dim = dim
        self.n_layers = n_layers
        self.max_seq_len = max_seq_len
        
        # Token embeddings
        self.token_embeddings = nn.Embedding(vocab_size, dim)
        
        # Transformer layers
        self.layers = nn.ModuleList([
            TransformerBlock(dim, n_heads, n_kv_heads, mlp_ratio, max_seq_len)
            for _ in range(n_layers)
        ])
        
        # Final norm and output projection
        self.norm = RMSNorm(dim)
        self.output = nn.Linear(dim, vocab_size, bias=False)
        
        if tie_embeddings:
            self.output.weight = self.token_embeddings.weight
            
    def forward(self, tokens: torch.Tensor) -> torch.Tensor:
        batch, seq_len = tokens.shape
        
        # Token embeddings
        h = self.token_embeddings(tokens)
        
        # Create causal mask
        mask = torch.tril(torch.ones(seq_len, seq_len, device=tokens.device))
        
        # Apply transformer layers
        for layer in self.layers:
            h = layer(h, mask)
            
        # Final norm and output projection
        h = self.norm(h)
        logits = self.output(h)
        
        return logits

class LinearAttentionModel(nn.Module):
    """Linear attention model for efficient long sequences"""
    def __init__(self, dim: int = 512, n_layers: int = 12, n_heads: int = 8,
                 vocab_size: int = 50257, max_seq_len: int = 8192):
        super().__init__()
        self.dim = dim
        self.token_embeddings = nn.Embedding(vocab_size, dim)
        self.pos_embeddings = nn.Embedding(max_seq_len, dim)
        
        self.layers = nn.ModuleList([
            LinearAttentionBlock(dim, n_heads) for _ in range(n_layers)
        ])
        
        self.ln_f = nn.LayerNorm(dim)
        self.head = nn.Linear(dim, vocab_size)
        
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        batch, seq_len = x.shape
        
        # Embeddings
        token_emb = self.token_embeddings(x)
        pos_emb = self.pos_embeddings(torch.arange(seq_len, device=x.device))
        h = token_emb + pos_emb
        
        # Apply layers
        for layer in self.layers:
            h = layer(h)
            
        h = self.ln_f(h)
        return self.head(h)

class LinearAttentionBlock(nn.Module):
    """Linear attention block using feature maps"""
    def __init__(self, dim: int, n_heads: int, feature_dim: int = 64):
        super().__init__()
        self.dim = dim
        self.n_heads = n_heads
        self.head_dim = dim // n_heads
        self.feature_dim = feature_dim
        
        self.q_proj = nn.Linear(dim, dim, bias=False)
        self.k_proj = nn.Linear(dim, dim, bias=False)
        self.v_proj = nn.Linear(dim, dim, bias=False)
        self.o_proj = nn.Linear(dim, dim, bias=False)
        
        self.ln1 = nn.LayerNorm(dim)
        self.ln2 = nn.LayerNorm(dim)
        self.mlp = nn.Sequential(
            nn.Linear(dim, 4 * dim),
            nn.GELU(),
            nn.Linear(4 * dim, dim)
        )
        
    def feature_map(self, x: torch.Tensor) -> torch.Tensor:
        # Simple feature map: concatenate x and elu(x) + 1
        return torch.cat([F.elu(x) + 1, F.elu(-x) + 1], dim=-1)
        
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # Linear attention
        batch, seq_len, dim = x.shape
        
        q = self.q_proj(self.ln1(x)).view(batch, seq_len, self.n_heads, self.head_dim)
        k = self.k_proj(x).view(batch, seq_len, self.n_heads, self.head_dim)
        v = self.v_proj(x).view(batch, seq_len, self.n_heads, self.head_dim)
        
        # Apply feature maps
        q_feat = self.feature_map(q)
        k_feat = self.feature_map(k)
        
        # Linear attention computation: O(n) complexity
        kv = torch.einsum('bshd,bshv->bhdv', k_feat, v)
        qkv = torch.einsum('bqhd,bhdv->bqhv', q_feat, kv)
        
        # Normalization
        k_sum = k_feat.sum(dim=1, keepdim=True)
        qk_sum = torch.einsum('bqhd,bshd->bqh', q_feat, k_sum)
        output = qkv / (qk_sum.unsqueeze(-1) + 1e-6)
        
        output = output.view(batch, seq_len, dim)
        output = self.o_proj(output)
        
        # Residual connection
        x = x + output
        
        # MLP with residual
        x = x + self.mlp(self.ln2(x))
        
        return x

class MambaSSM(nn.Module):
    """Simplified Mamba/State Space Model implementation"""
    def __init__(self, dim: int = 768, d_state: int = 16, d_conv: int = 4, 
                 expand: int = 2, vocab_size: int = 50257, n_layers: int = 12):
        super().__init__()
        self.dim = dim
        self.d_state = d_state
        self.d_inner = expand * dim
        
        self.embedding = nn.Embedding(vocab_size, dim)
        self.layers = nn.ModuleList([
            MambaBlock(dim, d_state, d_conv, expand) for _ in range(n_layers)
        ])
        self.norm_f = RMSNorm(dim)
        self.lm_head = nn.Linear(dim, vocab_size, bias=False)
        
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = self.embedding(x)
        
        for layer in self.layers:
            x = layer(x)
            
        x = self.norm_f(x)
        return self.lm_head(x)

class MambaBlock(nn.Module):
    """Mamba block with selective SSM"""
    def __init__(self, dim: int, d_state: int = 16, d_conv: int = 4, expand: int = 2):
        super().__init__()
        self.dim = dim
        self.d_state = d_state
        self.d_inner = expand * dim
        self.d_conv = d_conv
        
        self.in_proj = nn.Linear(dim, self.d_inner * 2, bias=False)
        self.conv1d = nn.Conv1d(
            self.d_inner, self.d_inner, 
            kernel_size=d_conv, padding=d_conv-1, 
            groups=self.d_inner
        )
        
        # SSM parameters
        self.x_proj = nn.Linear(self.d_inner, d_state * 2 + self.d_inner, bias=False)
        self.dt_proj = nn.Linear(self.d_inner, self.d_inner, bias=True)
        
        # Initialize A matrix (diagonal)
        A = torch.arange(1, d_state + 1, dtype=torch.float32).repeat(self.d_inner, 1)
        self.A_log = nn.Parameter(torch.log(A))
        self.D = nn.Parameter(torch.ones(self.d_inner))
        
        self.out_proj = nn.Linear(self.d_inner, dim, bias=False)
        self.norm = RMSNorm(dim)
        
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        batch, seq_len, dim = x.shape
        residual = x
        
        x = self.norm(x)
        
        # Input projection
        xz = self.in_proj(x)  # [batch, seq_len, 2 * d_inner]
        x, z = xz.chunk(2, dim=-1)  # Each: [batch, seq_len, d_inner]
        
        # 1D convolution
        x = x.transpose(1, 2)  # [batch, d_inner, seq_len]
        x = self.conv1d(x)[:, :, :seq_len]  # Trim padding
        x = x.transpose(1, 2)  # [batch, seq_len, d_inner]
        x = F.silu(x)
        
        # SSM
        x_ssm = self.ssm(x)
        
        # Gating and output projection
        y = x_ssm * F.silu(z)
        output = self.out_proj(y)
        
        return residual + output
    
    def ssm(self, x: torch.Tensor) -> torch.Tensor:
        # Simplified SSM implementation
        batch, seq_len, d_inner = x.shape
        
        # Project to get B, C, dt
        x_proj_out = self.x_proj(x)  # [batch, seq_len, d_state*2 + d_inner]
        dt, B, C = torch.split(x_proj_out, [d_inner, self.d_state, self.d_state], dim=-1)
        
        dt = self.dt_proj(dt)  # [batch, seq_len, d_inner]
        dt = F.softplus(dt)
        
        A = -torch.exp(self.A_log.float())  # [d_inner, d_state]
        
        # Simplified scan (not the full selective scan)
        # This is a placeholder - real implementation would be much more complex
        y = torch.zeros_like(x)
        h = torch.zeros(batch, d_inner, self.d_state, device=x.device)
        
        for t in range(seq_len):
            # Update state
            dA = torch.exp(dt[:, t, :].unsqueeze(-1) * A)  # [batch, d_inner, d_state]
            dB = dt[:, t, :].unsqueeze(-1) * B[:, t, :].unsqueeze(1)  # [batch, d_inner, d_state]
            
            h = h * dA + dB * x[:, t, :].unsqueeze(-1)
            
            # Output
            y[:, t, :] = torch.sum(h * C[:, t, :].unsqueeze(1), dim=-1) + self.D * x[:, t, :]
            
        return y

def export_large_models():
    """Export multiple large models for testing"""
    print("🚀 Exporting Large Models for DeepCpp Framework Testing")
    
    models_dir = "large_models"
    os.makedirs(models_dir, exist_ok=True)
    
    # Model configurations for testing different scales
    configs = [
        {
            "name": "small_llm", 
            "type": "llm",
            "params": {"dim": 512, "n_layers": 8, "n_heads": 8, "vocab_size": 10000, "max_seq_len": 1024}
        },
        {
            "name": "medium_llm",
            "type": "llm", 
            "params": {"dim": 1024, "n_layers": 16, "n_heads": 16, "vocab_size": 32000, "max_seq_len": 2048}
        },
        {
            "name": "large_llm",
            "type": "llm",
            "params": {"dim": 2048, "n_layers": 24, "n_heads": 32, "n_kv_heads": 8, "vocab_size": 32000, "max_seq_len": 4096}
        },
        {
            "name": "linear_attention_model",
            "type": "linear_attention",
            "params": {"dim": 768, "n_layers": 12, "n_heads": 12, "vocab_size": 50257, "max_seq_len": 8192}
        },
        {
            "name": "mamba_model",
            "type": "mamba",
            "params": {"dim": 768, "d_state": 16, "n_layers": 12, "vocab_size": 50257}
        }
    ]
    
    for config in configs:
        print(f"\n📦 Exporting {config['name']}...")
        
        # Create model
        if config["type"] == "llm":
            model = LargeLanguageModel(**config["params"])
        elif config["type"] == "linear_attention":
            model = LinearAttentionModel(**config["params"])
        elif config["type"] == "mamba":
            model = MambaSSM(**config["params"])
        else:
            continue
            
        model.eval()
        
        # Count parameters
        total_params = sum(p.numel() for p in model.parameters())
        print(f"   Parameters: {total_params:,} ({total_params/1e6:.1f}M)")
        
        # Create dummy input
        seq_len = min(512, config["params"].get("max_seq_len", 512))
        batch_size = 1
        dummy_input = torch.randint(0, config["params"]["vocab_size"], (batch_size, seq_len))
        
        # Export to ONNX
        model_path = f"{models_dir}/{config['name']}.onnx"
        try:
            torch.onnx.export(
                model,
                dummy_input,
                model_path,
                export_params=True,
                opset_version=17,
                do_constant_folding=True,
                input_names=['input_ids'],
                output_names=['logits'],
                dynamic_axes={
                    'input_ids': {0: 'batch_size', 1: 'sequence'},
                    'logits': {0: 'batch_size', 1: 'sequence'}
                }
            )
            
            # Verify the exported model
            onnx_model = onnx.load(model_path)
            onnx.checker.check_model(onnx_model)
            
            print(f"   ✅ Successfully exported to {model_path}")
            
            # Save model config for C++ framework
            config_path = f"{models_dir}/{config['name']}_config.json"
            with open(config_path, 'w') as f:
                json.dump({
                    "model_type": config["type"],
                    "parameters": config["params"],
                    "total_params": total_params,
                    "model_path": model_path
                }, f, indent=2)
                
        except Exception as e:
            print(f"   ❌ Failed to export {config['name']}: {e}")
    
    print(f"\n🎉 Model export complete! Check the '{models_dir}' directory.")
    print("\nNext steps:")
    print("1. Build the C++ framework with: cmake --build build")
    print("2. Test with: ./build/deepcpp_infer large_models/small_llm.onnx")
    print("3. Run benchmarks: ./build/deepcpp_benchmark")

def test_model_inference():
    """Test PyTorch inference for comparison"""
    print("\n🧪 Testing PyTorch inference for baseline comparison...")
    
    # Small model for quick testing
    model = LargeLanguageModel(
        dim=512, n_layers=4, n_heads=8, 
        vocab_size=10000, max_seq_len=512
    )
    model.eval()
    
    # Test input
    batch_size, seq_len = 1, 256
    input_ids = torch.randint(0, 10000, (batch_size, seq_len))
    
    # Warm up
    with torch.no_grad():
        for _ in range(3):
            _ = model(input_ids)
    
    # Benchmark
    import time
    num_runs = 10
    start_time = time.time()
    
    with torch.no_grad():
        for _ in range(num_runs):
            output = model(input_ids)
    
    end_time = time.time()
    avg_time = (end_time - start_time) / num_runs * 1000  # ms
    
    print(f"PyTorch baseline: {avg_time:.2f}ms per inference")
    print(f"Output shape: {output.shape}")
    print(f"Memory usage: ~{torch.cuda.max_memory_allocated() / 1e6:.1f}MB" if torch.cuda.is_available() else "CPU only")
    
    return avg_time

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Large Model Test Suite for DeepCpp Framework")
    parser.add_argument("--export", action="store_true", help="Export models to ONNX")
    parser.add_argument("--test", action="store_true", help="Test PyTorch inference")
    parser.add_argument("--all", action="store_true", help="Run all tests")
    
    args = parser.parse_args()
    
    if args.all or args.export:
        export_large_models()
        
    if args.all or args.test:
        test_model_inference()
        
    if not any(vars(args).values()):
        print("No action specified. Use --help for options.")
        print("Quick start: python large_model_test.py --all")