100 lines
2.9 KiB
Rust
100 lines
2.9 KiB
Rust
//! Distributed training benchmarks
|
|
|
|
use crate::{
|
|
BenchmarkConfig, BenchmarkMeasurement, BenchmarkResults, BenchmarkStatistics,
|
|
time_benchmark_async,
|
|
};
|
|
use anyhow::Result;
|
|
use std::collections::HashMap;
|
|
use tracing::info;
|
|
|
|
pub struct DistributedTrainingBenchmarks;
|
|
|
|
impl DistributedTrainingBenchmarks {
|
|
pub fn new() -> Self {
|
|
Self
|
|
}
|
|
|
|
pub async fn run_all(&self, config: &BenchmarkConfig) -> Result<BenchmarkResults> {
|
|
info!("Starting distributed training benchmarks");
|
|
|
|
let mut results = BenchmarkResults::new("DistributedTraining".to_string(), config.clone());
|
|
let mut all_measurements: HashMap<String, Vec<BenchmarkMeasurement>> = HashMap::new();
|
|
|
|
// Multi-GPU scaling benchmarks
|
|
all_measurements.insert(
|
|
"multi_gpu_scaling".to_string(),
|
|
self.benchmark_multi_gpu_scaling(config).await?,
|
|
);
|
|
|
|
// Gradient synchronization
|
|
all_measurements.insert(
|
|
"gradient_sync".to_string(),
|
|
self.benchmark_gradient_sync(config).await?,
|
|
);
|
|
|
|
// Calculate statistics
|
|
for (name, measurements) in all_measurements {
|
|
match BenchmarkStatistics::from_measurements(&measurements) {
|
|
Ok(stats) => results.add_benchmark(name, stats),
|
|
Err(e) => results.add_warning(format!("Failed to calculate stats: {e}")),
|
|
}
|
|
}
|
|
|
|
results.mark_complete();
|
|
Ok(results)
|
|
}
|
|
|
|
async fn benchmark_multi_gpu_scaling(
|
|
&self,
|
|
config: &BenchmarkConfig,
|
|
) -> Result<Vec<BenchmarkMeasurement>> {
|
|
let mut measurements = Vec::with_capacity(config.measurement_iterations);
|
|
|
|
for _ in 0..config.measurement_iterations {
|
|
let measurement = time_benchmark_async("multi_gpu_scaling".to_string(), || async {
|
|
self.multi_gpu_training_step().await?;
|
|
Ok(())
|
|
})
|
|
.await?;
|
|
measurements.push(measurement);
|
|
}
|
|
|
|
Ok(measurements)
|
|
}
|
|
|
|
async fn benchmark_gradient_sync(
|
|
&self,
|
|
config: &BenchmarkConfig,
|
|
) -> Result<Vec<BenchmarkMeasurement>> {
|
|
let mut measurements = Vec::with_capacity(config.measurement_iterations);
|
|
|
|
for _ in 0..config.measurement_iterations {
|
|
let measurement = time_benchmark_async("gradient_sync".to_string(), || async {
|
|
self.gradient_synchronization().await?;
|
|
Ok(())
|
|
})
|
|
.await?;
|
|
measurements.push(measurement);
|
|
}
|
|
|
|
Ok(measurements)
|
|
}
|
|
|
|
async fn multi_gpu_training_step(&self) -> Result<()> {
|
|
tokio::time::sleep(std::time::Duration::from_millis(200)).await;
|
|
Ok(())
|
|
}
|
|
|
|
async fn gradient_synchronization(&self) -> Result<()> {
|
|
tokio::time::sleep(std::time::Duration::from_millis(50)).await;
|
|
Ok(())
|
|
}
|
|
}
|
|
|
|
impl Default for DistributedTrainingBenchmarks {
|
|
fn default() -> Self {
|
|
Self::new()
|
|
}
|
|
}
|