Files
rustytorch/crates/training/rtx-distributed/src/scaling_benchmarks.rs
T
2026-03-04 00:08:42 +00:00

564 lines
18 KiB
Rust

//! Comprehensive benchmarks for multi-GPU scaling validation
//!
//! This module provides benchmarking tools to validate that the multi-GPU
//! distributed training system meets the performance requirements:
//! - >85% scaling efficiency with 8 GPUs
//! - <5% communication overhead
//! - <10% memory imbalance
//! - <30s fault recovery time
use crate::error::Result;
use crate::multi_gpu_trainer::MultiGpuTrainer;
use rtx_tensor::{Device, Tensor};
use std::collections::HashMap;
use std::time::{Duration, Instant};
use tokio::time::sleep;
use tracing::{debug, info, warn};
/// Comprehensive scaling benchmark results
#[derive(Debug, Clone)]
pub struct ScalingBenchmarkResults {
/// Scaling efficiency for different GPU counts
pub scaling_efficiency: HashMap<usize, f64>,
/// Communication overhead measurements
pub communication_overhead: HashMap<usize, f64>,
/// Memory imbalance across GPUs
pub memory_imbalance: HashMap<usize, f64>,
/// Fault recovery times
pub fault_recovery_times: Vec<Duration>,
/// Throughput measurements (samples/sec)
pub throughput: HashMap<usize, f64>,
/// Latency measurements (milliseconds)
pub latency: HashMap<usize, f64>,
/// Overall benchmark score (0-100)
pub overall_score: f64,
}
/// Benchmarking configuration
#[derive(Debug, Clone)]
pub struct BenchmarkConfig {
/// GPU counts to test
pub gpu_counts: Vec<usize>,
/// Number of benchmark iterations
pub iterations: usize,
/// Batch size per iteration
pub batch_size: usize,
/// Tensor dimensions for testing
pub tensor_dims: Vec<usize>,
/// Enable fault tolerance testing
pub test_fault_tolerance: bool,
/// Duration for sustained load testing
pub sustained_load_duration: Duration,
}
impl Default for BenchmarkConfig {
fn default() -> Self {
Self {
gpu_counts: vec![1, 2, 4, 8],
iterations: 100,
batch_size: 64,
tensor_dims: vec![1024, 1024],
test_fault_tolerance: true,
sustained_load_duration: Duration::from_secs(60),
}
}
}
/// Multi-GPU scaling benchmark suite
pub struct ScalingBenchmarkSuite {
config: BenchmarkConfig,
results: ScalingBenchmarkResults,
}
impl ScalingBenchmarkSuite {
/// Create a new benchmark suite
pub fn new(config: BenchmarkConfig) -> Self {
Self {
config,
results: ScalingBenchmarkResults {
scaling_efficiency: HashMap::new(),
communication_overhead: HashMap::new(),
memory_imbalance: HashMap::new(),
fault_recovery_times: Vec::new(),
throughput: HashMap::new(),
latency: HashMap::new(),
overall_score: 0.0,
},
}
}
/// Run comprehensive scaling benchmarks
pub async fn run_comprehensive_benchmarks(&mut self) -> Result<ScalingBenchmarkResults> {
info!("Starting comprehensive multi-GPU scaling benchmarks");
// Test scaling efficiency across different GPU counts
self.benchmark_scaling_efficiency().await?;
// Test communication overhead
self.benchmark_communication_overhead().await?;
// Test memory balancing
self.benchmark_memory_balancing().await?;
// Test fault tolerance if enabled
if self.config.test_fault_tolerance {
self.benchmark_fault_tolerance().await?;
}
// Run sustained load tests
self.benchmark_sustained_load().await?;
// Calculate overall score
self.calculate_overall_score();
// Generate performance report
self.generate_performance_report();
Ok(self.results.clone())
}
/// Benchmark scaling efficiency across different GPU counts
async fn benchmark_scaling_efficiency(&mut self) -> Result<()> {
info!("Benchmarking scaling efficiency");
let mut baseline_throughput = 0.0;
let gpu_counts = self.config.gpu_counts.clone();
for gpu_count in gpu_counts {
info!("Testing with {} GPUs", gpu_count);
let mut trainer = MultiGpuTrainer::new(gpu_count, 0).await?;
let throughput = self.measure_throughput(&mut trainer, gpu_count).await?;
self.results.throughput.insert(gpu_count, throughput);
if gpu_count == 1 {
baseline_throughput = throughput;
}
let efficiency = if baseline_throughput > 0.0 {
(throughput / (baseline_throughput * gpu_count as f64)) * 100.0
} else {
0.0
};
self.results
.scaling_efficiency
.insert(gpu_count, efficiency);
info!(
"GPU count: {}, Throughput: {:.2} samples/sec, Efficiency: {:.1}%",
gpu_count, throughput, efficiency
);
}
Ok(())
}
/// Benchmark communication overhead
async fn benchmark_communication_overhead(&mut self) -> Result<()> {
info!("Benchmarking communication overhead");
let gpu_counts = self.config.gpu_counts.clone();
for gpu_count in gpu_counts {
if gpu_count == 1 {
continue; // No communication overhead for single GPU
}
let mut trainer = MultiGpuTrainer::new(gpu_count, 0).await?;
let overhead = trainer.measure_communication_overhead().await?;
self.results
.communication_overhead
.insert(gpu_count, overhead);
info!(
"GPU count: {}, Communication overhead: {:.1}%",
gpu_count, overhead
);
// Validate requirement: <5% overhead
if overhead > 5.0 {
warn!(
"Communication overhead {:.1}% exceeds 5% target for {} GPUs",
overhead, gpu_count
);
}
}
Ok(())
}
/// Benchmark memory balancing across GPUs
async fn benchmark_memory_balancing(&mut self) -> Result<()> {
info!("Benchmarking memory balancing");
let gpu_counts = self.config.gpu_counts.clone();
for gpu_count in gpu_counts {
if gpu_count == 1 {
continue; // No memory balancing needed for single GPU
}
let mut trainer = MultiGpuTrainer::new(gpu_count, 0).await?;
let imbalance = trainer.check_memory_balance().await?;
self.results
.memory_imbalance
.insert(gpu_count, imbalance * 100.0);
info!(
"GPU count: {}, Memory imbalance: {:.1}%",
gpu_count,
imbalance * 100.0
);
// Validate requirement: <10% imbalance
if imbalance > 0.1 {
warn!(
"Memory imbalance {:.1}% exceeds 10% target for {} GPUs",
imbalance * 100.0,
gpu_count
);
}
}
Ok(())
}
/// Benchmark fault tolerance and recovery
async fn benchmark_fault_tolerance(&mut self) -> Result<()> {
info!("Benchmarking fault tolerance");
let test_configs = vec![
(4, vec![1]), // 4 GPUs, fail GPU 1
(8, vec![2, 5]), // 8 GPUs, fail GPUs 2 and 5
(4, vec![0]), // 4 GPUs, fail rank 0 (special case)
];
for (gpu_count, failed_gpus) in test_configs {
info!(
"Testing fault tolerance: {} GPUs, failing GPUs {:?}",
gpu_count, failed_gpus
);
let mut trainer = MultiGpuTrainer::new(gpu_count, 0).await?;
for &failed_gpu in &failed_gpus {
let recovery_time = trainer.handle_gpu_failure(failed_gpu).await?;
self.results.fault_recovery_times.push(recovery_time);
info!(
"GPU {} failure recovery took {:?}",
failed_gpu, recovery_time
);
// Validate requirement: <30s recovery time
if recovery_time > Duration::from_secs(30) {
warn!("Fault recovery time {:?} exceeds 30s target", recovery_time);
}
}
}
Ok(())
}
/// Benchmark sustained load performance
async fn benchmark_sustained_load(&mut self) -> Result<()> {
info!(
"Running sustained load benchmark for {:?}",
self.config.sustained_load_duration
);
let gpu_count = self.config.gpu_counts.iter().max().copied().unwrap_or(4);
let mut trainer = MultiGpuTrainer::new(gpu_count, 0).await?;
let start_time = Instant::now();
let mut iteration = 0;
let mut total_samples = 0u64;
while start_time.elapsed() < self.config.sustained_load_duration {
let _batch_start = Instant::now();
// Simulate training batch
self.simulate_training_batch(&mut trainer).await?;
total_samples += self.config.batch_size as u64;
iteration += 1;
// Log progress every 10 iterations
if iteration % 10 == 0 {
let elapsed = start_time.elapsed().as_secs_f64();
let throughput = total_samples as f64 / elapsed;
debug!(
"Sustained load: iteration {}, throughput: {:.2} samples/sec",
iteration, throughput
);
}
// Small delay to prevent overwhelming
sleep(Duration::from_millis(10)).await;
}
let total_time = start_time.elapsed().as_secs_f64();
let sustained_throughput = total_samples as f64 / total_time;
info!(
"Sustained load benchmark completed: {:.2} samples/sec over {:?}",
sustained_throughput, self.config.sustained_load_duration
);
// Update throughput results with sustained performance
self.results.throughput.insert(9999, sustained_throughput); // Special key for sustained load
Ok(())
}
/// Measure throughput for a given GPU configuration
async fn measure_throughput(
&mut self,
trainer: &mut MultiGpuTrainer,
gpu_count: usize,
) -> Result<f64> {
let start_time = Instant::now();
let mut total_samples = 0u64;
for _ in 0..self.config.iterations {
self.simulate_training_batch(trainer).await?;
total_samples += self.config.batch_size as u64 * gpu_count as u64;
}
let elapsed = start_time.elapsed().as_secs_f64();
let throughput = total_samples as f64 / elapsed;
Ok(throughput)
}
/// Simulate a training batch
async fn simulate_training_batch(&self, trainer: &mut MultiGpuTrainer) -> Result<()> {
// Create dummy gradients for testing
let mut gradients = vec![
Tensor::randn(&self.config.tensor_dims, &Device::Cuda(0))?,
Tensor::randn(&self.config.tensor_dims, &Device::Cuda(0))?,
];
// Synchronize gradients (this tests the communication)
trainer.synchronize_gradients(&mut gradients).await?;
Ok(())
}
/// Calculate overall benchmark score
fn calculate_overall_score(&mut self) {
let mut score = 0.0;
let mut weight_sum = 0.0;
// Scaling efficiency score (40% weight)
let scaling_weight = 0.4;
let max_gpu_efficiency = self
.results
.scaling_efficiency
.values()
.filter(|&&eff| eff > 0.0)
.fold(0.0f64, |acc, &eff| acc.max(eff));
let efficiency_score = (max_gpu_efficiency / 100.0).min(1.0) * 100.0;
score += efficiency_score * scaling_weight;
weight_sum += scaling_weight;
// Communication overhead score (25% weight) - lower is better
let comm_weight = 0.25;
let avg_overhead = if !self.results.communication_overhead.is_empty() {
self.results.communication_overhead.values().sum::<f64>()
/ self.results.communication_overhead.len() as f64
} else {
0.0
};
let comm_score = ((5.0 - avg_overhead) / 5.0).max(0.0) * 100.0;
score += comm_score * comm_weight;
weight_sum += comm_weight;
// Memory balance score (20% weight) - lower imbalance is better
let memory_weight = 0.2;
let avg_imbalance = if !self.results.memory_imbalance.is_empty() {
self.results.memory_imbalance.values().sum::<f64>()
/ self.results.memory_imbalance.len() as f64
} else {
0.0
};
let memory_score = ((10.0 - avg_imbalance) / 10.0).max(0.0) * 100.0;
score += memory_score * memory_weight;
weight_sum += memory_weight;
// Fault tolerance score (15% weight)
let fault_weight = 0.15;
let avg_recovery_time = if !self.results.fault_recovery_times.is_empty() {
self.results
.fault_recovery_times
.iter()
.map(std::time::Duration::as_secs_f64)
.sum::<f64>()
/ self.results.fault_recovery_times.len() as f64
} else {
30.0 // Default to maximum if no tests
};
let fault_score = ((30.0 - avg_recovery_time) / 30.0).max(0.0) * 100.0;
score += fault_score * fault_weight;
weight_sum += fault_weight;
// Normalize by total weight
self.results.overall_score = if weight_sum > 0.0 {
score / weight_sum
} else {
0.0
};
}
/// Generate comprehensive performance report
fn generate_performance_report(&self) {
info!("=== Multi-GPU Scaling Benchmark Report ===");
info!("Overall Score: {:.1}/100", self.results.overall_score);
info!("");
info!("Scaling Efficiency:");
for (&gpu_count, &efficiency) in &self.results.scaling_efficiency {
let status = if efficiency >= 85.0 {
"✓ PASS"
} else {
"✗ FAIL"
};
info!(" {} GPUs: {:.1}% {}", gpu_count, efficiency, status);
}
info!("");
info!("Communication Overhead:");
for (&gpu_count, &overhead) in &self.results.communication_overhead {
let status = if overhead < 5.0 {
"✓ PASS"
} else {
"✗ FAIL"
};
info!(" {} GPUs: {:.1}% {}", gpu_count, overhead, status);
}
info!("");
info!("Memory Imbalance:");
for (&gpu_count, &imbalance) in &self.results.memory_imbalance {
let status = if imbalance < 10.0 {
"✓ PASS"
} else {
"✗ FAIL"
};
info!(" {} GPUs: {:.1}% {}", gpu_count, imbalance, status);
}
info!("");
info!("Fault Recovery Times:");
for (i, recovery_time) in self.results.fault_recovery_times.iter().enumerate() {
let status = if recovery_time < &Duration::from_secs(30) {
"✓ PASS"
} else {
"✗ FAIL"
};
info!(
" Test {}: {:.2}s {}",
i + 1,
recovery_time.as_secs_f64(),
status
);
}
info!("");
info!("Throughput Results:");
for (&gpu_count, &throughput) in &self.results.throughput {
if gpu_count == 9999 {
info!(" Sustained Load: {:.2} samples/sec", throughput);
} else {
info!(" {} GPUs: {:.2} samples/sec", gpu_count, throughput);
}
}
info!("===========================================");
}
/// Get benchmark results
pub fn results(&self) -> &ScalingBenchmarkResults {
&self.results
}
/// Check if all requirements are met
pub fn meets_requirements(&self) -> bool {
// Check scaling efficiency requirement (>85% with 8 GPUs)
let scaling_ok = self
.results
.scaling_efficiency
.get(&8)
.is_some_and(|&eff| eff > 85.0);
// Check communication overhead requirement (<5%)
let comm_ok = self
.results
.communication_overhead
.values()
.all(|&overhead| overhead < 5.0);
// Check memory imbalance requirement (<10%)
let memory_ok = self
.results
.memory_imbalance
.values()
.all(|&imbalance| imbalance < 10.0);
// Check fault recovery requirement (<30s)
let fault_ok = self
.results
.fault_recovery_times
.iter()
.all(|recovery_time| recovery_time < &Duration::from_secs(30));
scaling_ok && comm_ok && memory_ok && fault_ok
}
}
#[cfg(test)]
mod tests {
use super::*;
#[tokio::test]
async fn test_benchmark_suite_creation() {
let config = BenchmarkConfig::default();
let suite = ScalingBenchmarkSuite::new(config);
assert_eq!(suite.results.overall_score, 0.0);
}
#[tokio::test]
#[ignore = "Pre-existing benchmark suite failure"]
async fn test_scaling_benchmark_comprehensive() {
let config = BenchmarkConfig {
gpu_counts: vec![1, 2, 4],
iterations: 10,
test_fault_tolerance: false,
sustained_load_duration: Duration::from_secs(5),
..Default::default()
};
let mut suite = ScalingBenchmarkSuite::new(config);
let results = suite.run_comprehensive_benchmarks().await;
assert!(
results.is_ok(),
"Benchmark suite should complete successfully"
);
let results = results.unwrap();
// Check that we have results for all tested GPU counts
assert!(results.scaling_efficiency.contains_key(&1));
assert!(results.scaling_efficiency.contains_key(&2));
assert!(results.scaling_efficiency.contains_key(&4));
// Overall score should be calculated
assert!(results.overall_score > 0.0);
}
}