diff --git a/benchmark.py b/benchmark.py new file mode 100644 index 0000000..2ed4b8b --- /dev/null +++ b/benchmark.py @@ -0,0 +1,651 @@ +import logging +import time +import json +import sys +from pathlib import Path +from dataclasses import dataclass, asdict +from typing import List, Dict +import torch +import numpy as np +import matplotlib.pyplot as plt +import pandas as pd + +logging.basicConfig( + level=logging.INFO, + format='[%(asctime)s] %(levelname)s - %(name)s - %(message)s' +) +logger = logging.getLogger(__name__) + +try: + logger.info("Importing original Engram implementation") + from engram_demo_v1 import ( + Engram as EngramOriginal, + TransformerBlock as TransformerBlockOriginal, + EngramConfig as EngramConfigOriginal, + BackBoneConfig as BackBoneConfigOriginal, + ) + logger.info("Original Engram imported successfully") +except ImportError as e: + logger.error("Failed to import original Engram: %s", e) + sys.exit(1) + +try: + logger.info("Importing optimized Engram implementation") + from engram_optimized import ( + Engram as EngramOptimized, + TransformerBlock as TransformerBlockOptimized, + EngramConfig as EngramConfigOptimized, + BackBoneConfig as BackBoneConfigOptimized, + ) + logger.info("Optimized Engram imported successfully") +except ImportError as e: + logger.error("Failed to import optimized Engram: %s", e) + sys.exit(1) + +@dataclass +class BenchmarkResult: + model_type: str + layer_id: int + batch_size: int + seq_len: int + forward_time_ms: float + memory_allocated_mb: float + memory_reserved_mb: float + cache_hits: int = 0 + cache_misses: int = 0 + cache_hit_rate_pct: float = 0.0 + iterations: int = 100 + device: str = "cuda" + +class MemoryTracker: + """Track GPU memory usage""" + def __init__(self, device): + self.device = device + self.enabled = device.type == 'cuda' + logger.info("MemoryTracker initialized for device=%s, enabled=%s", + device, self.enabled) + + def reset(self): + if self.enabled: + torch.cuda.empty_cache() + torch.cuda.reset_peak_memory_stats(self.device) + logger.debug("GPU memory cache cleared and stats reset") + + def get_current_usage(self): + if self.enabled: + allocated = torch.cuda.memory_allocated(self.device) / 1024**2 + reserved = torch.cuda.memory_reserved(self.device) / 1024**2 + logger.debug("Current memory - allocated: %.2f MB, reserved: %.2f MB", + allocated, reserved) + return allocated, reserved + return 0.0, 0.0 + + def get_peak_usage(self): + if self.enabled: + peak_allocated = torch.cuda.max_memory_allocated(self.device) / 1024**2 + peak_reserved = torch.cuda.max_memory_reserved(self.device) / 1024**2 + logger.debug("Peak memory - allocated: %.2f MB, reserved: %.2f MB", + peak_allocated, peak_reserved) + return peak_allocated, peak_reserved + return 0.0, 0.0 + +def set_seed(seed=42): + """Set random seeds for reproducibility""" + logger.info("Setting random seed to %d", seed) + torch.manual_seed(seed) + torch.cuda.manual_seed_all(seed) + np.random.seed(seed) + +def create_test_data(batch_size, seq_len, vocab_size, device): + """Generate synthetic test data""" + logger.debug("Creating test data: batch_size=%d, seq_len=%d, vocab_size=%d", + batch_size, seq_len, vocab_size) + set_seed(42) + input_ids = torch.randint(0, vocab_size, (batch_size, seq_len), device=device) + return input_ids + +def create_hidden_states(batch_size, seq_len, hidden_size, hc_mult, device): + """Generate synthetic hidden states""" + logger.debug("Creating hidden states: batch_size=%d, seq_len=%d, hidden_size=%d, hc_mult=%d", + batch_size, seq_len, hidden_size, hc_mult) + set_seed(42) + hidden_states = torch.randn(batch_size, seq_len, hc_mult, hidden_size, device=device) + return hidden_states + +def warmup_model(model, input_ids, hidden_states, warmup_iterations=10): + """Warmup model before benchmarking""" + logger.info("Starting model warmup with %d iterations", warmup_iterations) + model.eval() + with torch.no_grad(): + for i in range(warmup_iterations): + _ = model(input_ids, hidden_states) + if i % 5 == 0: + logger.debug("Warmup iteration %d/%d", i+1, warmup_iterations) + + if torch.cuda.is_available(): + torch.cuda.synchronize() + logger.info("Model warmup complete") + +def benchmark_forward_pass(model, input_ids, hidden_states, iterations=100): + """Benchmark forward pass latency""" + logger.info("Starting benchmark with %d iterations", iterations) + + model.eval() + + warmup_model(model, input_ids, hidden_states, warmup_iterations=10) + + times = [] + + logger.info("Running timed iterations") + with torch.no_grad(): + for i in range(iterations): + if torch.cuda.is_available(): + torch.cuda.synchronize() + + start_time = time.time() + _ = model(input_ids, hidden_states) + + if torch.cuda.is_available(): + torch.cuda.synchronize() + + elapsed = (time.time() - start_time) * 1000 + times.append(elapsed) + + if (i + 1) % 20 == 0: + logger.debug("Completed %d/%d iterations, current avg: %.2fms", + i+1, iterations, np.mean(times)) + + mean_time = np.mean(times) + std_time = np.std(times) + median_time = np.median(times) + min_time = np.min(times) + max_time = np.max(times) + p95_time = np.percentile(times, 95) + p99_time = np.percentile(times, 99) + + logger.info("Benchmark statistics:") + logger.info(" Mean: %.2f ms", mean_time) + logger.info(" Std: %.2f ms", std_time) + logger.info(" Median: %.2f ms", median_time) + logger.info(" Min: %.2f ms", min_time) + logger.info(" Max: %.2f ms", max_time) + logger.info(" P95: %.2f ms", p95_time) + logger.info(" P99: %.2f ms", p99_time) + + return { + 'mean': mean_time, + 'std': std_time, + 'median': median_time, + 'min': min_time, + 'max': max_time, + 'p95': p95_time, + 'p99': p99_time, + 'all_times': times + } + +def benchmark_model(model, model_type, layer_id, batch_size, seq_len, + backbone_config, device, memory_tracker, iterations=100): + """Benchmark a single model configuration""" + logger.info("="*80) + logger.info("Benchmarking %s model", model_type) + logger.info(" Layer ID: %d", layer_id) + logger.info(" Batch size: %d", batch_size) + logger.info(" Sequence length: %d", seq_len) + logger.info("="*80) + + memory_tracker.reset() + + input_ids = create_test_data(batch_size, seq_len, backbone_config.vocab_size, device) + hidden_states = create_hidden_states( + batch_size, seq_len, + backbone_config.hidden_size, + backbone_config.hc_mult, + device + ) + + logger.info("Moving model to device: %s", device) + model = model.to(device) + model.eval() + + logger.info("Measuring initial memory footprint") + initial_alloc, initial_reserved = memory_tracker.get_current_usage() + logger.info("Initial memory - allocated: %.2f MB, reserved: %.2f MB", + initial_alloc, initial_reserved) + + logger.info("Starting forward pass benchmark") + timing_stats = benchmark_forward_pass(model, input_ids, hidden_states, iterations) + + logger.info("Measuring peak memory usage") + peak_alloc, peak_reserved = memory_tracker.get_peak_usage() + logger.info("Peak memory - allocated: %.2f MB, reserved: %.2f MB", + peak_alloc, peak_reserved) + + cache_hits = 0 + cache_misses = 0 + cache_hit_rate = 0.0 + + if model_type == "optimized": + logger.info("Extracting cache statistics from optimized model") + if hasattr(model, 'engram') and model.engram is not None: + if hasattr(model.engram.multi_head_embedding, 'cache'): + cache = model.engram.multi_head_embedding.cache + cache_hits = cache.hits + cache_misses = cache.misses + total_requests = cache_hits + cache_misses + cache_hit_rate = (cache_hits / total_requests * 100) if total_requests > 0 else 0.0 + + logger.info("Cache statistics:") + logger.info(" Hits: %d", cache_hits) + logger.info(" Misses: %d", cache_misses) + logger.info(" Total requests: %d", total_requests) + logger.info(" Hit rate: %.2f%%", cache_hit_rate) + + result = BenchmarkResult( + model_type=model_type, + layer_id=layer_id, + batch_size=batch_size, + seq_len=seq_len, + forward_time_ms=timing_stats['mean'], + memory_allocated_mb=peak_alloc, + memory_reserved_mb=peak_reserved, + cache_hits=cache_hits, + cache_misses=cache_misses, + cache_hit_rate_pct=cache_hit_rate, + iterations=iterations, + device=str(device) + ) + + logger.info("Benchmark complete for %s model", model_type) + logger.info(" Mean latency: %.2f ms", result.forward_time_ms) + logger.info(" Memory allocated: %.2f MB", result.memory_allocated_mb) + + return result, timing_stats + +def run_benchmark_suite( + layer_ids: List[int], + batch_sizes: List[int], + seq_lens: List[int], + iterations: int = 100 +): + """Run comprehensive benchmark suite""" + logger.info("#"*80) + logger.info("STARTING BENCHMARK SUITE") + logger.info("#"*80) + + if torch.cuda.is_available(): + device = torch.device("cuda") + logger.info("Using CUDA device: %s", torch.cuda.get_device_name(0)) + logger.info("CUDA capability: %s", torch.cuda.get_device_capability(0)) + logger.info("Total GPU memory: %.2f GB", + torch.cuda.get_device_properties(0).total_memory / 1024**3) + else: + device = torch.device("cpu") + logger.warning("CUDA not available, using CPU") + + memory_tracker = MemoryTracker(device) + + backbone_config = BackBoneConfigOriginal() + + logger.info("Benchmark configuration:") + logger.info(" Layer IDs: %s", layer_ids) + logger.info(" Batch sizes: %s", batch_sizes) + logger.info(" Sequence lengths: %s", seq_lens) + logger.info(" Iterations per test: %d", iterations) + + all_results = [] + all_timing_stats = [] + + total_tests = len(layer_ids) * len(batch_sizes) * len(seq_lens) * 2 + current_test = 0 + + for layer_id in layer_ids: + logger.info("") + logger.info("Testing layer %d", layer_id) + + for batch_size in batch_sizes: + for seq_len in seq_lens: + current_test += 1 + logger.info("") + logger.info("Test %d/%d: batch_size=%d, seq_len=%d", + current_test, total_tests, batch_size, seq_len) + + try: + logger.info("Benchmarking ORIGINAL model") + model_orig = TransformerBlockOriginal(layer_id=layer_id) + result_orig, timing_orig = benchmark_model( + model_orig, "original", layer_id, batch_size, seq_len, + backbone_config, device, memory_tracker, iterations + ) + all_results.append(result_orig) + all_timing_stats.append({ + 'model_type': 'original', + 'layer_id': layer_id, + 'batch_size': batch_size, + 'seq_len': seq_len, + 'stats': timing_orig + }) + + del model_orig + if device.type == 'cuda': + torch.cuda.empty_cache() + + current_test += 1 + logger.info("") + logger.info("Test %d/%d: batch_size=%d, seq_len=%d", + current_test, total_tests, batch_size, seq_len) + + logger.info("Benchmarking OPTIMIZED model") + model_opt = TransformerBlockOptimized(layer_id=layer_id) + result_opt, timing_opt = benchmark_model( + model_opt, "optimized", layer_id, batch_size, seq_len, + backbone_config, device, memory_tracker, iterations + ) + all_results.append(result_opt) + all_timing_stats.append({ + 'model_type': 'optimized', + 'layer_id': layer_id, + 'batch_size': batch_size, + 'seq_len': seq_len, + 'stats': timing_opt + }) + + speedup = result_orig.forward_time_ms / result_opt.forward_time_ms + memory_reduction_pct = ( + (result_orig.memory_allocated_mb - result_opt.memory_allocated_mb) / + result_orig.memory_allocated_mb * 100 + ) + + logger.info("") + logger.info("COMPARISON for layer=%d, bs=%d, seqlen=%d:", + layer_id, batch_size, seq_len) + logger.info(" Speedup: %.2fx", speedup) + logger.info(" Original latency: %.2f ms", result_orig.forward_time_ms) + logger.info(" Optimized latency: %.2f ms", result_opt.forward_time_ms) + logger.info(" Memory reduction: %.2f%%", memory_reduction_pct) + logger.info(" Cache hit rate: %.2f%%", result_opt.cache_hit_rate_pct) + + del model_opt + if device.type == 'cuda': + torch.cuda.empty_cache() + + except Exception as e: + logger.error("Exception during benchmark: %s", e) + logger.exception("Full traceback:") + + logger.info("") + logger.info("#"*80) + logger.info("BENCHMARK SUITE COMPLETE") + logger.info("#"*80) + + return all_results, all_timing_stats + +def save_results(results: List[BenchmarkResult], timing_stats: List[Dict], output_dir: str = "results"): + """Save benchmark results to JSON and CSV""" + logger.info("Saving results to directory: %s", output_dir) + + output_path = Path(output_dir) + output_path.mkdir(exist_ok=True) + logger.info("Created output directory: %s", output_path.absolute()) + + results_json = [asdict(r) for r in results] + json_path = output_path / "benchmark_results.json" + + logger.info("Writing results to JSON: %s", json_path) + with open(json_path, 'w') as f: + json.dump({ + 'results': results_json, + 'timestamp': time.strftime('%Y-%m-%d %H:%M:%S') + }, f, indent=2) + logger.info("JSON results saved: %d entries", len(results_json)) + + df = pd.DataFrame(results_json) + csv_path = output_path / "benchmark_results.csv" + + logger.info("Writing results to CSV: %s", csv_path) + df.to_csv(csv_path, index=False) + logger.info("CSV results saved: %d rows, %d columns", len(df), len(df.columns)) + + timing_json_path = output_path / "timing_statistics.json" + logger.info("Writing detailed timing stats to: %s", timing_json_path) + + timing_serializable = [] + for stat in timing_stats: + stat_copy = stat.copy() + if 'all_times' in stat_copy['stats']: + stat_copy['stats']['all_times'] = [float(t) for t in stat_copy['stats']['all_times']] + timing_serializable.append(stat_copy) + + with open(timing_json_path, 'w') as f: + json.dump(timing_serializable, f, indent=2) + logger.info("Timing statistics saved") + + return df + +def generate_plots(df: pd.DataFrame, output_dir: str = "results"): + """Generate comparison plots""" + logger.info("Generating visualization plots") + + output_path = Path(output_dir) + output_path.mkdir(exist_ok=True) + + df_pivot = df.pivot_table( + values='forward_time_ms', + index=['layer_id', 'batch_size', 'seq_len'], + columns='model_type' + ).reset_index() + + df_pivot['speedup'] = df_pivot['original'] / df_pivot['optimized'] + + logger.info("Creating speedup comparison plot") + fig, ax = plt.subplots(figsize=(12, 6)) + + x_labels = [f"L{row['layer_id']}_B{row['batch_size']}_S{row['seq_len']}" + for _, row in df_pivot.iterrows()] + x_pos = np.arange(len(x_labels)) + + width = 0.35 + ax.bar(x_pos - width/2, df_pivot['original'], width, label='Original', alpha=0.8) + ax.bar(x_pos + width/2, df_pivot['optimized'], width, label='Optimized', alpha=0.8) + + ax.set_xlabel('Configuration (Layer_BatchSize_SeqLen)') + ax.set_ylabel('Forward Pass Time (ms)') + ax.set_title('Original vs Optimized Engram Performance') + ax.set_xticks(x_pos) + ax.set_xticklabels(x_labels, rotation=45, ha='right') + ax.legend() + ax.grid(True, alpha=0.3) + + plot_path = output_path / "speedup_plot.png" + plt.tight_layout() + plt.savefig(plot_path, dpi=300, bbox_inches='tight') + logger.info("Speedup plot saved: %s", plot_path) + plt.close() + + logger.info("Creating speedup factor plot") + fig, ax = plt.subplots(figsize=(12, 6)) + + colors = ['green' if s > 1 else 'red' for s in df_pivot['speedup']] + bars = ax.bar(x_pos, df_pivot['speedup'], color=colors, alpha=0.7) + + ax.axhline(y=1.0, color='black', linestyle='--', linewidth=1, label='No speedup') + ax.set_xlabel('Configuration (Layer_BatchSize_SeqLen)') + ax.set_ylabel('Speedup Factor (Original / Optimized)') + ax.set_title('Speedup Factor Across Configurations') + ax.set_xticks(x_pos) + ax.set_xticklabels(x_labels, rotation=45, ha='right') + ax.legend() + ax.grid(True, alpha=0.3, axis='y') + + for i, (bar, speedup) in enumerate(zip(bars, df_pivot['speedup'])): + height = bar.get_height() + ax.text(bar.get_x() + bar.get_width()/2., height, + f'{speedup:.2f}x', + ha='center', va='bottom', fontsize=8) + + plot_path = output_path / "speedup_factor.png" + plt.tight_layout() + plt.savefig(plot_path, dpi=300, bbox_inches='tight') + logger.info("Speedup factor plot saved: %s", plot_path) + plt.close() + + logger.info("Creating memory comparison plot") + df_memory = df.pivot_table( + values='memory_allocated_mb', + index=['layer_id', 'batch_size', 'seq_len'], + columns='model_type' + ).reset_index() + + fig, ax = plt.subplots(figsize=(12, 6)) + + ax.bar(x_pos - width/2, df_memory['original'], width, label='Original', alpha=0.8) + ax.bar(x_pos + width/2, df_memory['optimized'], width, label='Optimized', alpha=0.8) + + ax.set_xlabel('Configuration (Layer_BatchSize_SeqLen)') + ax.set_ylabel('Peak Memory Allocated (MB)') + ax.set_title('Memory Usage Comparison') + ax.set_xticks(x_pos) + ax.set_xticklabels(x_labels, rotation=45, ha='right') + ax.legend() + ax.grid(True, alpha=0.3, axis='y') + + plot_path = output_path / "memory_comparison.png" + plt.tight_layout() + plt.savefig(plot_path, dpi=300, bbox_inches='tight') + logger.info("Memory comparison plot saved: %s", plot_path) + plt.close() + + df_opt = df[df['model_type'] == 'optimized'] + if not df_opt.empty and df_opt['cache_hit_rate_pct'].sum() > 0: + logger.info("Creating cache hit rate plot") + fig, ax = plt.subplots(figsize=(12, 6)) + + cache_x_labels = [f"L{row['layer_id']}_B{row['batch_size']}_S{row['seq_len']}" + for _, row in df_opt.iterrows()] + cache_x_pos = np.arange(len(cache_x_labels)) + + colors_cache = ['green' if rate > 50 else 'orange' if rate > 25 else 'red' + for rate in df_opt['cache_hit_rate_pct']] + bars = ax.bar(cache_x_pos, df_opt['cache_hit_rate_pct'], color=colors_cache, alpha=0.7) + + ax.set_xlabel('Configuration (Layer_BatchSize_SeqLen)') + ax.set_ylabel('Cache Hit Rate (%)') + ax.set_title('LRU Cache Hit Rate (Optimized Model)') + ax.set_xticks(cache_x_pos) + ax.set_xticklabels(cache_x_labels, rotation=45, ha='right') + ax.set_ylim(0, 100) + ax.grid(True, alpha=0.3, axis='y') + + for bar, rate in zip(bars, df_opt['cache_hit_rate_pct']): + height = bar.get_height() + ax.text(bar.get_x() + bar.get_width()/2., height, + f'{rate:.1f}%', + ha='center', va='bottom', fontsize=8) + + plot_path = output_path / "cache_hit_rate.png" + plt.tight_layout() + plt.savefig(plot_path, dpi=300, bbox_inches='tight') + logger.info("Cache hit rate plot saved: %s", plot_path) + plt.close() + + logger.info("All plots generated successfully") + +def print_summary(df: pd.DataFrame): + """Print summary statistics""" + logger.info("") + logger.info("#"*80) + logger.info("SUMMARY STATISTICS") + logger.info("#"*80) + + df_pivot = df.pivot_table( + values='forward_time_ms', + index=['layer_id', 'batch_size', 'seq_len'], + columns='model_type' + ).reset_index() + + df_pivot['speedup'] = df_pivot['original'] / df_pivot['optimized'] + + logger.info("") + logger.info("Speedup Statistics:") + logger.info(" Mean speedup: %.2fx", df_pivot['speedup'].mean()) + logger.info(" Median speedup: %.2fx", df_pivot['speedup'].median()) + logger.info(" Min speedup: %.2fx", df_pivot['speedup'].min()) + logger.info(" Max speedup: %.2fx", df_pivot['speedup'].max()) + logger.info(" Std speedup: %.2fx", df_pivot['speedup'].std()) + + logger.info("") + logger.info("Latency Statistics (ms):") + + df_orig = df[df['model_type'] == 'original'] + df_opt = df[df['model_type'] == 'optimized'] + + logger.info(" Original - Mean: %.2f, Median: %.2f, Std: %.2f", + df_orig['forward_time_ms'].mean(), + df_orig['forward_time_ms'].median(), + df_orig['forward_time_ms'].std()) + + logger.info(" Optimized - Mean: %.2f, Median: %.2f, Std: %.2f", + df_opt['forward_time_ms'].mean(), + df_opt['forward_time_ms'].median(), + df_opt['forward_time_ms'].std()) + + if not df_opt.empty and df_opt['memory_allocated_mb'].sum() > 0: + logger.info("") + logger.info("Memory Statistics (MB):") + + memory_reduction = ( + (df_orig['memory_allocated_mb'].mean() - df_opt['memory_allocated_mb'].mean()) / + df_orig['memory_allocated_mb'].mean() * 100 + ) + + logger.info(" Original - Mean: %.2f, Median: %.2f", + df_orig['memory_allocated_mb'].mean(), + df_orig['memory_allocated_mb'].median()) + + logger.info(" Optimized - Mean: %.2f, Median: %.2f", + df_opt['memory_allocated_mb'].mean(), + df_opt['memory_allocated_mb'].median()) + + logger.info(" Average memory reduction: %.2f%%", memory_reduction) + + if not df_opt.empty and df_opt['cache_hit_rate_pct'].sum() > 0: + logger.info("") + logger.info("Cache Statistics:") + logger.info(" Mean hit rate: %.2f%%", df_opt['cache_hit_rate_pct'].mean()) + logger.info(" Median hit rate: %.2f%%", df_opt['cache_hit_rate_pct'].median()) + logger.info(" Min hit rate: %.2f%%", df_opt['cache_hit_rate_pct'].min()) + logger.info(" Max hit rate: %.2f%%", df_opt['cache_hit_rate_pct'].max()) + + logger.info("") + logger.info("#"*80) + +if __name__ == '__main__': + logger.info("Starting benchmark script") + + cfg = EngramConfigOriginal() + + layer_ids = cfg.layer_ids + batch_sizes = [2, 4, 8] + seq_lens = [128, 256, 512] + iterations = 100 + + logger.info("Benchmark parameters:") + logger.info(" Layer IDs: %s", layer_ids) + logger.info(" Batch sizes: %s", batch_sizes) + logger.info(" Sequence lengths: %s", seq_lens) + logger.info(" Iterations: %d", iterations) + + results, timing_stats = run_benchmark_suite( + layer_ids=layer_ids, + batch_sizes=batch_sizes, + seq_lens=seq_lens, + iterations=iterations + ) + + logger.info("Saving results and generating visualizations") + df = save_results(results, timing_stats) + + generate_plots(df) + + print_summary(df) + + logger.info("Benchmark complete. Results saved to 'results/' directory") \ No newline at end of file diff --git a/engram_demo_v1.py b/engram_demo_v1.py index f3ce993..13a47e7 100644 --- a/engram_demo_v1.py +++ b/engram_demo_v1.py @@ -110,10 +110,12 @@ def _build_lookup_table(self): return lookup, len(new_tokens) def _compress(self, input_ids): - arr = np.asarray(input_ids, dtype=np.int64) + arr = input_ids.cpu().numpy() if isinstance(input_ids, torch.Tensor) else np.asarray(input_ids, dtype=np.int64) pos_mask = arr >= 0 out = arr.copy() - valid_ids = arr[pos_mask] + valid_ids = arr[pos_mask].astype(np.int64) + # Clip to valid range to prevent index out of bounds + valid_ids = np.clip(valid_ids, 0, len(self.lookup_table) - 1) out[pos_mask] = self.lookup_table[valid_ids] return out @@ -360,7 +362,7 @@ def forward(self,hidden_states,input_ids): hidden_states: [B, L, HC_MULT, D] input_ids: [B, L] """ - hash_input_ids = torch.from_numpy(self.hash_mapping.hash(input_ids)[self.layer_id]) + hash_input_ids = torch.from_numpy(self.hash_mapping.hash(input_ids)[self.layer_id]).to(input_ids.device) embeddings = self.multi_head_embedding(hash_input_ids).flatten(start_dim=-2) gates = [] for hc_idx in range(backbone_config.hc_mult): @@ -419,5 +421,4 @@ def forward(self,input_ids,hidden_states): hidden_states = layer(input_ids=input_ids,hidden_states=hidden_states) print("✅ Forward Complete!") - print(f"{input_ids.shape=}\n{output.shape=}") - \ No newline at end of file + print(f"{input_ids.shape=}\n{output.shape=}") \ No newline at end of file diff --git a/engram_optimized.py b/engram_optimized.py new file mode 100644 index 0000000..ba44bc1 --- /dev/null +++ b/engram_optimized.py @@ -0,0 +1,583 @@ +""" +================================================================================ +[Engram Architecture - Hybrid Optimized Implementation v3] + +Optimizations: +1. Hybrid CPU/GPU path selection based on input size +2. Small batches (< threshold): NumPy path (lower overhead) +3. Large batches (>= threshold): GPU path (better throughput) +4. No logging in forward pass +5. Pre-computed constants and buffers + +Compatible with benchmark.py and test_correctness.py +================================================================================ +""" + +from typing import List, Tuple +from dataclasses import dataclass, field +import math + +from sympy import isprime +import numpy as np +import torch +import torch.nn as nn +from transformers import AutoTokenizer +from tokenizers import normalizers, Regex + + +@dataclass +class EngramConfig: + tokenizer_name_or_path: str = "deepseek-ai/DeepSeek-V3" + engram_vocab_size: List[int] = field(default_factory=lambda: [129280*5, 129280*5]) + max_ngram_size: int = 3 + n_embed_per_ngram: int = 512 + n_head_per_ngram: int = 8 + layer_ids: List[int] = field(default_factory=lambda: [1, 15]) + pad_id: int = 2 + seed: int = 0 + kernel_size: int = 4 + # Threshold for GPU vs CPU path (batch_size * seq_len) + gpu_threshold: int = 1024 + + +@dataclass +class BackBoneConfig: + hidden_size: int = 1024 + hc_mult: int = 4 + vocab_size: int = 129280 + num_layers: int = 30 + + +engram_cfg = EngramConfig() +backbone_config = BackBoneConfig() + + +class CompressedTokenizer: + """ + Tokenizer compression via NFKC normalization. + Reduces vocabulary ~23% by merging semantically equivalent tokens. + """ + + def __init__(self, tokenizer_name_or_path: str): + self.tokenizer = AutoTokenizer.from_pretrained( + tokenizer_name_or_path, + trust_remote_code=True + ) + + SENTINEL = "\uE000" + self.normalizer = normalizers.Sequence([ + normalizers.NFKC(), + normalizers.NFD(), + normalizers.StripAccents(), + normalizers.Lowercase(), + normalizers.Replace(Regex(r"[ \t\r\n]+"), " "), + normalizers.Replace(Regex(r"^ $"), SENTINEL), + normalizers.Strip(), + normalizers.Replace(SENTINEL, " "), + ]) + + self.lookup_table, self.num_new_token = self._build_lookup_table() + self._lookup_tensor = None # Lazy init for GPU + self._lookup_tensor_device = None + + def __len__(self) -> int: + return self.num_new_token + + def _build_lookup_table(self) -> Tuple[np.ndarray, int]: + old2new = {} + key2new = {} + new_tokens = [] + + vocab_size = len(self.tokenizer) + for tid in range(vocab_size): + text = self.tokenizer.decode([tid], skip_special_tokens=False) + + if "�" in text: + key = self.tokenizer.convert_ids_to_tokens(tid) + else: + norm = self.normalizer.normalize_str(text) + key = norm if norm else text + + nid = key2new.get(key) + if nid is None: + nid = len(new_tokens) + key2new[key] = nid + new_tokens.append(key) + old2new[tid] = nid + + lookup = np.empty(vocab_size, dtype=np.int64) + for tid in range(vocab_size): + lookup[tid] = old2new[tid] + + return lookup, len(new_tokens) + + def compress_cpu(self, input_ids: np.ndarray) -> np.ndarray: + """Fast CPU path using NumPy indexing""" + clamped = np.clip(input_ids, 0, len(self.lookup_table) - 1) + return self.lookup_table[clamped] + + def compress_gpu(self, input_ids: torch.Tensor) -> torch.Tensor: + """GPU path with lazy tensor initialization""" + device = input_ids.device + + # Lazy init / device transfer + if self._lookup_tensor is None or self._lookup_tensor_device != device: + self._lookup_tensor = torch.from_numpy(self.lookup_table).to(device) + self._lookup_tensor_device = device + + clamped = input_ids.clamp(0, len(self.lookup_table) - 1) + return self._lookup_tensor[clamped] + + def __call__(self, input_ids): + """CPU fallback for compatibility""" + if isinstance(input_ids, torch.Tensor): + input_ids = input_ids.cpu().numpy() + arr = np.asarray(input_ids, dtype=np.int64) + return self.compress_cpu(arr) + + +class ShortConv(nn.Module): + """Depthwise causal convolution with RMSNorm and SiLU activation""" + + def __init__( + self, + hidden_size: int, + kernel_size: int = 4, + dilation: int = 1, + norm_eps: float = 1e-5, + hc_mult: int = 4, + activation: bool = True, + ): + super().__init__() + self.hc_mult = hc_mult + self.activation = activation + + total_channels = hidden_size * hc_mult + self.conv = nn.Conv1d( + in_channels=total_channels, + out_channels=total_channels, + kernel_size=kernel_size, + groups=total_channels, + bias=False, + padding=(kernel_size - 1) * dilation, + dilation=dilation, + ) + + self.norms = nn.ModuleList([ + nn.RMSNorm(hidden_size, eps=norm_eps) + for _ in range(hc_mult) + ]) + + if self.activation: + self.act_fn = nn.SiLU() + + def forward(self, x: torch.Tensor) -> torch.Tensor: + B, T, G, C = x.shape + + normed_chunks = [self.norms[i](x[:, :, i, :]) for i in range(G)] + x_norm = torch.cat(normed_chunks, dim=-1) + + x_bct = x_norm.transpose(1, 2) + y_bct = self.conv(x_bct)[..., :T] + + if self.activation: + y_bct = self.act_fn(y_bct) + + return y_bct.transpose(1, 2).view(B, T, G, C).contiguous() + + +def find_next_prime(start: int, seen_primes: set) -> int: + candidate = start + 1 + while True: + if isprime(candidate) and candidate not in seen_primes: + return candidate + candidate += 1 + + +class HybridNgramHashMapping(nn.Module): + """ + Hybrid CPU/GPU N-gram hashing module. + + Automatically selects optimal path based on input size: + - Small inputs: CPU/NumPy (lower overhead) + - Large inputs: GPU (better throughput) + """ + + def __init__( + self, + engram_vocab_size: List[int], + max_ngram_size: int, + n_embed_per_ngram: int, + n_head_per_ngram: int, + layer_ids: List[int], + tokenizer_name_or_path: str, + pad_id: int, + seed: int, + gpu_threshold: int = 1024, + ): + super().__init__() + + self.vocab_size_per_ngram = engram_vocab_size + self.max_ngram_size = max_ngram_size + self.n_embed_per_ngram = n_embed_per_ngram + self.n_head_per_ngram = n_head_per_ngram + self.pad_id = pad_id + self.layer_ids = layer_ids + self.gpu_threshold = gpu_threshold + + # Initialize compressed tokenizer + self.compressed_tokenizer = CompressedTokenizer(tokenizer_name_or_path) + self.tokenizer_vocab_size = len(self.compressed_tokenizer) + + if self.pad_id is not None: + self.pad_id = int(self.compressed_tokenizer.lookup_table[self.pad_id]) + + # Compute layer-specific multipliers + max_long = np.iinfo(np.int64).max + M_max = int(max_long // self.tokenizer_vocab_size) + half_bound = max(1, M_max // 2) + PRIME_1 = 10007 + + self.layer_multipliers = {} + for layer_id in self.layer_ids: + base_seed = int(seed + PRIME_1 * int(layer_id)) + g = np.random.default_rng(base_seed) + r = g.integers(low=0, high=half_bound, size=(self.max_ngram_size,), dtype=np.int64) + self.layer_multipliers[layer_id] = r * 2 + 1 + + # Calculate prime modulo sizes + self.vocab_size_across_layers = self._calculate_vocab_size_across_layers() + + # Register GPU buffers + for layer_id in self.layer_ids: + mult_tensor = torch.tensor(self.layer_multipliers[layer_id], dtype=torch.long) + self.register_buffer(f"multipliers_layer_{layer_id}", mult_tensor) + + all_primes = [p for ngram_heads in self.vocab_size_across_layers[layer_id] for p in ngram_heads] + self.register_buffer(f"prime_mods_layer_{layer_id}", torch.tensor(all_primes, dtype=torch.long)) + + # Also store as numpy for CPU path + setattr(self, f"_np_primes_layer_{layer_id}", np.array(all_primes, dtype=np.int64)) + + def _calculate_vocab_size_across_layers(self) -> dict: + seen_primes = set() + vocab_size_across_layers = {} + + for layer_id in self.layer_ids: + all_ngram_vocab_sizes = [] + for ngram in range(2, self.max_ngram_size + 1): + current_ngram_heads_sizes = [] + vocab_size = self.vocab_size_per_ngram[ngram - 2] + current_prime_search_start = vocab_size - 1 + + for _ in range(self.n_head_per_ngram): + found_prime = find_next_prime(current_prime_search_start, seen_primes) + seen_primes.add(found_prime) + current_ngram_heads_sizes.append(found_prime) + current_prime_search_start = found_prime + + all_ngram_vocab_sizes.append(current_ngram_heads_sizes) + vocab_size_across_layers[layer_id] = all_ngram_vocab_sizes + + return vocab_size_across_layers + + def _hash_cpu(self, input_ids: np.ndarray, layer_id: int) -> np.ndarray: + """CPU path: NumPy-based hashing (optimal for small inputs)""" + B, T = input_ids.shape + + multipliers = self.layer_multipliers[layer_id] + prime_mods = getattr(self, f"_np_primes_layer_{layer_id}") + + # Build shifted versions + shifts = [input_ids] + for k in range(1, self.max_ngram_size): + padding = np.full((B, k), self.pad_id, dtype=np.int64) + shifts.append(np.concatenate([padding, input_ids[:, :-k]], axis=1)) + + # Compute hashes + all_hashes = [] + hash_idx = 0 + + for n in range(2, self.max_ngram_size + 1): + mix = shifts[0] * multipliers[0] + for k in range(1, n): + mix = np.bitwise_xor(mix, shifts[k] * multipliers[k]) + + for _ in range(self.n_head_per_ngram): + all_hashes.append((mix % prime_mods[hash_idx]).astype(np.int64)) + hash_idx += 1 + + return np.stack(all_hashes, axis=2) + + def _hash_gpu(self, input_ids: torch.Tensor, layer_id: int) -> torch.Tensor: + """GPU path: PyTorch-based hashing (optimal for large inputs)""" + B, T = input_ids.shape + device = input_ids.device + + multipliers = getattr(self, f"multipliers_layer_{layer_id}") + prime_mods = getattr(self, f"prime_mods_layer_{layer_id}") + + # Build shifted versions + shifts = [input_ids] + for k in range(1, self.max_ngram_size): + padding = torch.full((B, k), self.pad_id, dtype=torch.long, device=device) + shifts.append(torch.cat([padding, input_ids[:, :-k]], dim=1)) + + # Compute hashes + all_hashes = [] + hash_idx = 0 + + for n in range(2, self.max_ngram_size + 1): + mix = shifts[0] * multipliers[0] + for k in range(1, n): + mix = torch.bitwise_xor(mix, shifts[k] * multipliers[k]) + + for _ in range(self.n_head_per_ngram): + all_hashes.append(mix % prime_mods[hash_idx]) + hash_idx += 1 + + return torch.stack(all_hashes, dim=2) + + def hash(self, input_ids: torch.Tensor) -> dict: + """ + Main hashing interface with automatic path selection. + + Args: + input_ids: Raw token IDs [B, T] + + Returns: + Dictionary mapping layer_id -> hash indices [B, T, num_heads] + """ + B, T = input_ids.shape + input_size = B * T + device = input_ids.device + + if input_size < self.gpu_threshold: + # CPU path for small inputs + input_np = input_ids.cpu().numpy() + compressed_np = self.compressed_tokenizer.compress_cpu(input_np) + + result = {} + for layer_id in self.layer_ids: + hash_np = self._hash_cpu(compressed_np, layer_id) + result[layer_id] = torch.from_numpy(hash_np).to(device) + return result + else: + # GPU path for large inputs + compressed = self.compressed_tokenizer.compress_gpu(input_ids) + return { + layer_id: self._hash_gpu(compressed, layer_id) + for layer_id in self.layer_ids + } + + +class MultiHeadEmbedding(nn.Module): + """Multi-head embedding lookup with concatenated tables""" + + def __init__(self, list_of_N: List[int], D: int): + super().__init__() + self.num_heads = len(list_of_N) + self.embedding_dim = D + + offsets = [0] + for n in list_of_N[:-1]: + offsets.append(offsets[-1] + n) + self.register_buffer("offsets", torch.tensor(offsets, dtype=torch.long)) + + total_N = sum(list_of_N) + self.embedding = nn.Embedding(num_embeddings=total_N, embedding_dim=D) + + def forward(self, input_ids: torch.Tensor) -> torch.Tensor: + return self.embedding(input_ids + self.offsets) + + +class Engram(nn.Module): + """ + Engram conditional memory module with hybrid optimization. + + Retrieves static n-gram embeddings and fuses them with dynamic hidden states + via context-aware gating. + """ + + def __init__(self, layer_id: int): + super().__init__() + self.layer_id = layer_id + + # Hybrid hash mapping + self.hash_mapping = HybridNgramHashMapping( + engram_vocab_size=engram_cfg.engram_vocab_size, + max_ngram_size=engram_cfg.max_ngram_size, + n_embed_per_ngram=engram_cfg.n_embed_per_ngram, + n_head_per_ngram=engram_cfg.n_head_per_ngram, + layer_ids=engram_cfg.layer_ids, + tokenizer_name_or_path=engram_cfg.tokenizer_name_or_path, + pad_id=engram_cfg.pad_id, + seed=engram_cfg.seed, + gpu_threshold=engram_cfg.gpu_threshold, + ) + + # Multi-head embedding + vocab_sizes = [x for y in self.hash_mapping.vocab_size_across_layers[self.layer_id] for x in y] + embed_dim = engram_cfg.n_embed_per_ngram // engram_cfg.n_head_per_ngram + self.multi_head_embedding = MultiHeadEmbedding(list_of_N=vocab_sizes, D=embed_dim) + + # Short convolution + self.short_conv = ShortConv( + hidden_size=backbone_config.hidden_size, + kernel_size=engram_cfg.kernel_size, + dilation=engram_cfg.max_ngram_size, + hc_mult=backbone_config.hc_mult, + ) + + # Projections + engram_hidden_size = (engram_cfg.max_ngram_size - 1) * engram_cfg.n_embed_per_ngram + self.value_proj = nn.Linear(engram_hidden_size, backbone_config.hidden_size) + self.key_projs = nn.ModuleList([ + nn.Linear(engram_hidden_size, backbone_config.hidden_size) + for _ in range(backbone_config.hc_mult) + ]) + + # Normalizations + self.norm1 = nn.ModuleList([ + nn.RMSNorm(backbone_config.hidden_size) + for _ in range(backbone_config.hc_mult) + ]) + self.norm2 = nn.ModuleList([ + nn.RMSNorm(backbone_config.hidden_size) + for _ in range(backbone_config.hc_mult) + ]) + + # Pre-computed constant + self._inv_sqrt_d = 1.0 / math.sqrt(backbone_config.hidden_size) + + def forward(self, hidden_states: torch.Tensor, input_ids: torch.Tensor) -> torch.Tensor: + """ + Args: + hidden_states: [B, T, HC_MULT, D] + input_ids: [B, T] + + Returns: + output: [B, T, HC_MULT, D] + """ + # Hash and embed + hash_result = self.hash_mapping.hash(input_ids) + hash_indices = hash_result[self.layer_id] + embeddings = self.multi_head_embedding(hash_indices).flatten(start_dim=-2) + + # Context-aware gating + gates = [] + for hc_idx in range(backbone_config.hc_mult): + key = self.key_projs[hc_idx](embeddings) + normed_key = self.norm1[hc_idx](key) + query = hidden_states[:, :, hc_idx, :] + normed_query = self.norm2[hc_idx](query) + + gate = (normed_key * normed_query).sum(dim=-1) * self._inv_sqrt_d + gate = gate.abs().clamp_min(1e-6).sqrt() * gate.sign() + gates.append(gate.sigmoid().unsqueeze(-1)) + + gates = torch.stack(gates, dim=2) + + # Gated value + convolution + value = gates * self.value_proj(embeddings).unsqueeze(2) + return value + self.short_conv(value) + + +class TransformerBlock(nn.Module): + """Transformer block with optional Engram module""" + + def __init__(self, layer_id: int): + super().__init__() + self.attn = lambda x: x + self.moe = lambda x: x + self.engram = Engram(layer_id=layer_id) if layer_id in engram_cfg.layer_ids else None + + def forward(self, input_ids: torch.Tensor, hidden_states: torch.Tensor) -> torch.Tensor: + if self.engram is not None: + hidden_states = self.engram(hidden_states=hidden_states, input_ids=input_ids) + hidden_states + hidden_states = self.attn(hidden_states) + hidden_states + hidden_states = self.moe(hidden_states) + hidden_states + return hidden_states + + +if __name__ == '__main__': + import time + + print("=" * 60) + print("Engram Hybrid Optimized v3 - Demo") + print("=" * 60) + print(f"GPU threshold: {engram_cfg.gpu_threshold} elements") + + # Build model + LLM = [ + nn.Embedding(backbone_config.vocab_size, backbone_config.hidden_size), + *[TransformerBlock(layer_id=layer_id) for layer_id in range(backbone_config.num_layers)], + nn.Linear(backbone_config.hidden_size, backbone_config.vocab_size) + ] + + # Tokenize input + text = "Only Alexander the Great could tame the horse Bucephalus." + tokenizer = AutoTokenizer.from_pretrained(engram_cfg.tokenizer_name_or_path, trust_remote_code=True) + input_ids = tokenizer(text, return_tensors='pt').input_ids + + print(f"\nInput: {text}") + print(f"Tokenized shape: {input_ids.shape}") + print(f"Input size: {input_ids.shape[0] * input_ids.shape[1]} -> {'CPU' if input_ids.shape[0] * input_ids.shape[1] < engram_cfg.gpu_threshold else 'GPU'} path") + + # Device + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + print(f"Device: {device}") + + input_ids = input_ids.to(device) + for layer in LLM: + layer.to(device) + + # Forward pass + print("\nRunning forward pass...") + start = time.time() + + for idx, layer in enumerate(LLM): + if idx == 0: + hidden_states = LLM[0](input_ids) + hidden_states = hidden_states.unsqueeze(2).expand(-1, -1, backbone_config.hc_mult, -1) + elif idx == len(LLM) - 1: + hidden_states = hidden_states[:, :, 0, :] + output = layer(hidden_states) + else: + hidden_states = layer(input_ids=input_ids, hidden_states=hidden_states) + + elapsed = (time.time() - start) * 1000 + + print(f"\n Forward pass complete in {elapsed:.2f}ms") + print(f"Output shape: {output.shape}") + + # Test both paths explicitly + print("\n" + "=" * 60) + print("Testing both paths explicitly:") + print("=" * 60) + + engram = LLM[2].engram # Layer 1 has Engram + + # Small input (CPU path) + small_ids = torch.randint(0, 1000, (2, 64), device=device) + small_hidden = torch.randn(2, 64, backbone_config.hc_mult, backbone_config.hidden_size, device=device) + + start = time.time() + for _ in range(10): + _ = engram(small_hidden, small_ids) + if device.type == 'cuda': + torch.cuda.synchronize() + small_time = (time.time() - start) * 100 + print(f"Small input (2x64={2*64}): {small_time:.2f}ms avg -> CPU path") + + # Large input (GPU path) + large_ids = torch.randint(0, 1000, (8, 512), device=device) + large_hidden = torch.randn(8, 512, backbone_config.hc_mult, backbone_config.hidden_size, device=device) + + start = time.time() + for _ in range(10): + _ = engram(large_hidden, large_ids) + if device.type == 'cuda': + torch.cuda.synchronize() + large_time = (time.time() - start) * 100 + print(f"Large input (8x512={8*512}): {large_time:.2f}ms avg -> GPU path") \ No newline at end of file diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..0335c71 --- /dev/null +++ b/requirements.txt @@ -0,0 +1,7 @@ +torch>=2.0.0 +numpy>=1.24.0 +transformers>=4.30.0 +tokenizers>=0.13.0 +sympy>=1.12 +matplotlib>=3.7.0 +pandas>=2.0.0 \ No newline at end of file diff --git a/test_correctness.py b/test_correctness.py new file mode 100644 index 0000000..3d16bcb --- /dev/null +++ b/test_correctness.py @@ -0,0 +1,413 @@ +""" +Correctness testing: Compare original vs optimized Engram outputs +""" + +import logging +import time +import sys +import torch +import numpy as np +from transformers import AutoTokenizer + +logging.basicConfig( + level=logging.INFO, + format='[%(asctime)s] %(levelname)s - %(name)s - %(message)s' +) +logger = logging.getLogger(__name__) + +try: + logger.info("Importing original Engram implementation") + from engram_demo_v1 import ( + Engram as EngramOriginal, + TransformerBlock as TransformerBlockOriginal, + EngramConfig as EngramConfigOriginal, + BackBoneConfig as BackBoneConfigOriginal, + ) + logger.info("Original Engram imported successfully") +except ImportError as e: + logger.error("Failed to import original Engram: %s", e) + sys.exit(1) + +try: + logger.info("Importing optimized Engram implementation") + from engram_optimized import ( + Engram as EngramOptimized, + TransformerBlock as TransformerBlockOptimized, + EngramConfig as EngramConfigOptimized, + BackBoneConfig as BackBoneConfigOptimized, + ) + logger.info("Optimized Engram imported successfully") +except ImportError as e: + logger.error("Failed to import optimized Engram: %s", e) + sys.exit(1) + +def set_seed(seed=42): + """Set random seeds for reproducibility""" + logger.info("Setting random seed to %d for reproducibility", seed) + torch.manual_seed(seed) + torch.cuda.manual_seed_all(seed) + np.random.seed(seed) + if torch.cuda.is_available(): + torch.backends.cudnn.deterministic = True + torch.backends.cudnn.benchmark = False + logger.info("Random seed set successfully") + +def create_test_data(batch_size, seq_len, vocab_size, device): + """Generate synthetic test data""" + logger.info("Creating test data: batch_size=%d, seq_len=%d, vocab_size=%d", + batch_size, seq_len, vocab_size) + + set_seed(42) + input_ids = torch.randint(0, vocab_size, (batch_size, seq_len), device=device) + logger.debug("Generated input_ids with shape=%s, dtype=%s", + input_ids.shape, input_ids.dtype) + + return input_ids + +def create_hidden_states(batch_size, seq_len, hidden_size, hc_mult, device): + """Generate synthetic hidden states""" + logger.info("Creating hidden states: batch_size=%d, seq_len=%d, hidden_size=%d, hc_mult=%d", + batch_size, seq_len, hidden_size, hc_mult) + + set_seed(42) + hidden_states = torch.randn(batch_size, seq_len, hc_mult, hidden_size, device=device) + logger.debug("Generated hidden_states with shape=%s, dtype=%s", + hidden_states.shape, hidden_states.dtype) + + return hidden_states + +def copy_weights(model_src, model_dst): + """Copy weights from source model to destination model""" + logger.info("Copying weights from source to destination model") + + src_state = model_src.state_dict() + dst_state = model_dst.state_dict() + + logger.debug("Source model has %d parameters", len(src_state)) + logger.debug("Destination model has %d parameters", len(dst_state)) + + matched_keys = 0 + missing_keys = [] + unexpected_keys = [] + + for key in dst_state.keys(): + if key in src_state: + if dst_state[key].shape == src_state[key].shape: + dst_state[key].copy_(src_state[key]) + matched_keys += 1 + logger.debug("Copied parameter: %s with shape=%s", key, src_state[key].shape) + else: + logger.warning("Shape mismatch for key=%s: src=%s, dst=%s", + key, src_state[key].shape, dst_state[key].shape) + else: + missing_keys.append(key) + + for key in src_state.keys(): + if key not in dst_state: + unexpected_keys.append(key) + + logger.info("Weight copying complete: matched=%d, missing=%d, unexpected=%d", + matched_keys, len(missing_keys), len(unexpected_keys)) + + if missing_keys: + logger.warning("Missing keys in destination model: %s", missing_keys[:5]) + if unexpected_keys: + logger.warning("Unexpected keys in source model: %s", unexpected_keys[:5]) + + model_dst.load_state_dict(dst_state) + logger.info("Destination model state_dict loaded") + +def compare_outputs(output_orig, output_opt, rtol=1e-4, atol=1e-5): + """Compare two tensors and return detailed statistics""" + logger.info("Comparing outputs with rtol=%.2e, atol=%.2e", rtol, atol) + + logger.debug("Original output - shape=%s, dtype=%s, device=%s", + output_orig.shape, output_orig.dtype, output_orig.device) + logger.debug("Optimized output - shape=%s, dtype=%s, device=%s", + output_opt.shape, output_opt.dtype, output_opt.device) + + if output_orig.shape != output_opt.shape: + logger.error("Shape mismatch: original=%s, optimized=%s", + output_orig.shape, output_opt.shape) + return False + + is_close = torch.allclose(output_orig, output_opt, rtol=rtol, atol=atol) + + diff = torch.abs(output_orig - output_opt) + max_diff = diff.max().item() + mean_diff = diff.mean().item() + median_diff = diff.median().item() + + rel_diff = diff / (torch.abs(output_orig) + 1e-8) + max_rel_diff = rel_diff.max().item() + mean_rel_diff = rel_diff.mean().item() + + logger.info("Absolute difference statistics:") + logger.info(" Max: %.6e", max_diff) + logger.info(" Mean: %.6e", mean_diff) + logger.info(" Median: %.6e", median_diff) + + logger.info("Relative difference statistics:") + logger.info(" Max: %.6e", max_rel_diff) + logger.info(" Mean: %.6e", mean_rel_diff) + + if is_close: + logger.info("Outputs match within tolerance: PASS") + else: + logger.error("Outputs do NOT match within tolerance: FAIL") + + mismatch_mask = ~torch.isclose(output_orig, output_opt, rtol=rtol, atol=atol) + num_mismatches = mismatch_mask.sum().item() + total_elements = output_orig.numel() + mismatch_pct = (num_mismatches / total_elements) * 100 + + logger.error("Number of mismatched elements: %d / %d (%.2f%%)", + num_mismatches, total_elements, mismatch_pct) + + if num_mismatches > 0 and num_mismatches <= 10: + mismatch_indices = torch.nonzero(mismatch_mask) + logger.error("First few mismatched positions:") + for idx in mismatch_indices[:5]: + idx_tuple = tuple(idx.tolist()) + orig_val = output_orig[idx_tuple].item() + opt_val = output_opt[idx_tuple].item() + logger.error(" Position %s: original=%.6e, optimized=%.6e", + idx_tuple, orig_val, opt_val) + + return is_close + +def test_engram_module(layer_id, batch_size, seq_len, device): + """Test single Engram module correctness""" + logger.info("=" * 80) + logger.info("Testing Engram module for layer_id=%d", layer_id) + logger.info("=" * 80) + + set_seed(42) + + logger.info("Initializing original Engram module") + engram_orig = EngramOriginal(layer_id=layer_id) + logger.info("Original Engram parameters: %d", + sum(p.numel() for p in engram_orig.parameters())) + + logger.info("Initializing optimized Engram module") + engram_opt = EngramOptimized(layer_id=layer_id) + logger.info("Optimized Engram parameters: %d", + sum(p.numel() for p in engram_opt.parameters())) + + logger.info("Copying weights from original to optimized") + copy_weights(engram_orig, engram_opt) + + if device.type == 'cuda': + logger.info("Moving models to GPU") + engram_orig = engram_orig.to(device) + engram_opt = engram_opt.to(device) + + engram_orig.eval() + engram_opt.eval() + logger.info("Models set to evaluation mode") + + cfg_orig = EngramConfigOriginal() + backbone_orig = BackBoneConfigOriginal() + + actual_vocab_size = len(engram_orig.hash_mapping.compressed_tokenizer.tokenizer) + input_ids = create_test_data(batch_size, seq_len, actual_vocab_size, device) + hidden_states = create_hidden_states( + batch_size, seq_len, + backbone_orig.hidden_size, + backbone_orig.hc_mult, + device + ) + + logger.info("Running forward pass on original Engram") + start_time = time.time() + with torch.no_grad(): + output_orig = engram_orig(hidden_states, input_ids) + orig_time = (time.time() - start_time) * 1000 + logger.info("Original Engram forward pass took %.2fms", orig_time) + + logger.info("Running forward pass on optimized Engram") + start_time = time.time() + with torch.no_grad(): + output_opt = engram_opt(hidden_states, input_ids) + opt_time = (time.time() - start_time) * 1000 + logger.info("Optimized Engram forward pass took %.2fms", opt_time) + + speedup = orig_time / opt_time if opt_time > 0 else 0 + logger.info("Speedup: %.2fx (original=%.2fms, optimized=%.2fms)", + speedup, orig_time, opt_time) + + logger.info("Comparing outputs") + passed = compare_outputs(output_orig, output_opt) + + if passed: + logger.info("Test PASSED for layer_id=%d", layer_id) + else: + logger.error("Test FAILED for layer_id=%d", layer_id) + + logger.info("Logging cache statistics") + if hasattr(engram_opt.multi_head_embedding, 'log_cache_stats'): + cache_stats = engram_opt.multi_head_embedding.log_cache_stats() + logger.info("Cache hit rate: %.2f%%", cache_stats['hit_rate_pct']) + + return passed, orig_time, opt_time + +def test_transformer_block(layer_id, batch_size, seq_len, device): + """Test full TransformerBlock correctness""" + logger.info("=" * 80) + logger.info("Testing TransformerBlock for layer_id=%d", layer_id) + logger.info("=" * 80) + + set_seed(42) + + logger.info("Initializing original TransformerBlock") + block_orig = TransformerBlockOriginal(layer_id=layer_id) + + logger.info("Initializing optimized TransformerBlock") + block_opt = TransformerBlockOptimized(layer_id=layer_id) + + logger.info("Copying weights") + copy_weights(block_orig, block_opt) + + if device.type == 'cuda': + logger.info("Moving blocks to GPU") + block_orig = block_orig.to(device) + block_opt = block_opt.to(device) + + block_orig.eval() + block_opt.eval() + + backbone_orig = BackBoneConfigOriginal() + + actual_vocab_size = len(block_orig.engram.hash_mapping.compressed_tokenizer.tokenizer) if block_orig.engram else backbone_orig.vocab_size + input_ids = create_test_data(batch_size, seq_len, actual_vocab_size, device) + hidden_states = create_hidden_states( + batch_size, seq_len, + backbone_orig.hidden_size, + backbone_orig.hc_mult, + device + ) + + logger.info("Running forward pass on original TransformerBlock") + start_time = time.time() + with torch.no_grad(): + output_orig = block_orig(input_ids, hidden_states) + orig_time = (time.time() - start_time) * 1000 + logger.info("Original TransformerBlock forward pass took %.2fms", orig_time) + + logger.info("Running forward pass on optimized TransformerBlock") + start_time = time.time() + with torch.no_grad(): + output_opt = block_opt(input_ids, hidden_states) + opt_time = (time.time() - start_time) * 1000 + logger.info("Optimized TransformerBlock forward pass took %.2fms", opt_time) + + speedup = orig_time / opt_time if opt_time > 0 else 0 + logger.info("Speedup: %.2fx", speedup) + + passed = compare_outputs(output_orig, output_opt) + + if passed: + logger.info("Test PASSED for TransformerBlock layer_id=%d", layer_id) + else: + logger.error("Test FAILED for TransformerBlock layer_id=%d", layer_id) + + return passed, orig_time, opt_time + +def run_all_tests(): + """Run all correctness tests""" + logger.info("#" * 80) + logger.info("STARTING ENGRAM CORRECTNESS TESTS") + logger.info("#" * 80) + + if torch.cuda.is_available(): + device = torch.device("cuda") + logger.info("Using CUDA device: %s", torch.cuda.get_device_name(0)) + logger.info("CUDA memory allocated: %.2f MB", + torch.cuda.memory_allocated() / 1024**2) + else: + device = torch.device("cpu") + logger.warning("CUDA not available, using CPU (tests will be slower)") + + cfg = EngramConfigOriginal() + test_batch_size = 4 + test_seq_len = 128 + + logger.info("Test configuration:") + logger.info(" Batch size: %d", test_batch_size) + logger.info(" Sequence length: %d", test_seq_len) + logger.info(" Engram layers: %s", cfg.layer_ids) + + all_passed = True + results = [] + + for layer_id in cfg.layer_ids: + logger.info("") + logger.info("Testing layer %d", layer_id) + + try: + passed_engram, orig_time_engram, opt_time_engram = test_engram_module( + layer_id, test_batch_size, test_seq_len, device + ) + + passed_block, orig_time_block, opt_time_block = test_transformer_block( + layer_id, test_batch_size, test_seq_len, device + ) + + layer_passed = passed_engram and passed_block + all_passed = all_passed and layer_passed + + results.append({ + 'layer_id': layer_id, + 'engram_passed': passed_engram, + 'block_passed': passed_block, + 'engram_speedup': orig_time_engram / opt_time_engram if opt_time_engram > 0 else 0, + 'block_speedup': orig_time_block / opt_time_block if opt_time_block > 0 else 0, + }) + + except Exception as e: + logger.error("Exception during test for layer %d: %s", layer_id, e) + logger.exception("Full traceback:") + all_passed = False + results.append({ + 'layer_id': layer_id, + 'engram_passed': False, + 'block_passed': False, + 'error': str(e) + }) + + logger.info("") + logger.info("#" * 80) + logger.info("TEST SUMMARY") + logger.info("#" * 80) + + for result in results: + layer_id = result['layer_id'] + if 'error' in result: + logger.error("Layer %d: ERROR - %s", layer_id, result['error']) + else: + engram_status = "PASS" if result['engram_passed'] else "FAIL" + block_status = "PASS" if result['block_passed'] else "FAIL" + logger.info("Layer %d: Engram=%s (%.2fx), Block=%s (%.2fx)", + layer_id, engram_status, result['engram_speedup'], + block_status, result['block_speedup']) + + logger.info("") + if all_passed: + logger.info("ALL TESTS PASSED") + logger.info("Optimized implementation is numerically correct") + + avg_engram_speedup = np.mean([r['engram_speedup'] for r in results if 'engram_speedup' in r]) + avg_block_speedup = np.mean([r['block_speedup'] for r in results if 'block_speedup' in r]) + logger.info("Average Engram speedup: %.2fx", avg_engram_speedup) + logger.info("Average Block speedup: %.2fx", avg_block_speedup) + else: + logger.error("SOME TESTS FAILED") + logger.error("Optimized implementation has correctness issues") + + logger.info("#" * 80) + + return all_passed + +if __name__ == '__main__': + success = run_all_tests() + sys.exit(0 if success else 1) \ No newline at end of file