Omar Sobh
623e6679d5
ring_allreduce() contained its own simulation that multiplied each gradient value by world_size (to fake an AllReduce sum), bypassing the ProcessGroup dispatch entirely. This meant the overlapped synchronization path never used NCCL or RNCCL, even when those features were compiled in. Replace the hand-rolled simulation with a call to self.process_group.allreduce(tensor, ReduceOp::Sum) so the overlapped path uses the same backend as synchronize_gradients_sequential. The communication latency sleep is kept for benchmarking purposes. Add test_ring_allreduce_matches_sequential_path to verify both paths produce identical gradient values under CPU simulation. Closes #10
847 lines
28 KiB
Rust
847 lines
28 KiB
Rust
//! Multi-GPU trainer implementation for distributed training
|
|
//!
|
|
//! This module provides the core MultiGpuTrainer implementation that coordinates
|
|
//! training across multiple GPUs with near-linear scaling efficiency.
|
|
|
|
use crate::Backend;
|
|
use crate::comm::{CommunicationPrimitive, ReduceOp};
|
|
use crate::error::{DistributedError, Result};
|
|
use crate::group::ProcessGroup;
|
|
use rtx_tensor::{Device, Tensor};
|
|
use std::sync::Arc;
|
|
use std::time::{Duration, Instant};
|
|
use tokio::sync::RwLock;
|
|
use tracing::{debug, info, warn};
|
|
|
|
/// Multi-GPU trainer for coordinating distributed training
|
|
#[derive(Debug)]
|
|
pub struct MultiGpuTrainer {
|
|
/// Local GPU rank (0 to world_size-1)
|
|
pub local_rank: usize,
|
|
/// Total number of GPUs across all nodes
|
|
pub world_size: usize,
|
|
/// Process group for communication
|
|
pub process_group: ProcessGroup,
|
|
/// Available GPU devices
|
|
pub devices: Vec<Device>,
|
|
/// Training metrics
|
|
pub metrics: Arc<RwLock<TrainingMetrics>>,
|
|
/// Load balancer for work distribution
|
|
pub load_balancer: LoadBalancer,
|
|
/// Fault tolerance manager
|
|
pub fault_tolerance: FaultTolerance,
|
|
/// Scaling optimizer for communication overlap
|
|
pub scaling_optimizer: ScalingOptimizer,
|
|
/// Dynamic loss scaling for mixed precision training
|
|
pub loss_scaler: Arc<RwLock<DynamicLossScaler>>,
|
|
}
|
|
|
|
/// Training metrics for performance monitoring
|
|
#[derive(Debug, Default, Clone)]
|
|
pub struct TrainingMetrics {
|
|
/// Time spent synchronizing gradients (milliseconds)
|
|
pub gradient_sync_time_ms: f64,
|
|
/// Communication overhead as percentage of total time
|
|
pub communication_overhead_percent: f64,
|
|
/// Memory imbalance ratio across GPUs (0.0 = perfect balance)
|
|
pub memory_imbalance_ratio: f64,
|
|
/// Scaling efficiency (1.0 = linear scaling)
|
|
pub scaling_efficiency: f64,
|
|
/// Time to recover from GPU failure (milliseconds)
|
|
pub fault_recovery_time_ms: f64,
|
|
/// Total training steps completed
|
|
pub training_steps: u64,
|
|
/// Average throughput (samples per second)
|
|
pub throughput_samples_per_sec: f64,
|
|
}
|
|
|
|
/// Load balancer for optimal work distribution across GPUs
|
|
#[derive(Debug)]
|
|
pub struct LoadBalancer {
|
|
/// Current GPU utilization percentages (0.0 to 1.0)
|
|
pub gpu_utilizations: Vec<f64>,
|
|
/// Memory usage per GPU in bytes
|
|
pub memory_usage: Vec<u64>,
|
|
/// Maximum allowed memory imbalance (0.1 = 10%)
|
|
pub imbalance_threshold: f64,
|
|
/// Work distribution weights per GPU
|
|
pub distribution_weights: Vec<f64>,
|
|
}
|
|
|
|
/// Fault tolerance manager for handling GPU failures
|
|
#[derive(Debug)]
|
|
pub struct FaultTolerance {
|
|
/// Health status of each GPU (true = healthy)
|
|
pub health_monitors: Vec<bool>,
|
|
/// Maximum time allowed for recovery
|
|
pub recovery_timeout: Duration,
|
|
/// List of failed GPU ranks
|
|
pub failed_gpus: Vec<usize>,
|
|
/// Checkpoint interval for fault recovery
|
|
pub checkpoint_interval: Duration,
|
|
/// Last checkpoint timestamp
|
|
pub last_checkpoint: Instant,
|
|
}
|
|
|
|
/// Scaling optimizer for communication/computation overlap
|
|
#[derive(Debug)]
|
|
pub struct ScalingOptimizer {
|
|
/// Whether communication overlap is enabled
|
|
pub communication_overlap_enabled: bool,
|
|
/// Number of CUDA streams for communication
|
|
pub communication_streams: usize,
|
|
/// Number of CUDA streams for computation
|
|
pub compute_streams: usize,
|
|
/// Ring all-reduce buffer size
|
|
pub ring_buffer_size: usize,
|
|
/// Gradient compression enabled
|
|
pub gradient_compression: bool,
|
|
}
|
|
|
|
/// Dynamic loss scaler for mixed precision training
|
|
#[derive(Debug)]
|
|
pub struct DynamicLossScaler {
|
|
/// Current loss scale factor
|
|
pub scale: f32,
|
|
/// Growth factor for increasing scale
|
|
pub growth_factor: f32,
|
|
/// Backoff factor for decreasing scale
|
|
pub backoff_factor: f32,
|
|
/// Number of steps without overflow
|
|
pub steps_without_overflow: usize,
|
|
/// Steps required before growing scale
|
|
pub growth_interval: usize,
|
|
/// Minimum allowed scale
|
|
pub min_scale: f32,
|
|
/// Maximum allowed scale
|
|
pub max_scale: f32,
|
|
/// Current mixed precision mode
|
|
pub precision_mode: String,
|
|
/// Whether to check for overflow
|
|
pub enabled: bool,
|
|
}
|
|
|
|
impl MultiGpuTrainer {
|
|
/// Create a new multi-GPU trainer
|
|
pub async fn new(world_size: usize, local_rank: usize) -> Result<Self> {
|
|
info!(
|
|
"Initializing MultiGpuTrainer with world_size={}, local_rank={}",
|
|
world_size, local_rank
|
|
);
|
|
|
|
// Validate input parameters
|
|
if world_size == 0 {
|
|
return Err(DistributedError::configuration("world_size must be > 0"));
|
|
}
|
|
if local_rank >= world_size {
|
|
return Err(DistributedError::configuration(format!(
|
|
"local_rank ({local_rank}) must be < world_size ({world_size})"
|
|
)));
|
|
}
|
|
|
|
// Initialize process group with NCCL backend
|
|
use crate::group::WorldInfo;
|
|
let world_info = WorldInfo::new(world_size as i32, local_rank as i32, Backend::Nccl);
|
|
let process_group = ProcessGroup::new(Backend::Nccl, world_info)?;
|
|
|
|
// Create GPU devices
|
|
let devices = (0..world_size).map(Device::Cuda).collect();
|
|
|
|
// Initialize load balancer
|
|
let load_balancer = LoadBalancer::new(world_size);
|
|
|
|
// Initialize fault tolerance
|
|
let fault_tolerance = FaultTolerance::new(world_size);
|
|
|
|
// Initialize scaling optimizer
|
|
let scaling_optimizer = ScalingOptimizer::new();
|
|
|
|
// Initialize dynamic loss scaler
|
|
let loss_scaler = DynamicLossScaler::new();
|
|
|
|
let trainer = Self {
|
|
local_rank,
|
|
world_size,
|
|
process_group,
|
|
devices,
|
|
metrics: Arc::new(RwLock::new(TrainingMetrics::default())),
|
|
load_balancer,
|
|
fault_tolerance,
|
|
scaling_optimizer,
|
|
loss_scaler: Arc::new(RwLock::new(loss_scaler)),
|
|
};
|
|
|
|
info!("MultiGpuTrainer initialized successfully");
|
|
Ok(trainer)
|
|
}
|
|
|
|
/// Synchronize gradients across all GPUs using ring all-reduce
|
|
pub async fn synchronize_gradients(&mut self, gradients: &mut [Tensor]) -> Result<()> {
|
|
let start_time = Instant::now();
|
|
debug!(
|
|
"Starting gradient synchronization across {} GPUs",
|
|
self.world_size
|
|
);
|
|
|
|
// Check for failed GPUs before synchronization
|
|
self.fault_tolerance.check_gpu_health().await?;
|
|
|
|
if self.scaling_optimizer.communication_overlap_enabled {
|
|
self.synchronize_gradients_overlapped(gradients).await?;
|
|
} else {
|
|
self.synchronize_gradients_sequential(gradients).await?;
|
|
}
|
|
|
|
// Update metrics
|
|
let sync_time = start_time.elapsed().as_millis() as f64;
|
|
{
|
|
let mut metrics = self.metrics.write().await;
|
|
metrics.gradient_sync_time_ms = sync_time;
|
|
}
|
|
|
|
debug!("Gradient synchronization completed in {:.2}ms", sync_time);
|
|
Ok(())
|
|
}
|
|
|
|
/// Sequential gradient synchronization (no overlap)
|
|
async fn synchronize_gradients_sequential(&mut self, gradients: &mut [Tensor]) -> Result<()> {
|
|
let num_gradients = gradients.len();
|
|
for (i, gradient) in gradients.iter_mut().enumerate() {
|
|
debug!(
|
|
"Synchronizing gradient tensor {} of {}",
|
|
i + 1,
|
|
num_gradients
|
|
);
|
|
|
|
// Perform all-reduce to sum gradients across all GPUs
|
|
self.process_group
|
|
.allreduce(gradient, ReduceOp::Sum)
|
|
.await?;
|
|
|
|
// Average the gradients by dividing by world size
|
|
let world_size = self.world_size as f32;
|
|
*gradient = gradient.div_scalar(world_size)?;
|
|
}
|
|
Ok(())
|
|
}
|
|
|
|
/// Overlapped gradient synchronization (communication + computation)
|
|
async fn synchronize_gradients_overlapped(&mut self, gradients: &mut [Tensor]) -> Result<()> {
|
|
info!("Using ring AllReduce with communication overlap");
|
|
|
|
let world_size = self.world_size;
|
|
let rank = self.local_rank;
|
|
let buffer_size = self.scaling_optimizer.ring_buffer_size;
|
|
let gradient_count = gradients.len();
|
|
|
|
for (i, gradient) in gradients.iter_mut().enumerate() {
|
|
debug!(
|
|
"Ring AllReduce for gradient tensor {} of {}",
|
|
i + 1,
|
|
gradient_count
|
|
);
|
|
|
|
// Apply gradient compression if enabled
|
|
if self.scaling_optimizer.gradient_compression {
|
|
self.apply_gradient_compression(gradient).await?;
|
|
}
|
|
|
|
// Perform ring AllReduce algorithm
|
|
self.ring_allreduce(gradient, rank, world_size, buffer_size)
|
|
.await?;
|
|
|
|
// Average the gradients
|
|
let world_size_f32 = world_size as f32;
|
|
*gradient = gradient.div_scalar(world_size_f32)?;
|
|
}
|
|
|
|
Ok(())
|
|
}
|
|
|
|
/// Ring AllReduce implementation with communication overlap
|
|
async fn ring_allreduce(
|
|
&mut self,
|
|
tensor: &mut Tensor,
|
|
_rank: usize,
|
|
world_size: usize,
|
|
buffer_size: usize,
|
|
) -> Result<()> {
|
|
let tensor_data = tensor.data()?;
|
|
let total_elements = tensor_data.len();
|
|
|
|
if total_elements == 0 {
|
|
return Ok(());
|
|
}
|
|
|
|
let chunk_size = total_elements.div_ceil(world_size);
|
|
debug!(
|
|
"Ring AllReduce: tensor size={}, chunk_size={}, world_size={}",
|
|
total_elements, chunk_size, world_size
|
|
);
|
|
|
|
// Delegate to the process group for the actual reduction.
|
|
// This dispatches to NCCL or RNCCL when the corresponding feature is compiled in,
|
|
// ensuring the overlapped path uses the same communication backend as the
|
|
// sequential path (synchronize_gradients_sequential). Without this delegation
|
|
// the ring simulation bypassed NCCL even when it was available.
|
|
self.process_group.allreduce(tensor, ReduceOp::Sum).await?;
|
|
|
|
// Simulate ring communication latency for benchmarking purposes
|
|
let total_comm_time_us = (total_elements as f64 / buffer_size as f64 * 20.0) as u64;
|
|
tokio::time::sleep(Duration::from_micros(total_comm_time_us.max(10))).await;
|
|
|
|
debug!("Ring AllReduce completed successfully");
|
|
Ok(())
|
|
}
|
|
|
|
/// Apply gradient compression if enabled
|
|
async fn apply_gradient_compression(&self, gradient: &mut Tensor) -> Result<()> {
|
|
if !self.scaling_optimizer.gradient_compression {
|
|
return Ok(());
|
|
}
|
|
|
|
debug!("Applying gradient compression");
|
|
|
|
// Simple gradient compression: quantize to reduce precision
|
|
let mut data = gradient.data()?;
|
|
|
|
// Apply simple quantization (reduce precision)
|
|
for value in &mut data {
|
|
// Quantize to reduce bandwidth (simple approach)
|
|
*value = (*value * 1000.0).round() / 1000.0;
|
|
}
|
|
|
|
*gradient = Tensor::from_data(data, gradient.shape().dims().to_vec(), gradient.device())?;
|
|
|
|
// Simulate compression overhead
|
|
tokio::time::sleep(Duration::from_micros(5)).await;
|
|
|
|
Ok(())
|
|
}
|
|
|
|
/// Measure scaling efficiency with current configuration
|
|
pub async fn measure_scaling_efficiency(&mut self, target_gpus: usize) -> Result<f64> {
|
|
info!("Measuring scaling efficiency for {} GPUs", target_gpus);
|
|
|
|
if target_gpus > self.world_size {
|
|
return Err(DistributedError::configuration(format!(
|
|
"Cannot measure {} GPU efficiency with only {} GPUs available",
|
|
target_gpus, self.world_size
|
|
)));
|
|
}
|
|
|
|
// Benchmark single GPU performance
|
|
let single_gpu_throughput = self.benchmark_single_gpu_throughput().await?;
|
|
debug!(
|
|
"Single GPU throughput: {:.2} samples/sec",
|
|
single_gpu_throughput
|
|
);
|
|
|
|
// Benchmark multi-GPU performance
|
|
let multi_gpu_throughput = self.benchmark_multi_gpu_throughput(target_gpus).await?;
|
|
debug!(
|
|
"Multi GPU ({}) throughput: {:.2} samples/sec",
|
|
target_gpus, multi_gpu_throughput
|
|
);
|
|
|
|
// Calculate efficiency: actual_speedup / ideal_speedup
|
|
let ideal_throughput = single_gpu_throughput * target_gpus as f64;
|
|
let efficiency = if ideal_throughput > 0.0 {
|
|
(multi_gpu_throughput / ideal_throughput) * 100.0
|
|
} else {
|
|
0.0
|
|
};
|
|
|
|
// Update metrics
|
|
{
|
|
let mut metrics = self.metrics.write().await;
|
|
metrics.scaling_efficiency = efficiency;
|
|
}
|
|
|
|
info!(
|
|
"Scaling efficiency with {} GPUs: {:.1}%",
|
|
target_gpus, efficiency
|
|
);
|
|
Ok(efficiency)
|
|
}
|
|
|
|
/// Handle GPU failure and attempt recovery
|
|
pub async fn handle_gpu_failure(&mut self, failed_gpu_rank: usize) -> Result<Duration> {
|
|
let start_time = Instant::now();
|
|
warn!("Handling failure of GPU rank {}", failed_gpu_rank);
|
|
|
|
// Validate failed GPU rank
|
|
if failed_gpu_rank >= self.world_size {
|
|
return Err(DistributedError::fault_tolerance(format!(
|
|
"Invalid failed GPU rank: {failed_gpu_rank}"
|
|
)));
|
|
}
|
|
|
|
// Mark GPU as failed
|
|
self.fault_tolerance.mark_gpu_failed(failed_gpu_rank)?;
|
|
|
|
// Attempt to checkpoint current state
|
|
self.create_checkpoint().await?;
|
|
|
|
// Redistribute work to remaining healthy GPUs
|
|
self.redistribute_work().await?;
|
|
|
|
// Update process group to exclude failed GPU
|
|
self.update_process_group_for_failure(failed_gpu_rank)
|
|
.await?;
|
|
|
|
let recovery_time = start_time.elapsed();
|
|
|
|
// Update metrics
|
|
{
|
|
let mut metrics = self.metrics.write().await;
|
|
metrics.fault_recovery_time_ms = recovery_time.as_millis() as f64;
|
|
}
|
|
|
|
info!(
|
|
"GPU {} failure recovery completed in {:?}",
|
|
failed_gpu_rank, recovery_time
|
|
);
|
|
Ok(recovery_time)
|
|
}
|
|
|
|
/// Check memory balance across GPUs
|
|
pub async fn check_memory_balance(&mut self) -> Result<f64> {
|
|
debug!("Checking memory balance across {} GPUs", self.world_size);
|
|
|
|
// Update memory usage statistics
|
|
self.load_balancer.update_memory_usage().await?;
|
|
|
|
let memory_usages = &self.load_balancer.memory_usage;
|
|
|
|
if memory_usages.is_empty() {
|
|
return Ok(0.0);
|
|
}
|
|
|
|
let max_usage = *memory_usages.iter().max().unwrap() as f64;
|
|
let min_usage = *memory_usages.iter().min().unwrap() as f64;
|
|
|
|
let imbalance_ratio = if max_usage > 0.0 {
|
|
(max_usage - min_usage) / max_usage
|
|
} else {
|
|
0.0
|
|
};
|
|
|
|
// Update metrics
|
|
{
|
|
let mut metrics = self.metrics.write().await;
|
|
metrics.memory_imbalance_ratio = imbalance_ratio;
|
|
}
|
|
|
|
debug!("Memory imbalance ratio: {:.1}%", imbalance_ratio * 100.0);
|
|
Ok(imbalance_ratio)
|
|
}
|
|
|
|
/// Measure communication overhead
|
|
pub async fn measure_communication_overhead(&mut self) -> Result<f64> {
|
|
debug!("Measuring communication overhead");
|
|
|
|
// Benchmark pure computation time
|
|
let computation_time = self.benchmark_computation_time().await?;
|
|
|
|
// Benchmark computation + communication time
|
|
let total_time = self.benchmark_computation_with_communication().await?;
|
|
|
|
let overhead_percent = if computation_time > 0.0 {
|
|
((total_time - computation_time) / computation_time) * 100.0
|
|
} else {
|
|
0.0
|
|
};
|
|
|
|
// Update metrics
|
|
{
|
|
let mut metrics = self.metrics.write().await;
|
|
metrics.communication_overhead_percent = overhead_percent;
|
|
}
|
|
|
|
debug!("Communication overhead: {:.1}%", overhead_percent);
|
|
Ok(overhead_percent)
|
|
}
|
|
|
|
/// Get current training metrics
|
|
pub async fn get_metrics(&self) -> TrainingMetrics {
|
|
self.metrics.read().await.clone()
|
|
}
|
|
|
|
/// Synchronize gradients with dynamic loss scaling
|
|
pub async fn synchronize_gradients_with_loss_scaling(
|
|
&mut self,
|
|
gradients: &mut [Tensor],
|
|
) -> Result<()> {
|
|
let start_time = Instant::now();
|
|
debug!("Starting gradient synchronization with loss scaling");
|
|
|
|
// Scale gradients before synchronization
|
|
{
|
|
let scaler = self.loss_scaler.read().await;
|
|
scaler.scale_gradients(gradients)?;
|
|
}
|
|
|
|
// Perform gradient synchronization
|
|
self.synchronize_gradients(gradients).await?;
|
|
|
|
// Check for overflow and update loss scale
|
|
let has_overflow = {
|
|
let scaler = self.loss_scaler.read().await;
|
|
scaler.has_overflow(gradients)?
|
|
};
|
|
|
|
// Unscale gradients if no overflow
|
|
if !has_overflow {
|
|
let scaler = self.loss_scaler.read().await;
|
|
scaler.unscale_gradients(gradients)?;
|
|
}
|
|
|
|
// Update loss scale based on overflow detection
|
|
{
|
|
let mut scaler = self.loss_scaler.write().await;
|
|
scaler.update_scale(has_overflow);
|
|
}
|
|
|
|
// Zero out gradients if overflow detected
|
|
if has_overflow {
|
|
warn!("Gradient overflow detected, zeroing gradients");
|
|
for gradient in gradients.iter_mut() {
|
|
*gradient = Tensor::zeros(gradient.shape().dims(), gradient.device())?;
|
|
}
|
|
}
|
|
|
|
let sync_time = start_time.elapsed().as_millis() as f64;
|
|
{
|
|
let mut metrics = self.metrics.write().await;
|
|
metrics.gradient_sync_time_ms = sync_time;
|
|
}
|
|
|
|
debug!(
|
|
"Gradient synchronization with loss scaling completed in {:.2}ms",
|
|
sync_time
|
|
);
|
|
Ok(())
|
|
}
|
|
|
|
/// Get current loss scale
|
|
pub async fn get_loss_scale(&self) -> Result<f32> {
|
|
let scaler = self.loss_scaler.read().await;
|
|
Ok(scaler.scale)
|
|
}
|
|
|
|
/// Set loss scale
|
|
pub async fn set_loss_scale(&mut self, scale: f32) -> Result<()> {
|
|
let mut scaler = self.loss_scaler.write().await;
|
|
scaler.scale = scale.max(scaler.min_scale).min(scaler.max_scale);
|
|
Ok(())
|
|
}
|
|
|
|
/// Set mixed precision mode
|
|
pub async fn set_mixed_precision_mode(&mut self, mode: &str) -> Result<()> {
|
|
let mut scaler = self.loss_scaler.write().await;
|
|
scaler.precision_mode = mode.to_string();
|
|
|
|
// Adjust default scale based on precision mode
|
|
match mode {
|
|
"fp16" => {
|
|
scaler.scale = 65536.0; // 2^16
|
|
scaler.enabled = true;
|
|
}
|
|
"bf16" => {
|
|
scaler.scale = 256.0; // Lower scale for bfloat16
|
|
scaler.enabled = true;
|
|
}
|
|
"fp32" => {
|
|
scaler.scale = 1.0; // No scaling needed for fp32
|
|
scaler.enabled = false;
|
|
}
|
|
_ => {
|
|
return Err(DistributedError::configuration(format!(
|
|
"Unknown precision mode: {mode}"
|
|
)));
|
|
}
|
|
}
|
|
|
|
info!(
|
|
"Mixed precision mode set to {} with scale {}",
|
|
mode, scaler.scale
|
|
);
|
|
Ok(())
|
|
}
|
|
|
|
// Private helper methods for benchmarking and recovery
|
|
|
|
async fn benchmark_single_gpu_throughput(&self) -> Result<f64> {
|
|
// Simulate single GPU benchmark
|
|
tokio::time::sleep(Duration::from_millis(10)).await;
|
|
Ok(1000.0) // samples/sec
|
|
}
|
|
|
|
async fn benchmark_multi_gpu_throughput(&self, _num_gpus: usize) -> Result<f64> {
|
|
// Simulate multi-GPU benchmark
|
|
tokio::time::sleep(Duration::from_millis(10)).await;
|
|
Ok(3500.0) // samples/sec for 4 GPUs (87.5% efficiency)
|
|
}
|
|
|
|
async fn benchmark_computation_time(&self) -> Result<f64> {
|
|
// Simulate computation benchmark
|
|
tokio::time::sleep(Duration::from_millis(5)).await;
|
|
Ok(100.0) // milliseconds
|
|
}
|
|
|
|
async fn benchmark_computation_with_communication(&self) -> Result<f64> {
|
|
// Simulate computation + communication benchmark
|
|
tokio::time::sleep(Duration::from_millis(8)).await;
|
|
Ok(104.0) // milliseconds (4% overhead)
|
|
}
|
|
|
|
async fn create_checkpoint(&mut self) -> Result<()> {
|
|
debug!("Creating checkpoint for fault tolerance");
|
|
self.fault_tolerance.last_checkpoint = Instant::now();
|
|
Ok(())
|
|
}
|
|
|
|
async fn redistribute_work(&mut self) -> Result<()> {
|
|
debug!("Redistributing work after GPU failure");
|
|
self.load_balancer
|
|
.rebalance_after_failure(&self.fault_tolerance.failed_gpus)
|
|
.await
|
|
}
|
|
|
|
async fn update_process_group_for_failure(&mut self, failed_rank: usize) -> Result<()> {
|
|
debug!(
|
|
"Updating process group after GPU failure on rank {}",
|
|
failed_rank
|
|
);
|
|
|
|
// Mark the failed rank as unavailable
|
|
self.mark_rank_failed(failed_rank).await?;
|
|
|
|
// Recalculate load distribution among remaining healthy ranks
|
|
self.rebalance_after_failure().await?;
|
|
|
|
// Update communication topology to exclude failed rank
|
|
self.update_communication_topology(failed_rank).await?;
|
|
|
|
info!(
|
|
"Process group reconfigured after failure on rank {}",
|
|
failed_rank
|
|
);
|
|
Ok(())
|
|
}
|
|
|
|
async fn mark_rank_failed(&mut self, failed_rank: usize) -> Result<()> {
|
|
info!(
|
|
"Marking rank {} as failed and removing from active pool",
|
|
failed_rank
|
|
);
|
|
// In a real implementation, this would update internal state
|
|
// to exclude the failed rank from future communications
|
|
Ok(())
|
|
}
|
|
|
|
async fn rebalance_after_failure(&mut self) -> Result<()> {
|
|
debug!("Rebalancing workload after GPU failure");
|
|
// Redistribute work among remaining healthy GPUs
|
|
// This would involve recalculating batch splits, gradient synchronization patterns, etc.
|
|
Ok(())
|
|
}
|
|
|
|
async fn update_communication_topology(&mut self, failed_rank: usize) -> Result<()> {
|
|
debug!(
|
|
"Updating communication topology to exclude rank {}",
|
|
failed_rank
|
|
);
|
|
// Update AllReduce trees, ring topologies, and other communication patterns
|
|
// to route around the failed GPU
|
|
Ok(())
|
|
}
|
|
}
|
|
|
|
impl LoadBalancer {
|
|
/// Create a new load balancer
|
|
fn new(num_gpus: usize) -> Self {
|
|
Self {
|
|
gpu_utilizations: vec![0.0; num_gpus],
|
|
memory_usage: vec![0; num_gpus],
|
|
imbalance_threshold: 0.1, // 10% max imbalance
|
|
distribution_weights: vec![1.0; num_gpus], // Equal weights initially
|
|
}
|
|
}
|
|
|
|
/// Update memory usage statistics for all GPUs
|
|
async fn update_memory_usage(&mut self) -> Result<()> {
|
|
// Simulate memory usage monitoring
|
|
for (i, usage) in self.memory_usage.iter_mut().enumerate() {
|
|
*usage = 1_000_000_000 + (i * 100_000_000) as u64; // 1GB + variance
|
|
}
|
|
Ok(())
|
|
}
|
|
|
|
/// Rebalance work distribution after GPU failures
|
|
async fn rebalance_after_failure(&mut self, failed_gpus: &[usize]) -> Result<()> {
|
|
debug!("Rebalancing work after GPU failures: {:?}", failed_gpus);
|
|
|
|
let healthy_gpu_count = self.distribution_weights.len() - failed_gpus.len();
|
|
if healthy_gpu_count == 0 {
|
|
return Err(DistributedError::fault_tolerance(
|
|
"No healthy GPUs remaining",
|
|
));
|
|
}
|
|
|
|
// Reset weights for failed GPUs and redistribute to healthy GPUs
|
|
let total_work = self.distribution_weights.iter().sum::<f64>();
|
|
let work_per_healthy_gpu = total_work / healthy_gpu_count as f64;
|
|
|
|
for (i, weight) in self.distribution_weights.iter_mut().enumerate() {
|
|
if failed_gpus.contains(&i) {
|
|
*weight = 0.0;
|
|
} else {
|
|
*weight = work_per_healthy_gpu;
|
|
}
|
|
}
|
|
|
|
debug!(
|
|
"New work distribution weights: {:?}",
|
|
self.distribution_weights
|
|
);
|
|
Ok(())
|
|
}
|
|
}
|
|
|
|
impl FaultTolerance {
|
|
/// Create a new fault tolerance manager
|
|
fn new(num_gpus: usize) -> Self {
|
|
Self {
|
|
health_monitors: vec![true; num_gpus],
|
|
recovery_timeout: Duration::from_secs(30),
|
|
failed_gpus: Vec::new(),
|
|
checkpoint_interval: Duration::from_secs(300), // 5 minutes
|
|
last_checkpoint: Instant::now(),
|
|
}
|
|
}
|
|
|
|
/// Check health status of all GPUs
|
|
async fn check_gpu_health(&mut self) -> Result<()> {
|
|
// Simulate GPU health monitoring
|
|
for (i, healthy) in self.health_monitors.iter().enumerate() {
|
|
if !healthy && !self.failed_gpus.contains(&i) {
|
|
warn!("Detected failed GPU: {}", i);
|
|
self.failed_gpus.push(i);
|
|
}
|
|
}
|
|
Ok(())
|
|
}
|
|
|
|
/// Mark a GPU as failed
|
|
|
|
fn mark_gpu_failed(&mut self, gpu_rank: usize) -> Result<()> {
|
|
if gpu_rank >= self.health_monitors.len() {
|
|
return Err(DistributedError::fault_tolerance(format!(
|
|
"Invalid GPU rank: {gpu_rank}"
|
|
)));
|
|
}
|
|
|
|
self.health_monitors[gpu_rank] = false;
|
|
if !self.failed_gpus.contains(&gpu_rank) {
|
|
self.failed_gpus.push(gpu_rank);
|
|
}
|
|
|
|
Ok(())
|
|
}
|
|
}
|
|
|
|
impl ScalingOptimizer {
|
|
/// Create a new scaling optimizer
|
|
fn new() -> Self {
|
|
Self {
|
|
communication_overlap_enabled: true,
|
|
communication_streams: 4,
|
|
compute_streams: 8,
|
|
ring_buffer_size: 1024 * 1024, // 1MB
|
|
gradient_compression: false,
|
|
}
|
|
}
|
|
}
|
|
|
|
impl DynamicLossScaler {
|
|
/// Create a new dynamic loss scaler
|
|
fn new() -> Self {
|
|
Self {
|
|
scale: 65536.0, // 2^16, typical starting scale for fp16
|
|
growth_factor: 2.0,
|
|
backoff_factor: 0.5,
|
|
steps_without_overflow: 0,
|
|
growth_interval: 2000, // Grow scale every 2000 steps without overflow
|
|
min_scale: 1.0,
|
|
max_scale: 65536.0 * 65536.0, // 2^32
|
|
precision_mode: "fp32".to_string(),
|
|
enabled: true,
|
|
}
|
|
}
|
|
|
|
/// Check if gradients have overflow (inf or nan values)
|
|
fn has_overflow(&self, gradients: &[Tensor]) -> Result<bool> {
|
|
for gradient in gradients {
|
|
let data = gradient.data()?;
|
|
for &value in &data {
|
|
if !value.is_finite() {
|
|
return Ok(true);
|
|
}
|
|
}
|
|
}
|
|
Ok(false)
|
|
}
|
|
|
|
/// Update the loss scale based on overflow detection
|
|
fn update_scale(&mut self, has_overflow: bool) {
|
|
if has_overflow {
|
|
// Reduce scale on overflow
|
|
self.scale *= self.backoff_factor;
|
|
self.scale = self.scale.max(self.min_scale);
|
|
self.steps_without_overflow = 0;
|
|
debug!("Loss scale reduced to {} due to overflow", self.scale);
|
|
} else {
|
|
// Increase steps without overflow
|
|
self.steps_without_overflow += 1;
|
|
|
|
// Grow scale if no overflow for growth_interval steps
|
|
if self.steps_without_overflow >= self.growth_interval {
|
|
self.scale *= self.growth_factor;
|
|
self.scale = self.scale.min(self.max_scale);
|
|
self.steps_without_overflow = 0;
|
|
debug!(
|
|
"Loss scale increased to {} after {} stable steps",
|
|
self.scale, self.growth_interval
|
|
);
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Scale gradients for mixed precision training
|
|
fn scale_gradients(&self, gradients: &mut [Tensor]) -> Result<()> {
|
|
if !self.enabled || self.scale == 1.0 {
|
|
return Ok(());
|
|
}
|
|
|
|
for gradient in gradients.iter_mut() {
|
|
*gradient = gradient.mul_scalar(self.scale)?;
|
|
}
|
|
Ok(())
|
|
}
|
|
|
|
/// Unscale gradients after all-reduce
|
|
fn unscale_gradients(&self, gradients: &mut [Tensor]) -> Result<()> {
|
|
if !self.enabled || self.scale == 1.0 {
|
|
return Ok(());
|
|
}
|
|
|
|
let inv_scale = 1.0 / self.scale;
|
|
for gradient in gradients.iter_mut() {
|
|
*gradient = gradient.mul_scalar(inv_scale)?;
|
|
}
|
|
Ok(())
|
|
}
|
|
}
|