Files
foxhunt/crates/ml/examples/optimize_batch_sizes.rs
jgrusewski 9c3d741a08 refactor: restructure repo — crates/, bin/, testing/ layout
Move 17 library crates into crates/, CLI binary into bin/fxt,
consolidate 10 test crates into testing/, split config crate
from deployment config files.

Root directory reduced from 38+ to ~17 directories.
All Cargo.toml paths and build.rs proto refs updated.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
2026-02-25 11:56:00 +01:00

662 lines
20 KiB
Rust

//! GPU Batch Size Optimization for RTX 3050 Ti (4GB VRAM)
//!
//! Tests multiple batch sizes for TFT, MAMBA-2, and Liquid models to find
//! optimal configurations that maximize throughput while staying under 4GB VRAM.
//!
//! ## Testing Strategy
//!
//! 1. **TFT**: Test batch sizes [16, 32, 64, 128]
//! 2. **MAMBA-2**: Test batch sizes [8, 16, 32]
//! 3. **Liquid**: Test batch sizes [16, 32, 64]
//! 4. Monitor VRAM usage with nvidia-smi integration
//! 5. Find optimal batch size (max throughput, <4GB VRAM)
//!
//! ## Usage
//!
//! ```bash
//! cargo run -p ml --example optimize_batch_sizes --release
//! ```
//!
//! ## Output
//!
//! Generates `BATCH_SIZE_OPTIMIZATION_REPORT.md` with:
//! - VRAM usage per model/batch size
//! - Throughput measurements
//! - Recommended optimal batch sizes
//! - Updated configuration snippets
use candle_core::{DType, Device, Tensor};
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::fs::File;
use std::io::Write;
use std::process::Command;
use std::time::Instant;
use sysinfo::{ProcessExt, System, SystemExt};
// Import model types
use ml::liquid::network::{LiquidNetwork, LiquidNetworkConfig};
use ml::mamba::{Mamba2Config, Mamba2SSM};
use ml::tft::{TFTConfig, TemporalFusionTransformer};
const VRAM_LIMIT_GB: f32 = 4.0;
const WARMUP_ITERATIONS: usize = 5;
const BENCHMARK_ITERATIONS: usize = 20;
#[derive(Debug, Clone, Serialize, Deserialize)]
struct BatchSizeResult {
model_name: String,
batch_size: usize,
vram_used_mb: f32,
vram_percent: f32,
throughput_samples_per_sec: f32,
latency_ms: f32,
oom_occurred: bool,
recommended: bool,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
struct OptimizationReport {
timestamp: String,
gpu_name: String,
vram_total_gb: f32,
results: Vec<BatchSizeResult>,
recommendations: HashMap<String, usize>,
notes: Vec<String>,
}
/// Query NVIDIA GPU VRAM usage in MB
fn query_nvidia_vram() -> Result<f32, Box<dyn std::error::Error>> {
let output = Command::new("nvidia-smi")
.args(&["--query-gpu=memory.used", "--format=csv,noheader,nounits"])
.output()?;
let vram_str = String::from_utf8(output.stdout)?;
let vram_mb: f32 = vram_str.trim().parse()?;
Ok(vram_mb)
}
/// Query NVIDIA GPU name
fn query_gpu_name() -> Result<String, Box<dyn std::error::Error>> {
let output = Command::new("nvidia-smi")
.args(&["--query-gpu=name", "--format=csv,noheader"])
.output()?;
Ok(String::from_utf8(output.stdout)?.trim().to_string())
}
/// Query total VRAM in GB
fn query_total_vram_gb() -> Result<f32, Box<dyn std::error::Error>> {
let output = Command::new("nvidia-smi")
.args(&["--query-gpu=memory.total", "--format=csv,noheader,nounits"])
.output()?;
let vram_mb: f32 = String::from_utf8(output.stdout)?.trim().parse()?;
Ok(vram_mb / 1024.0) // Convert to GB
}
/// Benchmark TFT model with specific batch size
fn benchmark_tft_batch(
batch_size: usize,
device: &Device,
) -> Result<BatchSizeResult, Box<dyn std::error::Error>> {
println!(" Testing TFT with batch_size={}", batch_size);
// Create TFT configuration
let config = TFTConfig {
input_dim: 64,
hidden_dim: 256,
num_heads: 8,
num_layers: 4,
prediction_horizon: 10,
sequence_length: 50,
num_quantiles: 9,
num_static_features: 5,
num_known_features: 10,
num_unknown_features: 20,
learning_rate: 1e-3,
batch_size,
dropout_rate: 0.1,
l2_regularization: 1e-4,
use_flash_attention: true,
mixed_precision: false,
memory_efficient: true,
max_inference_latency_us: 50,
target_throughput_pps: 100_000,
};
let mut model = match TemporalFusionTransformer::new(config.clone()) {
Ok(m) => m,
Err(e) => {
return Ok(BatchSizeResult {
model_name: "TFT".to_string(),
batch_size,
vram_used_mb: 0.0,
vram_percent: 0.0,
throughput_samples_per_sec: 0.0,
latency_ms: 0.0,
oom_occurred: true,
recommended: false,
});
},
};
// Prepare batch inputs
let static_features = vec![0.0f32; config.num_static_features * batch_size];
let historical_features =
vec![0.0f32; config.sequence_length * config.num_unknown_features * batch_size];
let future_features =
vec![0.0f32; config.prediction_horizon * config.num_known_features * batch_size];
// Warmup
for _ in 0..WARMUP_ITERATIONS {
let _ = model.predict_fast(&static_features, &historical_features, &future_features);
}
// Clear GPU cache
std::thread::sleep(std::time::Duration::from_millis(500));
// Measure baseline VRAM
let baseline_vram = query_nvidia_vram()?;
// Benchmark iterations
let start = Instant::now();
for _ in 0..BENCHMARK_ITERATIONS {
match model.predict_fast(&static_features, &historical_features, &future_features) {
Ok(_) => {},
Err(_) => {
return Ok(BatchSizeResult {
model_name: "TFT".to_string(),
batch_size,
vram_used_mb: 0.0,
vram_percent: 0.0,
throughput_samples_per_sec: 0.0,
latency_ms: 0.0,
oom_occurred: true,
recommended: false,
});
},
}
}
let elapsed = start.elapsed();
// Measure peak VRAM
let peak_vram = query_nvidia_vram()?;
let vram_used_mb = peak_vram - baseline_vram;
let total_vram_gb = query_total_vram_gb()?;
let vram_percent = (peak_vram / (total_vram_gb * 1024.0)) * 100.0;
let latency_ms = elapsed.as_secs_f32() * 1000.0 / BENCHMARK_ITERATIONS as f32;
let throughput = (batch_size * BENCHMARK_ITERATIONS) as f32 / elapsed.as_secs_f32();
let recommended = vram_percent < 90.0 && !false;
Ok(BatchSizeResult {
model_name: "TFT".to_string(),
batch_size,
vram_used_mb,
vram_percent,
throughput_samples_per_sec: throughput,
latency_ms,
oom_occurred: false,
recommended,
})
}
/// Benchmark MAMBA-2 model with specific batch size
fn benchmark_mamba_batch(
batch_size: usize,
device: &Device,
) -> Result<BatchSizeResult, Box<dyn std::error::Error>> {
println!(" Testing MAMBA-2 with batch_size={}", batch_size);
let config = Mamba2Config {
d_model: 256,
d_state: 64,
d_head: 32,
num_heads: 8,
expand: 2,
num_layers: 4,
dropout: 0.1,
use_ssd: true,
use_selective_state: true,
hardware_aware: true,
target_latency_us: 5,
max_seq_len: 512,
learning_rate: 1e-3,
weight_decay: 1e-4,
grad_clip: 1.0,
warmup_steps: 1000,
batch_size,
seq_len: 256,
};
let mut model = match Mamba2SSM::new(config.clone(), device) {
Ok(m) => m,
Err(e) => {
return Ok(BatchSizeResult {
model_name: "MAMBA-2".to_string(),
batch_size,
vram_used_mb: 0.0,
vram_percent: 0.0,
throughput_samples_per_sec: 0.0,
latency_ms: 0.0,
oom_occurred: true,
recommended: false,
});
},
};
// Prepare batch input
let input_tensor = match Tensor::randn(0f32, 1f32, (batch_size, config.d_model), device) {
Ok(t) => t,
Err(_) => {
return Ok(BatchSizeResult {
model_name: "MAMBA-2".to_string(),
batch_size,
vram_used_mb: 0.0,
vram_percent: 0.0,
throughput_samples_per_sec: 0.0,
latency_ms: 0.0,
oom_occurred: true,
recommended: false,
});
},
};
// Warmup
for _ in 0..WARMUP_ITERATIONS {
let _ = model.forward(&input_tensor);
}
std::thread::sleep(std::time::Duration::from_millis(500));
let baseline_vram = query_nvidia_vram()?;
// Benchmark
let start = Instant::now();
for _ in 0..BENCHMARK_ITERATIONS {
match model.forward(&input_tensor) {
Ok(_) => {},
Err(_) => {
return Ok(BatchSizeResult {
model_name: "MAMBA-2".to_string(),
batch_size,
vram_used_mb: 0.0,
vram_percent: 0.0,
throughput_samples_per_sec: 0.0,
latency_ms: 0.0,
oom_occurred: true,
recommended: false,
});
},
}
}
let elapsed = start.elapsed();
let peak_vram = query_nvidia_vram()?;
let vram_used_mb = peak_vram - baseline_vram;
let total_vram_gb = query_total_vram_gb()?;
let vram_percent = (peak_vram / (total_vram_gb * 1024.0)) * 100.0;
let latency_ms = elapsed.as_secs_f32() * 1000.0 / BENCHMARK_ITERATIONS as f32;
let throughput = (batch_size * BENCHMARK_ITERATIONS) as f32 / elapsed.as_secs_f32();
let recommended = vram_percent < 90.0 && !false;
Ok(BatchSizeResult {
model_name: "MAMBA-2".to_string(),
batch_size,
vram_used_mb,
vram_percent,
throughput_samples_per_sec: throughput,
latency_ms,
oom_occurred: false,
recommended,
})
}
/// Benchmark Liquid model with specific batch size
fn benchmark_liquid_batch(
batch_size: usize,
) -> Result<BatchSizeResult, Box<dyn std::error::Error>> {
println!(" Testing Liquid with batch_size={}", batch_size);
let config = LiquidNetworkConfig {
input_dim: 256,
hidden_dim: 512,
output_dim: 3,
num_layers: 4,
cell_type: ml::liquid::cells::CellType::LTC,
activation: ml::liquid::activation::ActivationType::Tanh,
solver: ml::liquid::ode_solvers::SolverType::Euler,
time_step: 0.001,
inference_steps: 10,
};
let model = match LiquidNetwork::new(config.clone()) {
Ok(m) => m,
Err(e) => {
return Ok(BatchSizeResult {
model_name: "Liquid".to_string(),
batch_size,
vram_used_mb: 0.0,
vram_percent: 0.0,
throughput_samples_per_sec: 0.0,
latency_ms: 0.0,
oom_occurred: true,
recommended: false,
});
},
};
// Prepare batch input (Note: Liquid uses CPU, no GPU VRAM)
let input = vec![ml::liquid::FixedPoint::zero(); config.input_dim * batch_size];
// Warmup
for _ in 0..WARMUP_ITERATIONS {
let _ = model.forward(&input[..config.input_dim]);
}
std::thread::sleep(std::time::Duration::from_millis(500));
let baseline_vram = query_nvidia_vram().unwrap_or(0.0);
// Benchmark
let start = Instant::now();
for i in 0..BENCHMARK_ITERATIONS {
let offset = (i % batch_size) * config.input_dim;
match model.forward(&input[offset..offset + config.input_dim]) {
Ok(_) => {},
Err(_) => {
return Ok(BatchSizeResult {
model_name: "Liquid".to_string(),
batch_size,
vram_used_mb: 0.0,
vram_percent: 0.0,
throughput_samples_per_sec: 0.0,
latency_ms: 0.0,
oom_occurred: true,
recommended: false,
});
},
}
}
let elapsed = start.elapsed();
let peak_vram = query_nvidia_vram().unwrap_or(0.0);
let vram_used_mb = peak_vram - baseline_vram;
let total_vram_gb = query_total_vram_gb().unwrap_or(4.0);
let vram_percent = (peak_vram / (total_vram_gb * 1024.0)) * 100.0;
let latency_ms = elapsed.as_secs_f32() * 1000.0 / BENCHMARK_ITERATIONS as f32;
let throughput = (batch_size * BENCHMARK_ITERATIONS) as f32 / elapsed.as_secs_f32();
// Liquid is CPU-based, always safe
let recommended = true;
Ok(BatchSizeResult {
model_name: "Liquid".to_string(),
batch_size,
vram_used_mb,
vram_percent,
throughput_samples_per_sec: throughput,
latency_ms,
oom_occurred: false,
recommended,
})
}
/// Generate markdown report
fn generate_report(report: &OptimizationReport) -> String {
let mut md = String::new();
md.push_str("# GPU Batch Size Optimization Report\n\n");
md.push_str(&format!("**Generated**: {}\n", report.timestamp));
md.push_str(&format!("**GPU**: {}\n", report.gpu_name));
md.push_str(&format!("**VRAM**: {:.1} GB\n\n", report.vram_total_gb));
md.push_str("## Optimization Summary\n\n");
md.push_str("| Model | Recommended Batch Size | VRAM Usage | Throughput |\n");
md.push_str("|-------|------------------------|------------|------------|\n");
for (model, batch) in &report.recommendations {
let result = report
.results
.iter()
.find(|r| r.model_name == *model && r.batch_size == *batch)
.unwrap();
md.push_str(&format!(
"| {} | {} | {:.0} MB ({:.1}%) | {:.1} samples/sec |\n",
model,
batch,
result.vram_used_mb,
result.vram_percent,
result.throughput_samples_per_sec
));
}
md.push_str("\n");
md.push_str("## Detailed Results\n\n");
// Group by model
let mut by_model: HashMap<String, Vec<&BatchSizeResult>> = HashMap::new();
for result in &report.results {
by_model
.entry(result.model_name.clone())
.or_insert_with(Vec::new)
.push(result);
}
for (model, results) in by_model {
md.push_str(&format!("### {} Model\n\n", model));
md.push_str("| Batch Size | VRAM (MB) | VRAM % | Latency (ms) | Throughput | Status |\n");
md.push_str("|------------|-----------|--------|--------------|------------|--------|\n");
for result in results {
let status = if result.oom_occurred {
"OOM"
} else if result.recommended {
"✅ Optimal"
} else {
"OK"
};
md.push_str(&format!(
"| {} | {:.0} | {:.1}% | {:.2} | {:.1}/sec | {} |\n",
result.batch_size,
result.vram_used_mb,
result.vram_percent,
result.latency_ms,
result.throughput_samples_per_sec,
status
));
}
md.push_str("\n");
}
md.push_str("## Updated Configuration Snippets\n\n");
for (model, batch) in &report.recommendations {
md.push_str(&format!("### {} Configuration\n\n", model));
md.push_str("```rust\n");
match model.as_str() {
"TFT" => {
md.push_str(&format!(
"TFTConfig {{\n batch_size: {},\n // ... other fields\n}}\n",
batch
));
},
"MAMBA-2" => {
md.push_str(&format!(
"Mamba2Config {{\n batch_size: {},\n // ... other fields\n}}\n",
batch
));
},
"Liquid" => {
md.push_str(&format!(
"// Note: Liquid processes samples sequentially\n// Batch size {} tested for CPU efficiency\n",
batch
));
},
_ => {},
}
md.push_str("```\n\n");
}
md.push_str("## Notes\n\n");
for note in &report.notes {
md.push_str(&format!("- {}\n", note));
}
md
}
fn main() -> Result<(), Box<dyn std::error::Error>> {
println!("=== GPU Batch Size Optimization for RTX 3050 Ti ===\n");
// Query GPU info
let gpu_name = query_gpu_name()?;
let total_vram_gb = query_total_vram_gb()?;
println!("GPU: {}", gpu_name);
println!("Total VRAM: {:.1} GB\n", total_vram_gb);
let device = Device::cuda_if_available(0).unwrap_or(Device::Cpu);
println!("Device: {:?}\n", device);
let mut results = Vec::new();
let mut notes = Vec::new();
// Test TFT with batch sizes: 16, 32, 64, 128
println!("Testing TFT model...");
for batch_size in [16, 32, 64, 128] {
match benchmark_tft_batch(batch_size, &device) {
Ok(result) => {
println!(
" Batch {}: VRAM={:.0}MB ({:.1}%), Throughput={:.1}/sec, Status={}",
result.batch_size,
result.vram_used_mb,
result.vram_percent,
result.throughput_samples_per_sec,
if result.oom_occurred { "OOM" } else { "OK" }
);
results.push(result);
},
Err(e) => {
eprintln!(" Error: {}", e);
},
}
}
// Test MAMBA-2 with batch sizes: 8, 16, 32
println!("\nTesting MAMBA-2 model...");
for batch_size in [8, 16, 32] {
match benchmark_mamba_batch(batch_size, &device) {
Ok(result) => {
println!(
" Batch {}: VRAM={:.0}MB ({:.1}%), Throughput={:.1}/sec, Status={}",
result.batch_size,
result.vram_used_mb,
result.vram_percent,
result.throughput_samples_per_sec,
if result.oom_occurred { "OOM" } else { "OK" }
);
results.push(result);
},
Err(e) => {
eprintln!(" Error: {}", e);
},
}
}
// Test Liquid with batch sizes: 16, 32, 64
println!("\nTesting Liquid model (CPU)...");
for batch_size in [16, 32, 64] {
match benchmark_liquid_batch(batch_size) {
Ok(result) => {
println!(
" Batch {}: Throughput={:.1}/sec, Status={}",
result.batch_size,
result.throughput_samples_per_sec,
if result.oom_occurred { "OOM" } else { "OK" }
);
results.push(result);
},
Err(e) => {
eprintln!(" Error: {}", e);
},
}
}
// Determine optimal batch sizes
let mut recommendations = HashMap::new();
// TFT: Find largest batch size with <90% VRAM and no OOM
if let Some(best) = results
.iter()
.filter(|r| r.model_name == "TFT" && !r.oom_occurred && r.vram_percent < 90.0)
.max_by_key(|r| r.batch_size)
{
recommendations.insert("TFT".to_string(), best.batch_size);
}
// MAMBA-2: Find largest batch size with <90% VRAM and no OOM
if let Some(best) = results
.iter()
.filter(|r| r.model_name == "MAMBA-2" && !r.oom_occurred && r.vram_percent < 90.0)
.max_by_key(|r| r.batch_size)
{
recommendations.insert("MAMBA-2".to_string(), best.batch_size);
}
// Liquid: Find highest throughput (CPU-based)
if let Some(best) = results
.iter()
.filter(|r| r.model_name == "Liquid" && !r.oom_occurred)
.max_by(|a, b| {
a.throughput_samples_per_sec
.partial_cmp(&b.throughput_samples_per_sec)
.unwrap()
})
{
recommendations.insert("Liquid".to_string(), best.batch_size);
}
// Add notes
notes.push(format!(
"Tested on {} with {:.1}GB VRAM",
gpu_name, total_vram_gb
));
notes.push("Batch sizes optimized for <90% VRAM usage".to_string());
notes.push("Liquid model runs on CPU (no GPU VRAM usage)".to_string());
notes.push(format!(
"Warmup iterations: {}, Benchmark iterations: {}",
WARMUP_ITERATIONS, BENCHMARK_ITERATIONS
));
let report = OptimizationReport {
timestamp: chrono::Utc::now().to_rfc3339(),
gpu_name,
vram_total_gb: total_vram_gb,
results,
recommendations: recommendations.clone(),
notes,
};
// Generate markdown report
let md_content = generate_report(&report);
// Write report
let mut file = File::create("BATCH_SIZE_OPTIMIZATION_REPORT.md")?;
file.write_all(md_content.as_bytes())?;
println!("\n=== Optimization Complete ===\n");
println!("Recommendations:");
for (model, batch) in &recommendations {
println!(" {} -> batch_size = {}", model, batch);
}
println!("\nReport written to: BATCH_SIZE_OPTIMIZATION_REPORT.md");
Ok(())
}