🚀 Refined BitTransformerLM: Organized codebase with best practices

bit_transformer/BTLM_Extensions/rle_compression.py

@@ -0,0 +1,660 @@
+"""
+RLE Compression Extension for BitTransformerLM
+==============================================
+
+Advanced Run-Length Encoding compression module with multiple encoding schemes,
+adaptive compression, and training integration for BitTransformerLM.
+
+Key features:
+- Multiple RLE encoding schemes (basic, delta, hierarchical)
+- Adaptive compression with quality thresholds
+- Training integration with compression-aware loss
+- Batch processing and vectorized operations
+- Compatible with BitTransformerLM's training infrastructure
+"""
+
+import torch
+import torch.nn.functional as F
+from typing import List, Tuple, Optional, Dict, Any, Union
+import warnings
+import math
+from collections import defaultdict
+import numpy as np
+
+
+class RLEEncoder:
+    """
+    Advanced Run-Length Encoder with multiple encoding schemes.
+    
+    Supports:
+    - Basic RLE: (value, count) pairs
+    - Delta RLE: Differences between consecutive runs
+    - Hierarchical RLE: Multi-level compression
+    - Adaptive RLE: Chooses best scheme based on data
+    """
+    
+    def __init__(
+        self,
+        scheme: str = "adaptive",
+        min_run_length: int = 2,
+        max_value: int = 255,
+        delta_threshold: float = 0.7,
+        hierarchical_levels: int = 2,
+    ):
+        """
+        Args:
+            scheme: Encoding scheme ('basic', 'delta', 'hierarchical', 'adaptive')
+            min_run_length: Minimum run length to compress
+            max_value: Maximum value for encoding
+            delta_threshold: Compression ratio threshold for delta encoding
+            hierarchical_levels: Number of levels for hierarchical encoding
+        """
+        self.scheme = scheme
+        self.min_run_length = min_run_length
+        self.max_value = max_value
+        self.delta_threshold = delta_threshold
+        self.hierarchical_levels = hierarchical_levels
+        
+        self.stats = {
+            "total_compressions": 0,
+            "total_original_size": 0,
+            "total_compressed_size": 0,
+            "scheme_usage": defaultdict(int),
+        }
+    
+    def encode_basic_rle(self, data: torch.Tensor) -> torch.Tensor:
+        """Basic run-length encoding: (value, count) pairs."""
+        if data.numel() == 0:
+            return torch.tensor([], dtype=torch.uint8)
+            
+        data_flat = data.flatten()
+        encoded = []
+        
+        current_val = data_flat[0].item()
+        current_count = 1
+        
+        for i in range(1, len(data_flat)):
+            val = data_flat[i].item()
+            if val == current_val and current_count < 255:
+                current_count += 1
+            else:
+                if current_count >= self.min_run_length:
+                    encoded.extend([current_val, current_count])
+                else:
+                    # Store individual values for short runs
+                    for _ in range(current_count):
+                        encoded.append(current_val)
+                current_val = val
+                current_count = 1
+        
+        # Handle last run
+        if current_count >= self.min_run_length:
+            encoded.extend([current_val, current_count])
+        else:
+            for _ in range(current_count):
+                encoded.append(current_val)
+        
+        return torch.tensor(encoded, dtype=torch.uint8)
+    
+    def decode_basic_rle(self, encoded: torch.Tensor, target_length: Optional[int] = None) -> torch.Tensor:
+        """Decode basic run-length encoded data."""
+        if encoded.numel() == 0:
+            return torch.tensor([], dtype=torch.long)
+            
+        decoded = []
+        i = 0
+        
+        while i < len(encoded):
+            if i + 1 < len(encoded):
+                val = encoded[i].item()
+                count = encoded[i + 1].item()
+                
+                # Check if this looks like a (value, count) pair
+                if count > 1 and count <= 255:
+                    decoded.extend([val] * count)
+                    i += 2
+                else:
+                    # Individual value
+                    decoded.append(val)
+                    i += 1
+            else:
+                decoded.append(encoded[i].item())
+                i += 1
+        
+        result = torch.tensor(decoded, dtype=torch.long)
+        
+        # Trim or pad to target length if specified
+        if target_length is not None:
+            if len(result) > target_length:
+                result = result[:target_length]
+            elif len(result) < target_length:
+                result = F.pad(result, (0, target_length - len(result)))
+        
+        return result
+    
+    def encode_delta_rle(self, data: torch.Tensor) -> torch.Tensor:
+        """Delta run-length encoding: encode differences between values."""
+        if data.numel() <= 1:
+            return self.encode_basic_rle(data)
+            
+        data_flat = data.flatten()
+        
+        # Compute deltas
+        deltas = torch.diff(data_flat, prepend=data_flat[0:1])
+        
+        # Apply basic RLE to deltas (shifted to handle negatives)
+        shifted_deltas = deltas + 128  # Shift to 0-255 range
+        shifted_deltas = torch.clamp(shifted_deltas, 0, 255)
+        
+        delta_encoded = self.encode_basic_rle(shifted_deltas)
+        
+        # Prepend original first value
+        result = torch.cat([data_flat[0:1].to(torch.uint8), delta_encoded])
+        return result
+    
+    def decode_delta_rle(self, encoded: torch.Tensor, target_length: Optional[int] = None) -> torch.Tensor:
+        """Decode delta run-length encoded data."""
+        if encoded.numel() <= 1:
+            return self.decode_basic_rle(encoded, target_length)
+            
+        # First value is the original value
+        first_val = encoded[0].item()
+        delta_encoded = encoded[1:]
+        
+        # Decode deltas
+        deltas = self.decode_basic_rle(delta_encoded)
+        
+        # Unshift deltas
+        deltas = deltas.float() - 128
+        
+        # Reconstruct original sequence
+        if deltas.numel() > 0:
+            deltas[0] = first_val  # Replace first delta with original value
+            result = torch.cumsum(deltas, dim=0).long()
+        else:
+            result = torch.tensor([first_val], dtype=torch.long)
+        
+        # Trim or pad to target length
+        if target_length is not None:
+            if len(result) > target_length:
+                result = result[:target_length]
+            elif len(result) < target_length:
+                result = F.pad(result, (0, target_length - len(result)))
+        
+        return result
+    
+    def encode_hierarchical_rle(self, data: torch.Tensor) -> torch.Tensor:
+        """Hierarchical RLE: Apply RLE recursively for better compression."""
+        current_data = data.clone()
+        
+        for level in range(self.hierarchical_levels):
+            encoded = self.encode_basic_rle(current_data)
+            
+            # Check if compression is beneficial
+            if encoded.numel() >= current_data.numel() * 0.9:
+                # Compression not beneficial, return previous level
+                break
+                
+            current_data = encoded
+        
+        return current_data
+    
+    def decode_hierarchical_rle(self, encoded: torch.Tensor, target_length: Optional[int] = None, levels: int = None) -> torch.Tensor:
+        """Decode hierarchical RLE data."""
+        if levels is None:
+            levels = self.hierarchical_levels
+            
+        current_data = encoded.clone()
+        
+        for level in range(levels):
+            try:
+                current_data = self.decode_basic_rle(current_data)
+            except Exception:
+                # If decoding fails, return current state
+                break
+        
+        # Final length adjustment
+        if target_length is not None and current_data.numel() != target_length:
+            if current_data.numel() > target_length:
+                current_data = current_data[:target_length]
+            else:
+                current_data = F.pad(current_data, (0, target_length - current_data.numel()))
+        
+        return current_data
+    
+    def encode(self, data: torch.Tensor) -> Tuple[torch.Tensor, Dict[str, Any]]:
+        """
+        Encode data using the configured scheme.
+        
+        Args:
+            data: Input tensor to compress
+            
+        Returns:
+            Tuple of (encoded_data, metadata)
+        """
+        original_shape = data.shape
+        original_size = data.numel()
+        
+        if self.scheme == "basic":
+            encoded = self.encode_basic_rle(data)
+            scheme_used = "basic"
+        elif self.scheme == "delta":
+            encoded = self.encode_delta_rle(data)
+            scheme_used = "delta"
+        elif self.scheme == "hierarchical":
+            encoded = self.encode_hierarchical_rle(data)
+            scheme_used = "hierarchical"
+        elif self.scheme == "adaptive":
+            # Try all schemes and pick the best one
+            basic_encoded = self.encode_basic_rle(data)
+            delta_encoded = self.encode_delta_rle(data)
+            hierarchical_encoded = self.encode_hierarchical_rle(data)
+            
+            candidates = {
+                "basic": basic_encoded,
+                "delta": delta_encoded,
+                "hierarchical": hierarchical_encoded,
+            }
+            
+            # Choose scheme with best compression ratio
+            best_scheme = min(candidates.keys(), key=lambda k: candidates[k].numel())
+            encoded = candidates[best_scheme]
+            scheme_used = best_scheme
+        else:
+            raise ValueError(f"Unknown encoding scheme: {self.scheme}")
+        
+        # Update statistics
+        self.stats["total_compressions"] += 1
+        self.stats["total_original_size"] += original_size
+        self.stats["total_compressed_size"] += encoded.numel()
+        self.stats["scheme_usage"][scheme_used] += 1
+        
+        metadata = {
+            "scheme": scheme_used,
+            "original_shape": original_shape,
+            "original_size": original_size,
+            "compressed_size": encoded.numel(),
+            "compression_ratio": encoded.numel() / original_size if original_size > 0 else 1.0,
+        }
+        
+        return encoded, metadata
+    
+    def decode(self, encoded: torch.Tensor, metadata: Dict[str, Any]) -> torch.Tensor:
+        """
+        Decode compressed data using metadata.
+        
+        Args:
+            encoded: Compressed data
+            metadata: Metadata from encoding
+            
+        Returns:
+            Decoded tensor
+        """
+        scheme = metadata["scheme"]
+        original_shape = metadata["original_shape"]
+        target_length = math.prod(original_shape) if original_shape else None
+        
+        if scheme == "basic":
+            decoded = self.decode_basic_rle(encoded, target_length)
+        elif scheme == "delta":
+            decoded = self.decode_delta_rle(encoded, target_length)
+        elif scheme == "hierarchical":
+            decoded = self.decode_hierarchical_rle(encoded, target_length)
+        else:
+            raise ValueError(f"Unknown decoding scheme: {scheme}")
+        
+        # Reshape to original shape
+        if original_shape and decoded.numel() >= math.prod(original_shape):
+            decoded = decoded[:math.prod(original_shape)].reshape(original_shape)
+        
+        return decoded
+    
+    def get_compression_stats(self) -> Dict[str, float]:
+        """Get compression statistics."""
+        if self.stats["total_original_size"] == 0:
+            return {"average_compression_ratio": 1.0, "total_savings": 0.0}
+            
+        avg_ratio = self.stats["total_compressed_size"] / self.stats["total_original_size"]
+        total_savings = self.stats["total_original_size"] - self.stats["total_compressed_size"]
+        
+        return {
+            "average_compression_ratio": avg_ratio,
+            "total_savings": total_savings,
+            "total_compressions": self.stats["total_compressions"],
+            "scheme_usage": dict(self.stats["scheme_usage"]),
+        }
+
+
+class CompressedBitDataset(torch.utils.data.Dataset):
+    """
+    Dataset wrapper that applies RLE compression on-the-fly during training.
+    
+    This allows for memory-efficient storage of large bit sequences while
+    maintaining fast access during training.
+    """
+    
+    def __init__(
+        self,
+        data: torch.Tensor,
+        encoder: RLEEncoder,
+        compress_probability: float = 0.5,
+        cache_size: int = 1000,
+    ):
+        """
+        Args:
+            data: Original bit sequence data
+            encoder: RLE encoder instance
+            compress_probability: Probability of returning compressed data
+            cache_size: Number of compressed items to cache
+        """
+        self.data = data
+        self.encoder = encoder
+        self.compress_probability = compress_probability
+        self.cache_size = cache_size
+        self.cache = {}
+        self.access_count = defaultdict(int)
+        
+    def __len__(self):
+        return len(self.data)
+    
+    def __getitem__(self, idx: int) -> Tuple[torch.Tensor, Dict[str, Any]]:
+        """
+        Get item with optional compression.
+        
+        Returns:
+            Tuple of (data, metadata) where metadata indicates if compressed
+        """
+        original_item = self.data[idx]
+        
+        # Randomly decide whether to compress
+        if torch.rand(1).item() < self.compress_probability:
+            # Check cache first
+            if idx in self.cache:
+                compressed, metadata = self.cache[idx]
+                self.access_count[idx] += 1
+                metadata["from_cache"] = True
+                return compressed, metadata
+            
+            # Compress item
+            compressed, metadata = self.encoder.encode(original_item)
+            
+            # Add to cache if there's room
+            if len(self.cache) < self.cache_size:
+                self.cache[idx] = (compressed, metadata)
+            elif self.access_count:
+                # Replace least accessed item
+                least_accessed = min(self.cache.keys(), key=lambda k: self.access_count[k])
+                del self.cache[least_accessed]
+                del self.access_count[least_accessed]
+                self.cache[idx] = (compressed, metadata)
+            
+            metadata["from_cache"] = False
+            return compressed, metadata
+        else:
+            # Return original data
+            metadata = {
+                "scheme": "uncompressed",
+                "original_shape": original_item.shape,
+                "compressed": False,
+                "from_cache": False,
+            }
+            return original_item, metadata
+
+
+def create_compression_aware_loss(
+    base_loss_fn,
+    compression_penalty: float = 0.01,
+    quality_threshold: float = 0.8,
+) -> callable:
+    """
+    Create a loss function that penalizes poor compression quality.
+    
+    Args:
+        base_loss_fn: Base loss function (e.g., CrossEntropyLoss)
+        compression_penalty: Penalty weight for compression artifacts
+        quality_threshold: Minimum compression quality threshold
+        
+    Returns:
+        Compression-aware loss function
+    """
+    def compression_aware_loss(
+        logits: torch.Tensor,
+        targets: torch.Tensor,
+        metadata_batch: Optional[List[Dict[str, Any]]] = None,
+    ) -> torch.Tensor:
+        """
+        Compute loss with compression quality penalty.
+        
+        Args:
+            logits: Model output logits
+            targets: Target labels
+            metadata_batch: Batch of compression metadata
+            
+        Returns:
+            Adjusted loss tensor
+        """
+        base_loss = base_loss_fn(logits, targets)
+        
+        if metadata_batch is None:
+            return base_loss
+        
+        # Compute compression quality penalty
+        penalty = 0.0
+        compressed_items = 0
+        
+        for metadata in metadata_batch:
+            if metadata.get("compressed", False):
+                compressed_items += 1
+                compression_ratio = metadata.get("compression_ratio", 1.0)
+                
+                # Penalty for poor compression
+                if compression_ratio > quality_threshold:
+                    quality_penalty = (compression_ratio - quality_threshold) ** 2
+                    penalty += quality_penalty
+        
+        if compressed_items > 0:
+            penalty = penalty / compressed_items  # Average penalty
+            total_loss = base_loss + compression_penalty * penalty
+        else:
+            total_loss = base_loss
+        
+        return total_loss
+    
+    return compression_aware_loss
+
+
+def integrate_rle_with_training(
+    model,
+    data: torch.Tensor,
+    encoder_config: Optional[Dict[str, Any]] = None,
+    compression_config: Optional[Dict[str, Any]] = None,
+) -> Tuple[CompressedBitDataset, callable]:
+    """
+    Integrate RLE compression with BitTransformerLM training.
+    
+    Args:
+        model: BitTransformerLM model
+        data: Training data tensor
+        encoder_config: Configuration for RLE encoder
+        compression_config: Configuration for compression-aware training
+        
+    Returns:
+        Tuple of (compressed_dataset, compression_aware_loss_fn)
+    """
+    # Default configurations
+    if encoder_config is None:
+        encoder_config = {
+            "scheme": "adaptive",
+            "min_run_length": 2,
+            "delta_threshold": 0.7,
+        }
+    
+    if compression_config is None:
+        compression_config = {
+            "compress_probability": 0.3,
+            "compression_penalty": 0.01,
+            "quality_threshold": 0.8,
+            "cache_size": 1000,
+        }
+    
+    # Create encoder and dataset
+    encoder = RLEEncoder(**encoder_config)
+    dataset = CompressedBitDataset(
+        data,
+        encoder,
+        compress_probability=compression_config["compress_probability"],
+        cache_size=compression_config["cache_size"],
+    )
+    
+    # Create compression-aware loss
+    base_loss = torch.nn.CrossEntropyLoss()
+    loss_fn = create_compression_aware_loss(
+        base_loss,
+        compression_penalty=compression_config["compression_penalty"],
+        quality_threshold=compression_config["quality_threshold"],
+    )
+    
+    return dataset, loss_fn
+
+
+def benchmark_compression_schemes(
+    test_data: torch.Tensor,
+    schemes: List[str] = ["basic", "delta", "hierarchical", "adaptive"],
+) -> Dict[str, Dict[str, float]]:
+    """
+    Benchmark different compression schemes on test data.
+    
+    Args:
+        test_data: Test data tensor
+        schemes: List of schemes to test
+        
+    Returns:
+        Dictionary with benchmark results for each scheme
+    """
+    results = {}
+    
+    for scheme in schemes:
+        encoder = RLEEncoder(scheme=scheme)
+        
+        # Test compression/decompression
+        try:
+            compressed, metadata = encoder.encode(test_data)
+            reconstructed = encoder.decode(compressed, metadata)
+            
+            # Compute metrics
+            compression_ratio = compressed.numel() / test_data.numel()
+            reconstruction_error = torch.mean((test_data.float() - reconstructed.float()) ** 2).item()
+            
+            results[scheme] = {
+                "compression_ratio": compression_ratio,
+                "reconstruction_error": reconstruction_error,
+                "compressed_size": compressed.numel(),
+                "original_size": test_data.numel(),
+                "success": True,
+            }
+        except Exception as e:
+            results[scheme] = {
+                "compression_ratio": 1.0,
+                "reconstruction_error": float("inf"),
+                "compressed_size": test_data.numel(),
+                "original_size": test_data.numel(),
+                "success": False,
+                "error": str(e),
+            }
+    
+    return results
+
+
+# Example usage and utilities
+def create_rle_training_config(
+    scheme: str = "adaptive",
+    compress_probability: float = 0.3,
+    compression_penalty: float = 0.01,
+    **kwargs
+) -> Dict[str, Any]:
+    """
+    Create configuration for RLE-enhanced training.
+    
+    Args:
+        scheme: RLE encoding scheme
+        compress_probability: Probability of compression during training
+        compression_penalty: Loss penalty for compression artifacts
+        **kwargs: Additional configuration options
+        
+    Returns:
+        Dictionary with RLE training configuration
+    """
+    config = {
+        "compression_type": "rle",
+        "encoder_config": {
+            "scheme": scheme,
+            "min_run_length": kwargs.get("min_run_length", 2),
+            "delta_threshold": kwargs.get("delta_threshold", 0.7),
+            "hierarchical_levels": kwargs.get("hierarchical_levels", 2),
+        },
+        "training_config": {
+            "compress_probability": compress_probability,
+            "compression_penalty": compression_penalty,
+            "quality_threshold": kwargs.get("quality_threshold", 0.8),
+            "cache_size": kwargs.get("cache_size", 1000),
+        },
+    }
+    
+    return config
+
+
+if __name__ == "__main__":
+    # Test the RLE compression module
+    print("Testing RLE Compression Module...")
+    
+    # Create test data
+    test_data = torch.randint(0, 2, (100,))
+    
+    # Add some runs for better compression
+    test_data[20:30] = 1
+    test_data[50:70] = 0
+    test_data[80:90] = 1
+    
+    print(f"Original data shape: {test_data.shape}")
+    print(f"Original data: {test_data[:20]}...")
+    
+    # Test different encoding schemes
+    schemes = ["basic", "delta", "hierarchical", "adaptive"]
+    
+    for scheme in schemes:
+        print(f"\nTesting {scheme} scheme:")
+        encoder = RLEEncoder(scheme=scheme)
+        
+        try:
+            # Encode
+            compressed, metadata = encoder.encode(test_data)
+            print(f"  Compressed size: {compressed.numel()}")
+            print(f"  Compression ratio: {metadata['compression_ratio']:.3f}")
+            
+            # Decode
+            reconstructed = encoder.decode(compressed, metadata)
+            
+            # Check reconstruction quality
+            error = torch.mean((test_data.float() - reconstructed.float()) ** 2)
+            print(f"  Reconstruction error: {error.item():.6f}")
+            
+            if error.item() < 1e-6:
+                print("  ✅ Perfect reconstruction")
+            else:
+                print("  ❌ Reconstruction error detected")
+                
+        except Exception as e:
+            print(f"  ❌ Error: {e}")
+    
+    # Benchmark all schemes
+    print("\nBenchmarking compression schemes...")
+    benchmark_results = benchmark_compression_schemes(test_data)
+    
+    for scheme, results in benchmark_results.items():
+        if results["success"]:
+            print(f"{scheme:12}: ratio={results['compression_ratio']:.3f}, "
+                  f"error={results['reconstruction_error']:.6f}")
+        else:
+            print(f"{scheme:12}: FAILED - {results.get('error', 'Unknown error')}")
+    
+    print("\nRLE Compression Module test completed!")