484 lines
14 KiB
Rust
484 lines
14 KiB
Rust
//! Sample data and configurations for FoundationForge demo.
|
|
|
|
use forge_shared::{
|
|
BenchmarkConfig, CompressionStep, DistillLoss, DistillationConfig, ExportConfig, ExportFormat,
|
|
ModelArchitecture, ModelInfo, PipelineConfig, PruneMethod, PruneSchedule, PruningConfig,
|
|
QuantMethod, QuantPrecision, QuantizationConfig, StudentConfig, TargetDevice,
|
|
};
|
|
|
|
// ============================================================================
|
|
// Model Configurations
|
|
// ============================================================================
|
|
|
|
/// Create a sample LLaMA 7B model info.
|
|
#[must_use]
|
|
pub fn llama_7b() -> ModelInfo {
|
|
ModelInfo {
|
|
name: "llama-7b".to_string(),
|
|
architecture: ModelArchitecture::Transformer,
|
|
num_parameters: 7_000_000_000,
|
|
num_layers: 32,
|
|
hidden_dim: 4096,
|
|
vocab_size: Some(32000),
|
|
image_size: None,
|
|
precision_bits: 16,
|
|
size_bytes: 14_000_000_000,
|
|
}
|
|
}
|
|
|
|
/// Create a sample LLaMA 13B model info.
|
|
#[must_use]
|
|
pub fn llama_13b() -> ModelInfo {
|
|
ModelInfo {
|
|
name: "llama-13b".to_string(),
|
|
architecture: ModelArchitecture::Transformer,
|
|
num_parameters: 13_000_000_000,
|
|
num_layers: 40,
|
|
hidden_dim: 5120,
|
|
vocab_size: Some(32000),
|
|
image_size: None,
|
|
precision_bits: 16,
|
|
size_bytes: 26_000_000_000,
|
|
}
|
|
}
|
|
|
|
/// Create a sample Mistral 7B model info.
|
|
#[must_use]
|
|
pub fn mistral_7b() -> ModelInfo {
|
|
ModelInfo {
|
|
name: "mistral-7b".to_string(),
|
|
architecture: ModelArchitecture::Transformer,
|
|
num_parameters: 7_300_000_000,
|
|
num_layers: 32,
|
|
hidden_dim: 4096,
|
|
vocab_size: Some(32000),
|
|
image_size: None,
|
|
precision_bits: 16,
|
|
size_bytes: 14_600_000_000,
|
|
}
|
|
}
|
|
|
|
/// Create a sample ViT-Large model info.
|
|
#[must_use]
|
|
pub fn vit_large() -> ModelInfo {
|
|
ModelInfo {
|
|
name: "vit-large".to_string(),
|
|
architecture: ModelArchitecture::ViT,
|
|
num_parameters: 307_000_000,
|
|
num_layers: 24,
|
|
hidden_dim: 1024,
|
|
vocab_size: None,
|
|
image_size: Some((224, 224)),
|
|
precision_bits: 32,
|
|
size_bytes: 1_228_000_000,
|
|
}
|
|
}
|
|
|
|
/// Create a sample BERT-Base model info.
|
|
#[must_use]
|
|
pub fn bert_base() -> ModelInfo {
|
|
ModelInfo {
|
|
name: "bert-base".to_string(),
|
|
architecture: ModelArchitecture::Transformer,
|
|
num_parameters: 110_000_000,
|
|
num_layers: 12,
|
|
hidden_dim: 768,
|
|
vocab_size: Some(30522),
|
|
image_size: None,
|
|
precision_bits: 32,
|
|
size_bytes: 440_000_000,
|
|
}
|
|
}
|
|
|
|
// ============================================================================
|
|
// Quantization Configurations
|
|
// ============================================================================
|
|
|
|
/// Create INT4 GPTQ quantization config.
|
|
#[must_use]
|
|
pub fn gptq_int4_config() -> QuantizationConfig {
|
|
QuantizationConfig {
|
|
method: QuantMethod::GPTQ,
|
|
precision: QuantPrecision::Int4,
|
|
group_size: Some(128),
|
|
quantize_activations: false,
|
|
quantize_embeddings: false,
|
|
calibration_size: 512,
|
|
symmetric: false,
|
|
per_channel: true,
|
|
skip_layers: vec![],
|
|
}
|
|
}
|
|
|
|
/// Create INT8 AWQ quantization config.
|
|
#[must_use]
|
|
pub fn awq_int8_config() -> QuantizationConfig {
|
|
QuantizationConfig {
|
|
method: QuantMethod::AWQ,
|
|
precision: QuantPrecision::Int8,
|
|
group_size: None,
|
|
quantize_activations: true,
|
|
quantize_embeddings: false,
|
|
calibration_size: 256,
|
|
symmetric: true,
|
|
per_channel: true,
|
|
skip_layers: vec![],
|
|
}
|
|
}
|
|
|
|
/// Create GGML quantization config.
|
|
#[must_use]
|
|
pub fn ggml_q4_k_m_config() -> QuantizationConfig {
|
|
QuantizationConfig {
|
|
method: QuantMethod::GGML,
|
|
precision: QuantPrecision::Int4,
|
|
group_size: Some(32),
|
|
quantize_activations: false,
|
|
quantize_embeddings: true,
|
|
calibration_size: 0,
|
|
symmetric: false,
|
|
per_channel: true,
|
|
skip_layers: vec![],
|
|
}
|
|
}
|
|
|
|
/// Create dynamic quantization config.
|
|
#[must_use]
|
|
pub fn dynamic_int8_config() -> QuantizationConfig {
|
|
QuantizationConfig {
|
|
method: QuantMethod::Dynamic,
|
|
precision: QuantPrecision::Int8,
|
|
group_size: None,
|
|
quantize_activations: true,
|
|
quantize_embeddings: false,
|
|
calibration_size: 0,
|
|
symmetric: true,
|
|
per_channel: false,
|
|
skip_layers: vec!["lm_head".to_string()],
|
|
}
|
|
}
|
|
|
|
// ============================================================================
|
|
// Pruning Configurations
|
|
// ============================================================================
|
|
|
|
/// Create 50% magnitude pruning config.
|
|
#[must_use]
|
|
pub fn magnitude_50_config() -> PruningConfig {
|
|
PruningConfig {
|
|
method: PruneMethod::Magnitude,
|
|
target_sparsity: 0.5,
|
|
schedule: PruneSchedule::Gradual,
|
|
num_steps: 10,
|
|
initial_sparsity: 0.0,
|
|
final_sparsity: 0.5,
|
|
retrain_epochs: 3,
|
|
n: None,
|
|
m: None,
|
|
exclude_layers: vec!["embedding".to_string(), "lm_head".to_string()],
|
|
}
|
|
}
|
|
|
|
/// Create 2:4 structured sparsity config.
|
|
#[must_use]
|
|
pub fn nm_2_4_config() -> PruningConfig {
|
|
PruningConfig {
|
|
method: PruneMethod::NM,
|
|
target_sparsity: 0.5,
|
|
schedule: PruneSchedule::OneShot,
|
|
num_steps: 1,
|
|
initial_sparsity: 0.5,
|
|
final_sparsity: 0.5,
|
|
retrain_epochs: 5,
|
|
n: Some(2),
|
|
m: Some(4),
|
|
exclude_layers: vec![],
|
|
}
|
|
}
|
|
|
|
/// Create aggressive 90% pruning config.
|
|
#[must_use]
|
|
pub fn aggressive_prune_config() -> PruningConfig {
|
|
PruningConfig {
|
|
method: PruneMethod::Movement,
|
|
target_sparsity: 0.9,
|
|
schedule: PruneSchedule::Cubic,
|
|
num_steps: 20,
|
|
initial_sparsity: 0.3,
|
|
final_sparsity: 0.9,
|
|
retrain_epochs: 10,
|
|
n: None,
|
|
m: None,
|
|
exclude_layers: vec!["embedding".to_string()],
|
|
}
|
|
}
|
|
|
|
// ============================================================================
|
|
// Distillation Configurations
|
|
// ============================================================================
|
|
|
|
/// Create standard KL distillation config.
|
|
#[must_use]
|
|
pub fn kl_distill_config() -> DistillationConfig {
|
|
DistillationConfig {
|
|
loss_type: DistillLoss::KL,
|
|
temperature: 4.0,
|
|
alpha: 0.5,
|
|
intermediate_matching: true,
|
|
layer_mapping: vec![(0, 0), (10, 1), (21, 2), (31, 3)],
|
|
epochs: 10,
|
|
batch_size: 32,
|
|
learning_rate: 5e-5,
|
|
progressive: false,
|
|
}
|
|
}
|
|
|
|
/// Create TinyBERT-style distillation config.
|
|
#[must_use]
|
|
pub fn tinybert_distill_config() -> DistillationConfig {
|
|
DistillationConfig {
|
|
loss_type: DistillLoss::Combined,
|
|
temperature: 1.0,
|
|
alpha: 0.7,
|
|
intermediate_matching: true,
|
|
layer_mapping: vec![(0, 0), (3, 1), (6, 2), (9, 3), (11, 4), (11, 5)],
|
|
epochs: 20,
|
|
batch_size: 64,
|
|
learning_rate: 3e-5,
|
|
progressive: true,
|
|
}
|
|
}
|
|
|
|
/// Create attention transfer distillation config.
|
|
#[must_use]
|
|
pub fn attention_distill_config() -> DistillationConfig {
|
|
DistillationConfig {
|
|
loss_type: DistillLoss::AttentionTransfer,
|
|
temperature: 2.0,
|
|
alpha: 0.3,
|
|
intermediate_matching: true,
|
|
layer_mapping: vec![],
|
|
epochs: 15,
|
|
batch_size: 16,
|
|
learning_rate: 1e-4,
|
|
progressive: false,
|
|
}
|
|
}
|
|
|
|
// ============================================================================
|
|
// Student Model Configurations
|
|
// ============================================================================
|
|
|
|
/// Create a 6-layer student config.
|
|
#[must_use]
|
|
pub fn student_6_layer() -> StudentConfig {
|
|
StudentConfig {
|
|
num_layers: 6,
|
|
hidden_dim: 768,
|
|
num_heads: 12,
|
|
intermediate_dim: 3072,
|
|
init_from_teacher: true,
|
|
copy_layers: vec![0, 2, 5, 8, 10, 11],
|
|
}
|
|
}
|
|
|
|
/// Create a 4-layer student config.
|
|
#[must_use]
|
|
pub fn student_4_layer() -> StudentConfig {
|
|
StudentConfig {
|
|
num_layers: 4,
|
|
hidden_dim: 512,
|
|
num_heads: 8,
|
|
intermediate_dim: 2048,
|
|
init_from_teacher: true,
|
|
copy_layers: vec![0, 4, 8, 11],
|
|
}
|
|
}
|
|
|
|
// ============================================================================
|
|
// Pipeline Configurations
|
|
// ============================================================================
|
|
|
|
/// Create a mobile deployment pipeline.
|
|
#[must_use]
|
|
pub fn mobile_pipeline() -> PipelineConfig {
|
|
PipelineConfig {
|
|
name: "mobile-optimized".to_string(),
|
|
steps: vec![
|
|
CompressionStep::Quantize(gptq_int4_config()),
|
|
CompressionStep::PruneHeads { target_heads: 8 },
|
|
],
|
|
eval_dataset: Some("wikitext".to_string()),
|
|
target_size: Some(2_000_000_000),
|
|
max_accuracy_loss: 0.05,
|
|
export: ExportConfig {
|
|
format: ExportFormat::CoreML,
|
|
optimize: true,
|
|
include_tokenizer: true,
|
|
target_device: TargetDevice::Mobile,
|
|
output_path: "output/mobile_model".to_string(),
|
|
},
|
|
}
|
|
}
|
|
|
|
/// Create a server deployment pipeline.
|
|
#[must_use]
|
|
pub fn server_pipeline() -> PipelineConfig {
|
|
PipelineConfig {
|
|
name: "server-optimized".to_string(),
|
|
steps: vec![CompressionStep::Quantize(awq_int8_config())],
|
|
eval_dataset: Some("wikitext".to_string()),
|
|
target_size: None,
|
|
max_accuracy_loss: 0.01,
|
|
export: ExportConfig {
|
|
format: ExportFormat::SafeTensors,
|
|
optimize: true,
|
|
include_tokenizer: true,
|
|
target_device: TargetDevice::CUDA,
|
|
output_path: "output/server_model".to_string(),
|
|
},
|
|
}
|
|
}
|
|
|
|
/// Create an edge deployment pipeline.
|
|
#[must_use]
|
|
pub fn edge_pipeline() -> PipelineConfig {
|
|
PipelineConfig {
|
|
name: "edge-optimized".to_string(),
|
|
steps: vec![
|
|
CompressionStep::Prune(nm_2_4_config()),
|
|
CompressionStep::Quantize(gptq_int4_config()),
|
|
CompressionStep::PruneVocab { keep_tokens: 16000 },
|
|
],
|
|
eval_dataset: Some("wikitext".to_string()),
|
|
target_size: Some(1_000_000_000),
|
|
max_accuracy_loss: 0.08,
|
|
export: ExportConfig {
|
|
format: ExportFormat::TFLite,
|
|
optimize: true,
|
|
include_tokenizer: true,
|
|
target_device: TargetDevice::Edge,
|
|
output_path: "output/edge_model".to_string(),
|
|
},
|
|
}
|
|
}
|
|
|
|
/// Create a maximum compression pipeline.
|
|
#[must_use]
|
|
pub fn max_compression_pipeline() -> PipelineConfig {
|
|
PipelineConfig {
|
|
name: "max-compression".to_string(),
|
|
steps: vec![
|
|
CompressionStep::Distill(kl_distill_config()),
|
|
CompressionStep::Prune(aggressive_prune_config()),
|
|
CompressionStep::Quantize(ggml_q4_k_m_config()),
|
|
],
|
|
eval_dataset: Some("wikitext".to_string()),
|
|
target_size: Some(500_000_000),
|
|
max_accuracy_loss: 0.15,
|
|
export: ExportConfig {
|
|
format: ExportFormat::GGUF,
|
|
optimize: true,
|
|
include_tokenizer: true,
|
|
target_device: TargetDevice::CPU,
|
|
output_path: "output/tiny_model".to_string(),
|
|
},
|
|
}
|
|
}
|
|
|
|
// ============================================================================
|
|
// Benchmark Configurations
|
|
// ============================================================================
|
|
|
|
/// Create a quick benchmark config.
|
|
#[must_use]
|
|
pub fn quick_benchmark() -> BenchmarkConfig {
|
|
BenchmarkConfig {
|
|
dataset: "wikitext".to_string(),
|
|
num_samples: 100,
|
|
batch_size: 1,
|
|
measure_latency: true,
|
|
measure_memory: true,
|
|
warmup_runs: 2,
|
|
benchmark_runs: 5,
|
|
}
|
|
}
|
|
|
|
/// Create a full benchmark config.
|
|
#[must_use]
|
|
pub fn full_benchmark() -> BenchmarkConfig {
|
|
BenchmarkConfig {
|
|
dataset: "wikitext".to_string(),
|
|
num_samples: 1000,
|
|
batch_size: 1,
|
|
measure_latency: true,
|
|
measure_memory: true,
|
|
warmup_runs: 5,
|
|
benchmark_runs: 20,
|
|
}
|
|
}
|
|
|
|
// ============================================================================
|
|
// Tests
|
|
// ============================================================================
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
|
|
#[test]
|
|
fn test_model_infos() {
|
|
let llama7 = llama_7b();
|
|
let llama13 = llama_13b();
|
|
|
|
assert!(llama13.num_parameters > llama7.num_parameters);
|
|
assert!(llama13.num_layers > llama7.num_layers);
|
|
}
|
|
|
|
#[test]
|
|
fn test_quant_configs() {
|
|
let gptq = gptq_int4_config();
|
|
let awq = awq_int8_config();
|
|
|
|
assert_eq!(gptq.precision, QuantPrecision::Int4);
|
|
assert_eq!(awq.precision, QuantPrecision::Int8);
|
|
}
|
|
|
|
#[test]
|
|
fn test_prune_configs() {
|
|
let mag = magnitude_50_config();
|
|
let nm = nm_2_4_config();
|
|
|
|
assert_eq!(mag.target_sparsity, 0.5);
|
|
assert_eq!(nm.n, Some(2));
|
|
assert_eq!(nm.m, Some(4));
|
|
}
|
|
|
|
#[test]
|
|
fn test_distill_configs() {
|
|
let kl = kl_distill_config();
|
|
let tinybert = tinybert_distill_config();
|
|
|
|
assert_eq!(kl.loss_type, DistillLoss::KL);
|
|
assert_eq!(tinybert.loss_type, DistillLoss::Combined);
|
|
}
|
|
|
|
#[test]
|
|
fn test_pipelines() {
|
|
let mobile = mobile_pipeline();
|
|
let edge = edge_pipeline();
|
|
|
|
assert!(mobile.target_size.unwrap() > edge.target_size.unwrap());
|
|
assert_eq!(mobile.export.target_device, TargetDevice::Mobile);
|
|
}
|
|
|
|
#[test]
|
|
fn test_benchmarks() {
|
|
let quick = quick_benchmark();
|
|
let full = full_benchmark();
|
|
|
|
assert!(quick.num_samples < full.num_samples);
|
|
assert!(quick.benchmark_runs < full.benchmark_runs);
|
|
}
|
|
}
|