465 lines
15 KiB
Rust
465 lines
15 KiB
Rust
use criterion::{BenchmarkId, Criterion, black_box, criterion_group, criterion_main};
|
|
use rtx_compress::{
|
|
checkpoint::{CheckpointCompressor, CheckpointFormat, CompressionConfig},
|
|
kv_cache::{CompressedKVCache, CompressionMethod, KVCacheConfig},
|
|
quantization::{
|
|
mixed_precision::{MixedPrecisionOptimizer, OptimizationObjective, PrecisionConfig},
|
|
product_quantization::{PQConfig, ProductQuantizer},
|
|
vector_quantization::{CodebookInitialization, DistanceMetric, VQConfig, VectorQuantizer},
|
|
},
|
|
};
|
|
use rtx_tensor::{DType, Device, Tensor};
|
|
use std::collections::HashMap;
|
|
|
|
fn bench_product_quantization(c: &mut Criterion) {
|
|
let device = Device::try_default().unwrap();
|
|
|
|
let mut group = c.benchmark_group("product_quantization");
|
|
|
|
// Test different data sizes
|
|
let sizes = vec![
|
|
("small", 1000, 128),
|
|
("medium", 10000, 256),
|
|
("large", 100000, 512),
|
|
];
|
|
|
|
for (name, num_vectors, dim) in sizes {
|
|
let data = Tensor::randn(&[num_vectors, dim], DType::F32, &device).unwrap();
|
|
|
|
// Benchmark codebook learning
|
|
group.bench_with_input(BenchmarkId::new("fit", name), &data, |b, data| {
|
|
b.iter(|| {
|
|
let config = PQConfig {
|
|
num_subquantizers: 8,
|
|
codebook_size: 256,
|
|
max_iterations: 20, // Reduced for benchmarking
|
|
tolerance: 1e-4,
|
|
};
|
|
let mut pq = ProductQuantizer::new(config);
|
|
black_box(pq.fit(data).unwrap());
|
|
});
|
|
});
|
|
|
|
// Pre-train a quantizer for encoding/decoding benchmarks
|
|
let config = PQConfig {
|
|
num_subquantizers: 8,
|
|
codebook_size: 256,
|
|
max_iterations: 10,
|
|
tolerance: 1e-3,
|
|
};
|
|
let mut pq = ProductQuantizer::new(config);
|
|
pq.fit(&data).unwrap();
|
|
|
|
// Benchmark encoding
|
|
group.bench_with_input(
|
|
BenchmarkId::new("encode", name),
|
|
&(&data, &pq),
|
|
|b, (data, pq)| {
|
|
b.iter(|| {
|
|
black_box(pq.encode(data).unwrap());
|
|
});
|
|
},
|
|
);
|
|
|
|
// Benchmark decoding
|
|
let codes = pq.encode(&data).unwrap();
|
|
group.bench_with_input(
|
|
BenchmarkId::new("decode", name),
|
|
&(&codes, &pq),
|
|
|b, (codes, pq)| {
|
|
b.iter(|| {
|
|
black_box(pq.decode(codes).unwrap());
|
|
});
|
|
},
|
|
);
|
|
}
|
|
|
|
group.finish();
|
|
}
|
|
|
|
fn bench_vector_quantization(c: &mut Criterion) {
|
|
let device = Device::try_default().unwrap();
|
|
|
|
let mut group = c.benchmark_group("vector_quantization");
|
|
|
|
let data = Tensor::randn(&[50000, 256], DType::F32, &device).unwrap();
|
|
|
|
// Test different codebook sizes
|
|
let codebook_sizes = vec![128, 256, 512, 1024];
|
|
|
|
for codebook_size in codebook_sizes {
|
|
let config = VQConfig {
|
|
codebook_size,
|
|
vector_dim: 256,
|
|
max_iterations: 20,
|
|
tolerance: 1e-4,
|
|
initialization: CodebookInitialization::KMeansPlusPlus,
|
|
distance_metric: DistanceMetric::Euclidean,
|
|
};
|
|
|
|
group.bench_with_input(
|
|
BenchmarkId::new("kmeans_fit", codebook_size),
|
|
&config,
|
|
|b, config| {
|
|
b.iter(|| {
|
|
let mut vq = VectorQuantizer::new(config.clone());
|
|
black_box(vq.fit(&data).unwrap());
|
|
});
|
|
},
|
|
);
|
|
|
|
// Pre-train for encoding benchmarks
|
|
let mut vq = VectorQuantizer::new(config);
|
|
vq.fit(&data).unwrap();
|
|
|
|
group.bench_with_input(
|
|
BenchmarkId::new("encode", codebook_size),
|
|
&(&data, &vq),
|
|
|b, (data, vq)| {
|
|
b.iter(|| {
|
|
black_box(vq.encode(data).unwrap());
|
|
});
|
|
},
|
|
);
|
|
}
|
|
|
|
group.finish();
|
|
}
|
|
|
|
fn bench_kv_cache_compression(c: &mut Criterion) {
|
|
let device = Device::try_default().unwrap();
|
|
|
|
let mut group = c.benchmark_group("kv_cache");
|
|
|
|
// Test different sequence lengths
|
|
let seq_lengths = vec![256, 512, 1024, 2048];
|
|
|
|
for seq_len in seq_lengths {
|
|
let keys = Tensor::randn(&[4, 12, seq_len, 64], DType::F32, &device).unwrap();
|
|
let values = Tensor::randn(&[4, 12, seq_len, 64], DType::F32, &device).unwrap();
|
|
|
|
// Product Quantization method
|
|
let pq_config = KVCacheConfig {
|
|
compression_method: CompressionMethod::ProductQuantization {
|
|
num_subquantizers: 8,
|
|
codebook_size: 256,
|
|
},
|
|
compression_ratio_target: 4.0,
|
|
quality_threshold: 0.90,
|
|
};
|
|
|
|
group.bench_with_input(
|
|
BenchmarkId::new("pq_insert", seq_len),
|
|
&(&keys, &values, &pq_config),
|
|
|b, (keys, values, config)| {
|
|
b.iter_with_setup(
|
|
|| CompressedKVCache::new(config.clone()),
|
|
|mut cache| {
|
|
black_box(cache.insert(0, keys, values).unwrap());
|
|
},
|
|
);
|
|
},
|
|
);
|
|
|
|
// Vector Quantization method
|
|
let vq_config = KVCacheConfig {
|
|
compression_method: CompressionMethod::VectorQuantization {
|
|
codebook_size: 1024,
|
|
update_frequency: 100,
|
|
},
|
|
compression_ratio_target: 3.0,
|
|
quality_threshold: 0.88,
|
|
};
|
|
|
|
group.bench_with_input(
|
|
BenchmarkId::new("vq_insert", seq_len),
|
|
&(&keys, &values, &vq_config),
|
|
|b, (keys, values, config)| {
|
|
b.iter_with_setup(
|
|
|| CompressedKVCache::new(config.clone()),
|
|
|mut cache| {
|
|
black_box(cache.insert(0, keys, values).unwrap());
|
|
},
|
|
);
|
|
},
|
|
);
|
|
|
|
// Benchmark retrieval
|
|
let mut cache = CompressedKVCache::new(pq_config);
|
|
cache.insert(0, &keys, &values).unwrap();
|
|
|
|
group.bench_with_input(
|
|
BenchmarkId::new("retrieve", seq_len),
|
|
&(&cache, seq_len),
|
|
|b, (cache, seq_len)| {
|
|
b.iter(|| {
|
|
black_box(cache.get(0, 0, *seq_len).unwrap());
|
|
});
|
|
},
|
|
);
|
|
}
|
|
|
|
group.finish();
|
|
}
|
|
|
|
fn bench_checkpoint_compression(c: &mut Criterion) {
|
|
let device = Device::try_default().unwrap();
|
|
|
|
let mut group = c.benchmark_group("checkpoint_compression");
|
|
|
|
// Create model of different sizes
|
|
let model_sizes = vec![
|
|
("small", vec![(512, 1024), (256, 512)]),
|
|
("medium", vec![(2048, 4096), (1024, 2048), (512, 1024)]),
|
|
(
|
|
"large",
|
|
vec![(4096, 8_192), (2048, 4096), (1024, 2048), (512, 1024)],
|
|
),
|
|
];
|
|
|
|
for (size_name, layer_dims) in model_sizes {
|
|
let mut state_dict = HashMap::new();
|
|
for (i, (in_dim, out_dim)) in layer_dims.iter().enumerate() {
|
|
state_dict.insert(
|
|
format!("layer_{}.weight", i),
|
|
Tensor::randn(&[*out_dim, *in_dim], DType::F32, &device).unwrap(),
|
|
);
|
|
}
|
|
|
|
// Test different compression formats
|
|
let formats = vec![
|
|
("lz4", CheckpointFormat::Lz4, 1),
|
|
("zstd_fast", CheckpointFormat::Zstd, 1),
|
|
("zstd_balanced", CheckpointFormat::Zstd, 6),
|
|
("zstd_best", CheckpointFormat::Zstd, 22),
|
|
];
|
|
|
|
for (format_name, format, compression_level) in formats {
|
|
let config = CompressionConfig {
|
|
format,
|
|
compression_level,
|
|
quantization_bits: 16, // fp16
|
|
exclude_patterns: vec![],
|
|
};
|
|
|
|
let compressor = CheckpointCompressor::new(config);
|
|
|
|
// Benchmark save
|
|
group.bench_with_input(
|
|
BenchmarkId::new(format!("save_{}_{}", size_name, format_name), ""),
|
|
&(&state_dict, &compressor),
|
|
|b, (state_dict, compressor)| {
|
|
b.iter(|| {
|
|
black_box(compressor.save(state_dict).unwrap());
|
|
});
|
|
},
|
|
);
|
|
|
|
// Pre-compress for load benchmark
|
|
let compressed_data = compressor.save(&state_dict).unwrap();
|
|
|
|
// Benchmark load
|
|
group.bench_with_input(
|
|
BenchmarkId::new(format!("load_{}_{}", size_name, format_name), ""),
|
|
&(&compressed_data, &compressor),
|
|
|b, (compressed_data, compressor)| {
|
|
b.iter(|| {
|
|
black_box(compressor.load(compressed_data).unwrap());
|
|
});
|
|
},
|
|
);
|
|
}
|
|
}
|
|
|
|
group.finish();
|
|
}
|
|
|
|
fn bench_mixed_precision_optimization(c: &mut Criterion) {
|
|
let device = Device::try_default().unwrap();
|
|
|
|
let mut group = c.benchmark_group("mixed_precision");
|
|
|
|
// Create different sized model architectures
|
|
let architectures = vec![
|
|
("transformer_small", 6, 512, 2048), // 6 layers, 512 hidden, 2048 ffn
|
|
("transformer_base", 12, 768, 3072), // 12 layers, 768 hidden, 3072 ffn
|
|
("transformer_large", 24, 1024, 4096), // 24 layers, 1024 hidden, 4096 ffn
|
|
];
|
|
|
|
for (arch_name, num_layers, hidden_dim, ffn_dim) in architectures {
|
|
let mut layers = HashMap::new();
|
|
|
|
for i in 0..num_layers {
|
|
layers.insert(
|
|
format!("layers.{}.attention.query", i),
|
|
Tensor::randn(&[hidden_dim, hidden_dim], DType::F32, &device).unwrap(),
|
|
);
|
|
layers.insert(
|
|
format!("layers.{}.attention.key", i),
|
|
Tensor::randn(&[hidden_dim, hidden_dim], DType::F32, &device).unwrap(),
|
|
);
|
|
layers.insert(
|
|
format!("layers.{}.attention.value", i),
|
|
Tensor::randn(&[hidden_dim, hidden_dim], DType::F32, &device).unwrap(),
|
|
);
|
|
layers.insert(
|
|
format!("layers.{}.ffn.up_proj", i),
|
|
Tensor::randn(&[hidden_dim, ffn_dim], DType::F32, &device).unwrap(),
|
|
);
|
|
layers.insert(
|
|
format!("layers.{}.ffn.down_proj", i),
|
|
Tensor::randn(&[ffn_dim, hidden_dim], DType::F32, &device).unwrap(),
|
|
);
|
|
}
|
|
|
|
let calibration_data = Tensor::randn(&[100, hidden_dim], DType::F32, &device).unwrap();
|
|
|
|
// Benchmark sensitivity analysis
|
|
let config = PrecisionConfig {
|
|
precision_bits: vec![4, 8, 12, 16],
|
|
sensitivity_threshold: 0.02,
|
|
performance_weight: 0.6,
|
|
quality_weight: 0.4,
|
|
};
|
|
|
|
group.bench_with_input(
|
|
BenchmarkId::new("sensitivity_analysis", arch_name),
|
|
&(&layers, &calibration_data, &config),
|
|
|b, (layers, calibration_data, config)| {
|
|
b.iter(|| {
|
|
let mut optimizer = MixedPrecisionOptimizer::new(config.clone());
|
|
black_box(
|
|
optimizer
|
|
.analyze_sensitivity(layers, calibration_data)
|
|
.unwrap(),
|
|
);
|
|
});
|
|
},
|
|
);
|
|
|
|
// Benchmark full optimization
|
|
let objective = OptimizationObjective {
|
|
target_compression_ratio: 4.0,
|
|
max_quality_loss: 0.05,
|
|
memory_constraint_mb: Some(1000),
|
|
};
|
|
|
|
group.bench_with_input(
|
|
BenchmarkId::new("full_optimization", arch_name),
|
|
&(&layers, &calibration_data, &config, &objective),
|
|
|b, (layers, calibration_data, config, objective)| {
|
|
b.iter(|| {
|
|
let mut optimizer = MixedPrecisionOptimizer::new(config.clone());
|
|
black_box(
|
|
optimizer
|
|
.optimize(layers, calibration_data, objective.clone())
|
|
.unwrap(),
|
|
);
|
|
});
|
|
},
|
|
);
|
|
}
|
|
|
|
group.finish();
|
|
}
|
|
|
|
fn bench_compression_throughput(c: &mut Criterion) {
|
|
let device = Device::try_default().unwrap();
|
|
|
|
let mut group = c.benchmark_group("throughput");
|
|
group.throughput(criterion::Throughput::Bytes(1024 * 1024 * 4)); // 1M f32 values
|
|
|
|
let data = Tensor::randn(&[1024, 1024], DType::F32, &device).unwrap();
|
|
|
|
// Benchmark different compression algorithms
|
|
let algorithms = vec![
|
|
("pq_8x256", |data: &Tensor| {
|
|
let config = PQConfig {
|
|
num_subquantizers: 8,
|
|
codebook_size: 256,
|
|
max_iterations: 10,
|
|
tolerance: 1e-3,
|
|
};
|
|
let mut pq = ProductQuantizer::new(config);
|
|
pq.fit(data).unwrap();
|
|
pq.encode(data).unwrap();
|
|
}),
|
|
("vq_1024", |data: &Tensor| {
|
|
let config = VQConfig {
|
|
codebook_size: 1024,
|
|
vector_dim: 1024,
|
|
max_iterations: 10,
|
|
tolerance: 1e-3,
|
|
initialization: CodebookInitialization::Random,
|
|
distance_metric: DistanceMetric::Euclidean,
|
|
};
|
|
let mut vq = VectorQuantizer::new(config);
|
|
vq.fit(data).unwrap();
|
|
vq.encode(data).unwrap();
|
|
}),
|
|
];
|
|
|
|
for (name, algorithm) in algorithms {
|
|
group.bench_function(name, |b| {
|
|
b.iter(|| {
|
|
black_box(algorithm(&data));
|
|
});
|
|
});
|
|
}
|
|
|
|
group.finish();
|
|
}
|
|
|
|
fn bench_memory_efficiency(c: &mut Criterion) {
|
|
let device = Device::try_default().unwrap();
|
|
|
|
let mut group = c.benchmark_group("memory_efficiency");
|
|
|
|
// Test compression with different memory constraints
|
|
let memory_limits = vec![100, 200, 500, 1000]; // MB
|
|
|
|
for memory_limit in memory_limits {
|
|
let data_size_mb = memory_limit * 2; // 2x memory pressure
|
|
let num_elements = (data_size_mb * 1024 * 1024) / 4; // f32 = 4 bytes
|
|
let data = Tensor::randn(&[num_elements], DType::F32, &device).unwrap();
|
|
|
|
group.bench_with_input(
|
|
BenchmarkId::new("constrained_compression", memory_limit),
|
|
&(&data, memory_limit),
|
|
|b, (data, memory_limit)| {
|
|
b.iter(|| {
|
|
// Simulate memory-constrained compression
|
|
let config = PQConfig {
|
|
num_subquantizers: 8,
|
|
codebook_size: 128, // Smaller codebook for memory constraint
|
|
max_iterations: 5, // Fewer iterations
|
|
tolerance: 1e-2, // Less strict tolerance
|
|
};
|
|
|
|
let mut pq = ProductQuantizer::new(config);
|
|
|
|
// In real implementation, this would respect memory limits
|
|
let reshaped = data.reshape(&[num_elements / 128, 128]).unwrap();
|
|
pq.fit(&reshaped).unwrap();
|
|
black_box(pq.encode(&reshaped).unwrap());
|
|
});
|
|
},
|
|
);
|
|
}
|
|
|
|
group.finish();
|
|
}
|
|
|
|
criterion_group!(
|
|
benches,
|
|
bench_product_quantization,
|
|
bench_vector_quantization,
|
|
bench_kv_cache_compression,
|
|
bench_checkpoint_compression,
|
|
bench_mixed_precision_optimization,
|
|
bench_compression_throughput,
|
|
bench_memory_efficiency
|
|
);
|
|
|
|
criterion_main!(benches);
|