Create inference/test_long_context.py

inference/test_long_context.py

@@ -0,0 +1,316 @@
+"""
+Test script to verify 250K context length support
+Tests RoPE scaling and long context handling
+"""
+
+import torch
+from transformers import AutoTokenizer, AutoModelForCausalLM, AutoConfig
+import logging
+from typing import Optional
+import time
+
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+
+class LongContextTester:
+    """Test long context capabilities of Helion-OSC"""
+    
+    def __init__(self, model_path: str = "./inference"):
+        """
+        Initialize tester
+        
+        Args:
+            model_path: Path to model inference directory
+        """
+        self.model_path = model_path
+        logger.info("Loading model configuration...")
+        
+        # Load config
+        self.config = AutoConfig.from_pretrained(model_path)
+        
+        # Verify context length
+        max_pos = self.config.max_position_embeddings
+        logger.info(f"Model max position embeddings: {max_pos:,}")
+        
+        if max_pos < 250000:
+            logger.warning(f"Context length ({max_pos:,}) is less than 250K!")
+        else:
+            logger.info(f"✓ Context length supports 250K+ tokens ({max_pos:,})")
+        
+        # Check RoPE scaling
+        rope_scaling = getattr(self.config, 'rope_scaling', None)
+        rope_theta = getattr(self.config, 'rope_theta', None)
+        
+        if rope_scaling:
+            logger.info(f"RoPE Scaling: {rope_scaling}")
+        if rope_theta:
+            logger.info(f"RoPE Theta: {rope_theta:,}")
+    
+    def test_tokenization_capacity(self, tokenizer_path: str = "DeepXR/Helion-OSC"):
+        """Test that tokenizer supports long sequences"""
+        logger.info("\n" + "="*80)
+        logger.info("TEST 1: Tokenizer Capacity")
+        logger.info("="*80)
+        
+        try:
+            tokenizer = AutoTokenizer.from_pretrained(tokenizer_path)
+            
+            max_length = tokenizer.model_max_length
+            logger.info(f"Tokenizer max length: {max_length:,}")
+            
+            if max_length >= 250000:
+                logger.info("✓ Tokenizer supports 250K+ tokens")
+            else:
+                logger.warning(f"✗ Tokenizer max length only {max_length:,}")
+            
+            # Test with a long sequence
+            test_tokens = 10000
+            test_text = "Hello world! " * (test_tokens // 2)
+            
+            logger.info(f"Testing tokenization of ~{test_tokens:,} tokens...")
+            encoded = tokenizer(test_text, return_tensors="pt", truncation=False)
+            actual_tokens = encoded['input_ids'].shape[1]
+            
+            logger.info(f"Successfully tokenized {actual_tokens:,} tokens")
+            logger.info("✓ Tokenization test passed")
+            
+            return True
+            
+        except Exception as e:
+            logger.error(f"✗ Tokenization test failed: {e}")
+            return False
+    
+    def test_position_embeddings(self):
+        """Test position embedding capacity"""
+        logger.info("\n" + "="*80)
+        logger.info("TEST 2: Position Embeddings")
+        logger.info("="*80)
+        
+        max_pos = self.config.max_position_embeddings
+        hidden_size = self.config.hidden_size
+        
+        logger.info(f"Max positions: {max_pos:,}")
+        logger.info(f"Hidden size: {hidden_size:,}")
+        
+        # Calculate memory requirement for position embeddings
+        if hasattr(self.config, 'rope_theta'):
+            logger.info("Using RoPE (Rotary Position Embeddings)")
+            logger.info("✓ RoPE scales efficiently to long contexts")
+            
+            # RoPE doesn't store position embeddings, it computes them
+            logger.info(f"RoPE Theta: {self.config.rope_theta:,}")
+            
+            if hasattr(self.config, 'rope_scaling'):
+                scaling = self.config.rope_scaling
+                logger.info(f"RoPE Scaling Configuration:")
+                logger.info(f"  Type: {scaling.get('type', 'N/A')}")
+                logger.info(f"  Factor: {scaling.get('factor', 'N/A')}")
+                
+                if scaling.get('factor', 0) >= 32:
+                    logger.info("✓ RoPE scaling factor supports 250K+ context (32x from 8K base)")
+                else:
+                    logger.warning("✗ RoPE scaling factor may be insufficient")
+            
+            return True
+        else:
+            # Learned position embeddings
+            pos_emb_size = max_pos * hidden_size * 2  # bfloat16
+            pos_emb_gb = pos_emb_size / (1024**3)
+            logger.info(f"Position embedding size: {pos_emb_gb:.2f} GB")
+            
+            if max_pos >= 250000:
+                logger.info("✓ Sufficient position embeddings for 250K context")
+                return True
+            else:
+                logger.warning("✗ Insufficient position embeddings")
+                return False
+    
+    def test_attention_computation(self, sequence_lengths: list = [1024, 8192, 32768, 131072]):
+        """Test attention computation at various lengths"""
+        logger.info("\n" + "="*80)
+        logger.info("TEST 3: Attention Computation Scaling")
+        logger.info("="*80)
+        
+        hidden_size = self.config.hidden_size
+        num_heads = self.config.num_attention_heads
+        head_dim = hidden_size // num_heads
+        
+        logger.info(f"Attention heads: {num_heads}")
+        logger.info(f"Head dimension: {head_dim}")
+        
+        for seq_len in sequence_lengths:
+            # Calculate attention matrix size
+            # For self-attention: (batch, heads, seq_len, seq_len)
+            attn_size = 1 * num_heads * seq_len * seq_len * 2  # bfloat16
+            attn_gb = attn_size / (1024**3)
+            
+            logger.info(f"\nSequence length: {seq_len:,} tokens")
+            logger.info(f"  Attention matrix: {attn_gb:.2f} GB")
+            
+            if seq_len <= 32768:
+                logger.info(f"  ✓ Manageable size")
+            elif seq_len <= 131072:
+                logger.info(f"  ⚠ Large - may need Flash Attention")
+            else:
+                logger.info(f"  ⚠ Very large - requires optimizations")
+        
+        # Check for Flash Attention support
+        use_flash = getattr(self.config, 'use_flash_attention_2', False)
+        if use_flash:
+            logger.info("\n✓ Flash Attention 2 enabled - efficient for long contexts")
+        else:
+            logger.warning("\n⚠ Flash Attention not configured - may be slow for long contexts")
+        
+        return True
+    
+    def test_memory_requirements(self):
+        """Calculate memory requirements for 250K context"""
+        logger.info("\n" + "="*80)
+        logger.info("TEST 4: Memory Requirements")
+        logger.info("="*80)
+        
+        context_length = 250000
+        batch_size = 1
+        hidden_size = self.config.hidden_size
+        num_layers = self.config.num_hidden_layers
+        
+        logger.info(f"Configuration:")
+        logger.info(f"  Context: {context_length:,} tokens")
+        logger.info(f"  Batch size: {batch_size}")
+        logger.info(f"  Hidden size: {hidden_size:,}")
+        logger.info(f"  Layers: {num_layers}")
+        
+        # Calculate activation memory (rough estimate)
+        # Main components: hidden states, attention outputs
+        hidden_states_size = batch_size * context_length * hidden_size * 2  # bfloat16
+        hidden_states_gb = hidden_states_size / (1024**3)
+        
+        # Per layer
+        layer_memory_gb = hidden_states_gb * 2  # rough estimate with attention
+        total_activation_gb = layer_memory_gb * num_layers
+        
+        logger.info(f"\nMemory estimates:")
+        logger.info(f"  Hidden states per layer: {hidden_states_gb:.2f} GB")
+        logger.info(f"  Total activation memory: {total_activation_gb:.2f} GB")
+        logger.info(f"  Model weights: ~349 GB")
+        logger.info(f"  Total (weights + activations): ~{349 + total_activation_gb:.2f} GB")
+        
+        logger.info(f"\nRecommendations:")
+        if total_activation_gb < 50:
+            logger.info("  ✓ Should fit on 8x A100 (80GB) GPUs")
+        elif total_activation_gb < 100:
+            logger.info("  ⚠ May need gradient checkpointing")
+        else:
+            logger.info("  ⚠ Will need aggressive optimizations (checkpointing, CPU offload)")
+        
+        return True
+    
+    def test_rope_frequencies(self):
+        """Test RoPE frequency calculations for long context"""
+        logger.info("\n" + "="*80)
+        logger.info("TEST 5: RoPE Frequency Analysis")
+        logger.info("="*80)
+        
+        rope_theta = getattr(self.config, 'rope_theta', 10000)
+        hidden_size = self.config.hidden_size
+        num_heads = self.config.num_attention_heads
+        head_dim = hidden_size // num_heads
+        
+        logger.info(f"RoPE theta: {rope_theta:,}")
+        logger.info(f"Head dimension: {head_dim}")
+        
+        # Calculate frequency range
+        # freqs = theta^(-2i/d) for i in [0, d/2]
+        min_freq = rope_theta ** (-2 * (head_dim-1) / head_dim)
+        max_freq = rope_theta ** 0
+        
+        logger.info(f"Frequency range: [{min_freq:.6f}, {max_freq:.6f}]")
+        
+        # Calculate wavelengths at different frequencies
+        wavelengths = [2 * 3.14159 / (rope_theta ** (-2 * i / head_dim)) 
+                      for i in range(0, head_dim // 2, head_dim // 8)]
+        
+        logger.info(f"\nWavelengths (in tokens):")
+        for i, wl in enumerate(wavelengths):
+            logger.info(f"  Frequency {i}: {wl:,.0f} tokens")
+        
+        max_wavelength = max(wavelengths)
+        if max_wavelength >= 250000:
+            logger.info(f"\n✓ Maximum wavelength ({max_wavelength:,.0f}) supports 250K context")
+        else:
+            logger.warning(f"\n⚠ Maximum wavelength ({max_wavelength:,.0f}) may be insufficient")
+        
+        return True
+    
+    def run_all_tests(self):
+        """Run all context length tests"""
+        logger.info("\n" + "="*80)
+        logger.info("HELION-OSC 250K CONTEXT LENGTH TEST SUITE")
+        logger.info("="*80)
+        
+        results = {
+            "tokenization": self.test_tokenization_capacity(),
+            "position_embeddings": self.test_position_embeddings(),
+            "attention_scaling": self.test_attention_computation(),
+            "memory_requirements": self.test_memory_requirements(),
+            "rope_frequencies": self.test_rope_frequencies()
+        }
+        
+        # Summary
+        logger.info("\n" + "="*80)
+        logger.info("TEST SUMMARY")
+        logger.info("="*80)
+        
+        for test_name, passed in results.items():
+            status = "✓ PASS" if passed else "✗ FAIL"
+            logger.info(f"{test_name}: {status}")
+        
+        all_passed = all(results.values())
+        
+        if all_passed:
+            logger.info("\n✓ All tests passed - Model supports 250K context length")
+        else:
+            logger.warning("\n⚠ Some tests failed - Check configuration")
+        
+        return all_passed
+
+
+def main():
+    """Main test script"""
+    import argparse
+    
+    parser = argparse.ArgumentParser(description="Test Helion-OSC 250K context support")
+    parser.add_argument(
+        "--model-path",
+        type=str,
+        default="./inference",
+        help="Path to model inference directory"
+    )
+    parser.add_argument(
+        "--test",
+        choices=["all", "tokenization", "position", "attention", "memory", "rope"],
+        default="all",
+        help="Which test to run"
+    )
+    
+    args = parser.parse_args()
+    
+    tester = LongContextTester(args.model_path)
+    
+    if args.test == "all":
+        tester.run_all_tests()
+    elif args.test == "tokenization":
+        tester.test_tokenization_capacity()
+    elif args.test == "position":
+        tester.test_position_embeddings()
+    elif args.test == "attention":
+        tester.test_attention_computation()
+    elif args.test == "memory":
+        tester.test_memory_requirements()
+    elif args.test == "rope":
+        tester.test_rope_frequencies()
+
+
+if __name__ == "__main__":
+    main()