520 lines
17 KiB
Rust
520 lines
17 KiB
Rust
#![cfg(feature = "disabled_tests")]
|
|
|
|
use anyhow::Result;
|
|
use rtx_compress::{
|
|
CompressedStorage, CompressionConfig,
|
|
checkpoint::CheckpointCompressor,
|
|
kv_cache::CompressedKVCache,
|
|
quantization::{
|
|
mixed_precision::MixedPrecisionOptimizer, product_quantization::ProductQuantizer,
|
|
},
|
|
};
|
|
use rtx_inference::{InferenceEngine, InferenceRequest};
|
|
use rtx_tensor::{DType, Device, Tensor};
|
|
use std::collections::HashMap;
|
|
|
|
#[test]
|
|
fn test_end_to_end_model_compression_pipeline() -> Result<()> {
|
|
let device = Device::try_default()?;
|
|
|
|
// Step 1: Create a mock transformer model
|
|
let mut model_weights = HashMap::new();
|
|
for layer in 0..6 {
|
|
// 6-layer transformer
|
|
model_weights.insert(
|
|
format!("layers.{}.attention.query", layer),
|
|
Tensor::randn(&[512, 512], &device)?,
|
|
);
|
|
model_weights.insert(
|
|
format!("layers.{}.attention.key", layer),
|
|
Tensor::randn(&[512, 512], &device)?,
|
|
);
|
|
model_weights.insert(
|
|
format!("layers.{}.attention.value", layer),
|
|
Tensor::randn(&[512, 512], &device)?,
|
|
);
|
|
model_weights.insert(
|
|
format!("layers.{}.mlp.up_proj", layer),
|
|
Tensor::randn(&[512, 2048], &device)?,
|
|
);
|
|
model_weights.insert(
|
|
format!("layers.{}.mlp.down_proj", layer),
|
|
Tensor::randn(&[2048, 512], &device)?,
|
|
);
|
|
}
|
|
|
|
// Step 2: Apply mixed precision optimization
|
|
let precision_config = rtx_compress::quantization::mixed_precision::PrecisionConfig {
|
|
precision_bits: vec![4, 8, 12, 16],
|
|
sensitivity_threshold: 0.02,
|
|
performance_weight: 0.6,
|
|
quality_weight: 0.4,
|
|
};
|
|
|
|
let mut precision_optimizer = MixedPrecisionOptimizer::new(precision_config);
|
|
let calibration_data = Tensor::randn(&[100, 512], &device)?;
|
|
|
|
let objective = rtx_compress::quantization::mixed_precision::OptimizationObjective {
|
|
target_compression_ratio: 4.0,
|
|
max_quality_loss: 0.05,
|
|
memory_constraint_mb: Some(200),
|
|
};
|
|
|
|
let optimal_precision =
|
|
precision_optimizer.optimize(&model_weights, &calibration_data, objective)?;
|
|
|
|
// Step 3: Compress model checkpoints
|
|
let checkpoint_config = rtx_compress::checkpoint::CompressionConfig {
|
|
format: rtx_compress::checkpoint::CheckpointFormat::Zstd,
|
|
compression_level: 6,
|
|
quantization_bits: 0, // Use mixed precision
|
|
exclude_patterns: vec!["*layer_norm*".to_string()],
|
|
};
|
|
|
|
let mut checkpoint_compressor = CheckpointCompressor::new(checkpoint_config);
|
|
|
|
// Apply optimal precisions
|
|
for (layer, bits) in &optimal_precision.layer_precisions {
|
|
checkpoint_compressor.add_quantization_rule(layer, *bits);
|
|
}
|
|
|
|
let compressed_model = checkpoint_compressor.save(&model_weights)?;
|
|
|
|
// Step 4: Set up compressed KV cache for inference
|
|
let kv_config = rtx_compress::kv_cache::KVCacheConfig {
|
|
compression_method: rtx_compress::kv_cache::CompressionMethod::ProductQuantization {
|
|
num_subquantizers: 8,
|
|
codebook_size: 256,
|
|
use_opq: true,
|
|
},
|
|
compression_ratio_target: 3.0,
|
|
quality_threshold: 0.92,
|
|
max_cache_size_mb: 1024,
|
|
enable_sliding_window: true,
|
|
window_size: 2048,
|
|
enable_attention_scoring: true,
|
|
prefetch_batch_size: 32,
|
|
enable_auto_tuning: false,
|
|
};
|
|
|
|
let mut kv_cache = CompressedKVCache::new(kv_config)?;
|
|
|
|
// Step 5: Verify full pipeline works
|
|
let loaded_model = checkpoint_compressor.load(&compressed_model)?;
|
|
assert_eq!(loaded_model.len(), model_weights.len());
|
|
|
|
// Simulate inference with compressed KV cache
|
|
for seq_id in 0..5 {
|
|
let keys = Tensor::randn(&[1, 8, 256, 64], &device)?;
|
|
let values = Tensor::randn(&[1, 8, 256, 64], &device)?;
|
|
|
|
kv_cache.insert(seq_id, &keys, &values)?;
|
|
}
|
|
|
|
// Verify compression metrics
|
|
let compression_stats = kv_cache.compression_stats();
|
|
assert!(compression_stats.compression_ratio >= 2.5);
|
|
|
|
let original_model_size = calculate_model_size(&model_weights);
|
|
let compressed_model_size = compressed_model.len();
|
|
let model_compression_ratio = original_model_size as f64 / compressed_model_size as f64;
|
|
|
|
assert!(
|
|
model_compression_ratio >= 3.0,
|
|
"Model compression should achieve 3x ratio, got {:.2}",
|
|
model_compression_ratio
|
|
);
|
|
|
|
println!("Pipeline results:");
|
|
println!(" Model compression: {:.2}x", model_compression_ratio);
|
|
println!(
|
|
" KV cache compression: {:.2}x",
|
|
compression_stats.compression_ratio
|
|
);
|
|
println!(
|
|
" Total memory savings: ~{:.1}%",
|
|
(1.0 - 1.0 / model_compression_ratio) * 100.0
|
|
);
|
|
|
|
Ok(())
|
|
}
|
|
|
|
#[test]
|
|
fn test_compressed_inference_integration() -> Result<()> {
|
|
let device = Device::try_default()?;
|
|
|
|
// Create compressed storage system
|
|
let storage_config = CompressionConfig {
|
|
default_compression_ratio: 4.0,
|
|
quality_threshold: 0.90,
|
|
memory_limit_mb: 1000,
|
|
adaptive_compression: true,
|
|
};
|
|
|
|
let mut storage = CompressedStorage::new(storage_config);
|
|
|
|
// Store model layers with different compression strategies
|
|
let attention_weights = Tensor::randn(&[768, 768], &device)?;
|
|
let embedding_weights = Tensor::randn(&[50000, 768], &device)?;
|
|
let mlp_weights = Tensor::randn(&[768, 3072], &device)?;
|
|
|
|
// High precision for attention (critical for quality)
|
|
storage.store_tensor(
|
|
"attention.query",
|
|
&attention_weights,
|
|
rtx_compress::CompressionLevel::Low,
|
|
)?;
|
|
|
|
// Aggressive compression for embeddings (large but redundant)
|
|
storage.store_tensor(
|
|
"embeddings.weight",
|
|
&embedding_weights,
|
|
rtx_compress::CompressionLevel::High,
|
|
)?;
|
|
|
|
// Balanced compression for MLP
|
|
storage.store_tensor(
|
|
"mlp.up_proj",
|
|
&mlp_weights,
|
|
rtx_compress::CompressionLevel::Medium,
|
|
)?;
|
|
|
|
// Simulate inference requests
|
|
let batch_size = 4;
|
|
let seq_len = 512;
|
|
|
|
for request_id in 0..10 {
|
|
// Load weights on demand (decompression)
|
|
let attention_loaded = storage.load_tensor("attention.query")?;
|
|
let embeddings_loaded = storage.load_tensor("embeddings.weight")?;
|
|
let mlp_loaded = storage.load_tensor("mlp.up_proj")?;
|
|
|
|
// Verify shapes are preserved
|
|
assert_eq!(attention_loaded.shape(), attention_weights.shape());
|
|
assert_eq!(embeddings_loaded.shape(), embedding_weights.shape());
|
|
assert_eq!(mlp_loaded.shape(), mlp_weights.shape());
|
|
|
|
// Simulate forward pass computation
|
|
let input_ids = Tensor::randint(0, 50000, &[batch_size, seq_len], &device)?;
|
|
let embeddings = embeddings_loaded.gather(&input_ids, 0)?;
|
|
|
|
// Mock attention computation
|
|
let queries = embeddings.matmul(&attention_loaded)?;
|
|
|
|
// Verify computation produces reasonable results
|
|
assert_eq!(queries.shape(), &[batch_size, seq_len, 768]);
|
|
|
|
println!(
|
|
"Request {}: processed batch_size={}, seq_len={}",
|
|
request_id, batch_size, seq_len
|
|
);
|
|
}
|
|
|
|
// Verify storage efficiency
|
|
let storage_stats = storage.get_statistics();
|
|
println!("Storage statistics:");
|
|
println!(
|
|
" Total compressed size: {} MB",
|
|
storage_stats.total_compressed_size_mb
|
|
);
|
|
println!(
|
|
" Average compression ratio: {:.2}x",
|
|
storage_stats.average_compression_ratio
|
|
);
|
|
println!(
|
|
" Memory usage: {:.1}%",
|
|
storage_stats.memory_utilization * 100.0
|
|
);
|
|
|
|
assert!(storage_stats.average_compression_ratio >= 2.0);
|
|
assert!(storage_stats.memory_utilization <= 1.0);
|
|
|
|
Ok(())
|
|
}
|
|
|
|
#[test]
|
|
fn test_multi_model_compression_sharing() -> Result<()> {
|
|
let device = Device::try_default()?;
|
|
|
|
// Simulate multiple model variants that share common layers
|
|
let shared_embeddings = Tensor::randn(&[30000, 512], &device)?;
|
|
let shared_attention = Tensor::randn(&[512, 512], &device)?;
|
|
|
|
// Model A: Base model
|
|
let mut model_a = HashMap::new();
|
|
model_a.insert("embeddings".to_string(), shared_embeddings.clone());
|
|
model_a.insert("attention".to_string(), shared_attention.clone());
|
|
model_a.insert(
|
|
"classifier".to_string(),
|
|
Tensor::randn(&[512, 1000], &device)?,
|
|
);
|
|
|
|
// Model B: Fine-tuned variant (shares embeddings and attention)
|
|
let mut model_b = HashMap::new();
|
|
model_b.insert("embeddings".to_string(), shared_embeddings.clone());
|
|
model_b.insert("attention".to_string(), shared_attention.clone());
|
|
model_b.insert(
|
|
"classifier".to_string(),
|
|
Tensor::randn(&[512, 100], &device)?,
|
|
);
|
|
|
|
let config = rtx_compress::checkpoint::CompressionConfig {
|
|
format: rtx_compress::checkpoint::CheckpointFormat::Zstd,
|
|
compression_level: 5,
|
|
quantization_bits: 8,
|
|
exclude_patterns: vec![],
|
|
};
|
|
|
|
let compressor = CheckpointCompressor::new(config);
|
|
|
|
// Save models with deduplication
|
|
let compressed_a = compressor.save(&model_a)?;
|
|
let compressed_b_incremental = compressor.save_with_base(&model_b, &model_a)?;
|
|
|
|
// Verify incremental compression is much smaller
|
|
let compressed_b_full = compressor.save(&model_b)?;
|
|
|
|
assert!(
|
|
compressed_b_incremental.len() < compressed_b_full.len() / 2,
|
|
"Incremental compression should be much smaller"
|
|
);
|
|
|
|
// Verify both models can be loaded correctly
|
|
let loaded_a = compressor.load(&compressed_a)?;
|
|
let loaded_b = compressor.load_with_base(&compressed_b_incremental, &model_a)?;
|
|
|
|
// Verify shared layers are identical
|
|
let embeddings_diff = (&loaded_a["embeddings"] - &loaded_b["embeddings"])?
|
|
.abs()?
|
|
.max()?
|
|
.to_scalar::<f32>()?;
|
|
assert!(
|
|
embeddings_diff < 1e-6,
|
|
"Shared embeddings should be identical"
|
|
);
|
|
|
|
let attention_diff = (&loaded_a["attention"] - &loaded_b["attention"])?
|
|
.abs()?
|
|
.max()?
|
|
.to_scalar::<f32>()?;
|
|
assert!(
|
|
attention_diff < 1e-6,
|
|
"Shared attention should be identical"
|
|
);
|
|
|
|
// But classifiers should be different
|
|
let classifier_diff = (&loaded_a["classifier"] - &loaded_b["classifier"])?
|
|
.abs()?
|
|
.mean(&[], false)?
|
|
.to_scalar::<f32>()?;
|
|
assert!(classifier_diff > 0.1, "Classifiers should be different");
|
|
|
|
println!("Model sharing results:");
|
|
println!(" Model A size: {} bytes", compressed_a.len());
|
|
println!(" Model B (full): {} bytes", compressed_b_full.len());
|
|
println!(
|
|
" Model B (incremental): {} bytes",
|
|
compressed_b_incremental.len()
|
|
);
|
|
println!(
|
|
" Sharing savings: {:.1}%",
|
|
(1.0 - compressed_b_incremental.len() as f64 / compressed_b_full.len() as f64) * 100.0
|
|
);
|
|
|
|
Ok(())
|
|
}
|
|
|
|
#[test]
|
|
fn test_adaptive_compression_under_load() -> Result<()> {
|
|
let device = Device::try_default()?;
|
|
|
|
let storage_config = CompressionConfig {
|
|
default_compression_ratio: 3.0,
|
|
quality_threshold: 0.88,
|
|
memory_limit_mb: 500, // Constrained memory
|
|
adaptive_compression: true,
|
|
};
|
|
|
|
let mut storage = CompressedStorage::new(storage_config);
|
|
storage.enable_load_monitoring(true);
|
|
|
|
// Store many tensors to exceed memory limit
|
|
let mut tensors = HashMap::new();
|
|
for i in 0..100 {
|
|
let tensor_name = format!("layer_{}", i);
|
|
let tensor = Tensor::randn(&[1024, 1024], &device)?;
|
|
tensors.insert(tensor_name.clone(), tensor.clone());
|
|
|
|
// Simulate varying access patterns
|
|
let access_frequency = if i < 20 {
|
|
rtx_compress::AccessPattern::VeryHigh
|
|
} else if i < 50 {
|
|
rtx_compress::AccessPattern::Medium
|
|
} else {
|
|
rtx_compress::AccessPattern::Low
|
|
};
|
|
|
|
storage.store_tensor_with_hint(&tensor_name, &tensor, access_frequency)?;
|
|
}
|
|
|
|
// Simulate high load with concurrent access
|
|
let mut handles = vec![];
|
|
for thread_id in 0..4 {
|
|
let storage_clone = storage.clone();
|
|
let handle = std::thread::spawn(move || -> Result<()> {
|
|
for request in 0..25 {
|
|
let tensor_idx = (thread_id * 25 + request) % 100;
|
|
let tensor_name = format!("layer_{}", tensor_idx);
|
|
|
|
// Load tensor (may trigger adaptive recompression)
|
|
let loaded = storage_clone.load_tensor(&tensor_name)?;
|
|
|
|
// Simulate computation
|
|
let result = loaded.sum(Some(0))?;
|
|
let scalar = result.to_vec::<f32>()?;
|
|
assert!(scalar[0].is_finite());
|
|
}
|
|
Ok(())
|
|
});
|
|
handles.push(handle);
|
|
}
|
|
|
|
// Wait for all threads
|
|
for handle in handles {
|
|
handle.join().unwrap()?;
|
|
}
|
|
|
|
// Verify adaptive behavior
|
|
let final_stats = storage.get_statistics();
|
|
println!("Adaptive compression results:");
|
|
println!(
|
|
" Memory usage: {:.1}%",
|
|
final_stats.memory_utilization * 100.0
|
|
);
|
|
println!(
|
|
" Average compression ratio: {:.2}x",
|
|
final_stats.average_compression_ratio
|
|
);
|
|
println!(" Cache hits: {:.1}%", final_stats.cache_hit_rate * 100.0);
|
|
println!(" Evictions: {}", final_stats.eviction_count);
|
|
|
|
// Should stay within memory limits
|
|
assert!(
|
|
final_stats.memory_utilization <= 1.0,
|
|
"Should not exceed memory limit"
|
|
);
|
|
|
|
// Should achieve reasonable cache performance under load
|
|
assert!(
|
|
final_stats.cache_hit_rate >= 0.7,
|
|
"Should achieve reasonable cache hit rate under load"
|
|
);
|
|
|
|
Ok(())
|
|
}
|
|
|
|
#[test]
|
|
fn test_compression_quality_vs_performance_tradeoffs() -> Result<()> {
|
|
let device = Device::try_default()?;
|
|
|
|
// Create test tensor with known characteristics
|
|
let test_tensor = Tensor::randn(&[2048, 2048], &device)?;
|
|
|
|
let compression_levels = vec![
|
|
(
|
|
"ultra_fast",
|
|
rtx_compress::CompressionLevel::UltraFast,
|
|
10.0,
|
|
), // 10ms target
|
|
("fast", rtx_compress::CompressionLevel::Fast, 50.0), // 50ms target
|
|
("balanced", rtx_compress::CompressionLevel::Balanced, 200.0), // 200ms target
|
|
("high", rtx_compress::CompressionLevel::High, 1000.0), // 1s target
|
|
];
|
|
|
|
let mut results = vec![];
|
|
|
|
for (name, level, time_budget_ms) in compression_levels {
|
|
let config = CompressionConfig {
|
|
default_compression_ratio: 0.0, // Auto-determine based on level
|
|
quality_threshold: 0.0, // Auto-determine
|
|
memory_limit_mb: 1000,
|
|
adaptive_compression: false,
|
|
};
|
|
|
|
let mut storage = CompressedStorage::new(config);
|
|
|
|
// Measure compression time
|
|
let start = std::time::Instant::now();
|
|
storage.store_tensor_with_level("test", &test_tensor, level)?;
|
|
let compression_time = start.elapsed();
|
|
|
|
// Measure decompression time
|
|
let start = std::time::Instant::now();
|
|
let decompressed = storage.load_tensor("test")?;
|
|
let decompression_time = start.elapsed();
|
|
|
|
// Measure quality
|
|
let mse = (&test_tensor - &decompressed)?
|
|
.pow(2.0)?
|
|
.mean(&[], false)?
|
|
.to_scalar::<f32>()?;
|
|
|
|
// Measure compression ratio
|
|
let original_size = test_tensor.numel() * 4; // f32
|
|
let stats = storage.get_tensor_stats("test")?;
|
|
let compression_ratio = original_size as f64 / stats.compressed_size as f64;
|
|
|
|
results.push((
|
|
name,
|
|
compression_time.as_millis() as f64,
|
|
decompression_time.as_millis() as f64,
|
|
compression_ratio,
|
|
mse,
|
|
));
|
|
|
|
println!(
|
|
"Level {}: compress={:.1}ms, decompress={:.1}ms, ratio={:.2}x, mse={:.6}",
|
|
name,
|
|
compression_time.as_millis() as f64,
|
|
decompression_time.as_millis() as f64,
|
|
compression_ratio,
|
|
mse
|
|
);
|
|
|
|
// Verify time budget is respected (with 50% margin)
|
|
assert!(compression_time.as_millis() as f64 < time_budget_ms * 1.5,
|
|
"Compression time should respect budget");
|
|
}
|
|
|
|
// Verify tradeoff relationships
|
|
// Higher compression levels should generally:
|
|
// 1. Take more time
|
|
// 2. Achieve better compression ratios
|
|
// 3. Have better quality (lower MSE)
|
|
|
|
for i in 1..results.len() {
|
|
let (prev_name, prev_comp_time, _, prev_ratio, prev_mse) = results[i - 1];
|
|
let (curr_name, curr_comp_time, _, curr_ratio, curr_mse) = results[i];
|
|
|
|
// Allow some variance due to algorithm differences
|
|
println!("Comparing {} vs {}", prev_name, curr_name);
|
|
|
|
// Generally expect better compression ratios at higher levels
|
|
if curr_ratio < prev_ratio * 0.8 {
|
|
println!(
|
|
"Warning: {} has lower compression ratio than {}",
|
|
curr_name, prev_name
|
|
);
|
|
}
|
|
}
|
|
|
|
Ok(())
|
|
}
|
|
|
|
fn calculate_model_size(model: &HashMap<String, Tensor>) -> usize {
|
|
model
|
|
.values()
|
|
.map(|t| t.numel() * t.dtype().size_in_bytes())
|
|
.sum()
|
|
}
|