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
415 lines
14 KiB
Rust
415 lines
14 KiB
Rust
//! Tests for ring-based AllReduce implementation
|
|
//!
|
|
//! These tests verify the ring-based AllReduce algorithm implementation
|
|
//! for efficient gradient synchronization across multiple GPUs.
|
|
|
|
use rtx_distributed::{Backend, MultiGpuTrainer, ProcessGroup, ReduceOp, Result, WorldInfo};
|
|
use rtx_tensor::{Device, Tensor};
|
|
use std::time::Duration;
|
|
|
|
/// Test ring AllReduce basic functionality
|
|
#[tokio::test]
|
|
async fn test_ring_allreduce_basic() -> Result<()> {
|
|
let world_size = 4;
|
|
let rank = 0;
|
|
|
|
let mut trainer = MultiGpuTrainer::new(world_size, rank).await?;
|
|
|
|
// Enable communication overlap to use ring AllReduce
|
|
trainer.scaling_optimizer.communication_overlap_enabled = true;
|
|
|
|
// Create test gradients with known values
|
|
let shape_dims = vec![1000];
|
|
let gradient1 = Tensor::full(
|
|
&shape_dims,
|
|
2.0,
|
|
&Device::cuda(0).unwrap_or(Device::default()),
|
|
)?;
|
|
let gradient2 = Tensor::full(
|
|
&shape_dims,
|
|
4.0,
|
|
&Device::cuda(0).unwrap_or(Device::default()),
|
|
)?;
|
|
let mut gradients = vec![gradient1, gradient2];
|
|
|
|
// Perform ring AllReduce synchronization
|
|
trainer.synchronize_gradients(&mut gradients).await?;
|
|
|
|
// Verify the gradients were processed correctly
|
|
let grad1_data = gradients[0].data()?;
|
|
let grad2_data = gradients[1].data()?;
|
|
|
|
println!(
|
|
"Grad1 first value: {}, last value: {}",
|
|
grad1_data[0],
|
|
grad1_data[grad1_data.len() - 1]
|
|
);
|
|
println!(
|
|
"Grad2 first value: {}, last value: {}",
|
|
grad2_data[0],
|
|
grad2_data[grad2_data.len() - 1]
|
|
);
|
|
|
|
// Check that all values within each tensor are consistent
|
|
for (i, &value) in grad1_data.iter().enumerate() {
|
|
if (value - grad1_data[0]).abs() >= 1e-6 {
|
|
println!(
|
|
"Inconsistent value at index {}: {} vs {}",
|
|
i, value, grad1_data[0]
|
|
);
|
|
}
|
|
assert!(
|
|
(value - grad1_data[0]).abs() < 1e-6,
|
|
"Ring AllReduce should produce consistent values"
|
|
);
|
|
}
|
|
|
|
for (i, &value) in grad2_data.iter().enumerate() {
|
|
if (value - grad2_data[0]).abs() >= 1e-6 {
|
|
println!(
|
|
"Inconsistent value at index {}: {} vs {}",
|
|
i, value, grad2_data[0]
|
|
);
|
|
}
|
|
assert!(
|
|
(value - grad2_data[0]).abs() < 1e-6,
|
|
"Ring AllReduce should produce consistent values"
|
|
);
|
|
}
|
|
|
|
Ok(())
|
|
}
|
|
|
|
/// Test ring AllReduce with different buffer sizes
|
|
#[tokio::test]
|
|
async fn test_ring_allreduce_buffer_sizes() -> Result<()> {
|
|
let world_size = 4;
|
|
let rank = 0;
|
|
|
|
let mut trainer = MultiGpuTrainer::new(world_size, rank).await?;
|
|
trainer.scaling_optimizer.communication_overlap_enabled = true;
|
|
|
|
// Test different buffer sizes
|
|
let buffer_sizes = vec![512, 1024, 2048, 4096];
|
|
|
|
for buffer_size in buffer_sizes {
|
|
trainer.scaling_optimizer.ring_buffer_size = buffer_size;
|
|
|
|
// Create gradients that will be chunked based on buffer size
|
|
let shape_dims = vec![buffer_size * 2]; // Larger than buffer to test chunking
|
|
let gradient = Tensor::full(
|
|
&shape_dims,
|
|
1.0,
|
|
&Device::cuda(0).unwrap_or(Device::default()),
|
|
)?;
|
|
let mut gradients = vec![gradient];
|
|
|
|
let start_time = std::time::Instant::now();
|
|
trainer.synchronize_gradients(&mut gradients).await?;
|
|
let duration = start_time.elapsed();
|
|
|
|
println!(
|
|
"Ring AllReduce with buffer size {}: {:.2}ms",
|
|
buffer_size,
|
|
duration.as_millis()
|
|
);
|
|
|
|
// Verify correctness
|
|
let data = gradients[0].data()?;
|
|
for &value in &data {
|
|
assert!(
|
|
(value - data[0]).abs() < 1e-6,
|
|
"Values should be consistent after ring AllReduce"
|
|
);
|
|
}
|
|
}
|
|
|
|
Ok(())
|
|
}
|
|
|
|
/// Test ring AllReduce communication overlap
|
|
#[tokio::test]
|
|
async fn test_ring_allreduce_overlap() -> Result<()> {
|
|
let world_size = 8;
|
|
let rank = 0;
|
|
|
|
let mut trainer = MultiGpuTrainer::new(world_size, rank).await?;
|
|
|
|
// Create large gradients to test overlap effectiveness
|
|
let shape_dims = vec![100_000];
|
|
let gradient = Tensor::ones(&shape_dims, &Device::cuda(0).unwrap_or(Device::default()))?;
|
|
let mut gradients = vec![gradient.clone(), gradient.clone(), gradient];
|
|
|
|
// Test without overlap
|
|
trainer.scaling_optimizer.communication_overlap_enabled = false;
|
|
let start_no_overlap = std::time::Instant::now();
|
|
trainer.synchronize_gradients(&mut gradients).await?;
|
|
let time_no_overlap = start_no_overlap.elapsed();
|
|
|
|
// Reset gradients
|
|
let gradient = Tensor::ones(&shape_dims, &Device::cuda(0).unwrap_or(Device::default()))?;
|
|
gradients = vec![gradient.clone(), gradient.clone(), gradient];
|
|
|
|
// Test with overlap
|
|
trainer.scaling_optimizer.communication_overlap_enabled = true;
|
|
let start_with_overlap = std::time::Instant::now();
|
|
trainer.synchronize_gradients(&mut gradients).await?;
|
|
let time_with_overlap = start_with_overlap.elapsed();
|
|
|
|
println!(
|
|
"Without overlap: {:.2}ms, With overlap: {:.2}ms",
|
|
time_no_overlap.as_millis(),
|
|
time_with_overlap.as_millis()
|
|
);
|
|
|
|
// With proper implementation, overlap should be faster or similar
|
|
// For simulation, we just verify both complete successfully
|
|
assert!(
|
|
time_no_overlap.as_millis() < 5000,
|
|
"Sync should complete in reasonable time"
|
|
);
|
|
assert!(
|
|
time_with_overlap.as_millis() < 5000,
|
|
"Overlapped sync should complete in reasonable time"
|
|
);
|
|
|
|
Ok(())
|
|
}
|
|
|
|
/// Test ring AllReduce with gradient compression
|
|
#[tokio::test]
|
|
async fn test_ring_allreduce_compression() -> Result<()> {
|
|
let world_size = 4;
|
|
let rank = 0;
|
|
|
|
let mut trainer = MultiGpuTrainer::new(world_size, rank).await?;
|
|
trainer.scaling_optimizer.communication_overlap_enabled = true;
|
|
|
|
// Create gradients for compression testing
|
|
let shape_dims = vec![10_000];
|
|
let gradient = Tensor::full(
|
|
&shape_dims,
|
|
3.14159,
|
|
&Device::cuda(0).unwrap_or(Device::default()),
|
|
)?;
|
|
let mut gradients = vec![gradient];
|
|
|
|
// Test without compression
|
|
trainer.scaling_optimizer.gradient_compression = false;
|
|
let start_uncompressed = std::time::Instant::now();
|
|
trainer.synchronize_gradients(&mut gradients).await?;
|
|
let time_uncompressed = start_uncompressed.elapsed();
|
|
|
|
// Reset gradient
|
|
let gradient = Tensor::full(
|
|
&shape_dims,
|
|
3.14159,
|
|
&Device::cuda(0).unwrap_or(Device::default()),
|
|
)?;
|
|
gradients = vec![gradient];
|
|
|
|
// Test with compression
|
|
trainer.scaling_optimizer.gradient_compression = true;
|
|
let start_compressed = std::time::Instant::now();
|
|
trainer.synchronize_gradients(&mut gradients).await?;
|
|
let time_compressed = start_compressed.elapsed();
|
|
|
|
println!(
|
|
"Uncompressed: {:.2}ms, Compressed: {:.2}ms",
|
|
time_uncompressed.as_millis(),
|
|
time_compressed.as_millis()
|
|
);
|
|
|
|
// Verify correctness is maintained with compression
|
|
let data = gradients[0].data()?;
|
|
for &value in &data {
|
|
// With compression, there might be small precision loss
|
|
assert!(
|
|
(value - data[0]).abs() < 1e-4,
|
|
"Compression should maintain reasonable precision"
|
|
);
|
|
}
|
|
|
|
Ok(())
|
|
}
|
|
|
|
/// Test ring AllReduce error handling and recovery
|
|
#[tokio::test]
|
|
async fn test_ring_allreduce_error_handling() -> Result<()> {
|
|
let world_size = 4;
|
|
let rank = 0;
|
|
|
|
let mut trainer = MultiGpuTrainer::new(world_size, rank).await?;
|
|
trainer.scaling_optimizer.communication_overlap_enabled = true;
|
|
|
|
// Test with empty gradients
|
|
let mut empty_gradients = Vec::new();
|
|
let result = trainer.synchronize_gradients(&mut empty_gradients).await;
|
|
assert!(
|
|
result.is_ok(),
|
|
"Empty gradients should be handled gracefully"
|
|
);
|
|
|
|
// Test with very large gradients (potential memory issues)
|
|
let large_shape = vec![10_000_000]; // 10M elements
|
|
let large_gradient_result =
|
|
Tensor::ones(&large_shape, &Device::cuda(0).unwrap_or(Device::default()));
|
|
|
|
if let Ok(large_gradient) = large_gradient_result {
|
|
let mut large_gradients = vec![large_gradient];
|
|
let result = trainer.synchronize_gradients(&mut large_gradients).await;
|
|
// Should either succeed or fail gracefully
|
|
match result {
|
|
Ok(_) => println!("Large gradient sync succeeded"),
|
|
Err(e) => println!("Large gradient sync failed gracefully: {}", e),
|
|
}
|
|
}
|
|
|
|
Ok(())
|
|
}
|
|
|
|
/// Test ring AllReduce with multiple communication streams
|
|
#[tokio::test]
|
|
async fn test_ring_allreduce_multi_stream() -> Result<()> {
|
|
let world_size = 4;
|
|
let rank = 0;
|
|
|
|
let mut trainer = MultiGpuTrainer::new(world_size, rank).await?;
|
|
trainer.scaling_optimizer.communication_overlap_enabled = true;
|
|
|
|
// Test different numbers of communication streams
|
|
let stream_counts = vec![1, 2, 4, 8];
|
|
|
|
for stream_count in stream_counts {
|
|
trainer.scaling_optimizer.communication_streams = stream_count;
|
|
trainer.scaling_optimizer.compute_streams = stream_count * 2;
|
|
|
|
// Create multiple gradients to utilize multiple streams
|
|
let shape_dims = vec![5000];
|
|
let gradients_count = stream_count;
|
|
let mut gradients = Vec::new();
|
|
|
|
for i in 0..gradients_count {
|
|
let gradient = Tensor::full(
|
|
&shape_dims,
|
|
(i + 1) as f32,
|
|
&Device::cuda(0).unwrap_or(Device::default()),
|
|
)?;
|
|
gradients.push(gradient);
|
|
}
|
|
|
|
let start_time = std::time::Instant::now();
|
|
trainer.synchronize_gradients(&mut gradients).await?;
|
|
let duration = start_time.elapsed();
|
|
|
|
println!(
|
|
"Ring AllReduce with {} streams: {:.2}ms",
|
|
stream_count,
|
|
duration.as_millis()
|
|
);
|
|
|
|
// Verify all gradients were processed correctly
|
|
for (i, gradient) in gradients.iter().enumerate() {
|
|
let data = gradient.data()?;
|
|
for &value in &data {
|
|
assert!(
|
|
(value - data[0]).abs() < 1e-6,
|
|
"Gradient {} values should be consistent",
|
|
i
|
|
);
|
|
}
|
|
}
|
|
}
|
|
|
|
Ok(())
|
|
}
|
|
|
|
/// Benchmark ring AllReduce performance vs sequential AllReduce
|
|
#[tokio::test]
|
|
async fn benchmark_ring_vs_sequential_allreduce() -> Result<()> {
|
|
let world_size = 8;
|
|
let rank = 0;
|
|
|
|
// Test different gradient sizes
|
|
let sizes = vec![1_000, 10_000, 100_000];
|
|
|
|
for size in sizes {
|
|
let mut trainer = MultiGpuTrainer::new(world_size, rank).await?;
|
|
let shape_dims = vec![size];
|
|
|
|
// Benchmark sequential AllReduce
|
|
trainer.scaling_optimizer.communication_overlap_enabled = false;
|
|
let gradient = Tensor::ones(&shape_dims, &Device::cuda(0).unwrap_or(Device::default()))?;
|
|
let mut gradients = vec![gradient];
|
|
|
|
let start_sequential = std::time::Instant::now();
|
|
trainer.synchronize_gradients(&mut gradients).await?;
|
|
let time_sequential = start_sequential.elapsed();
|
|
|
|
// Benchmark ring AllReduce
|
|
trainer.scaling_optimizer.communication_overlap_enabled = true;
|
|
let gradient = Tensor::ones(&shape_dims, &Device::cuda(0).unwrap_or(Device::default()))?;
|
|
let mut gradients = vec![gradient];
|
|
|
|
let start_ring = std::time::Instant::now();
|
|
trainer.synchronize_gradients(&mut gradients).await?;
|
|
let time_ring = start_ring.elapsed();
|
|
|
|
let speedup = time_sequential.as_micros() as f64 / time_ring.as_micros() as f64;
|
|
|
|
println!("Size: {} elements", size);
|
|
println!(" Sequential: {:.2}ms", time_sequential.as_millis());
|
|
println!(" Ring: {:.2}ms", time_ring.as_millis());
|
|
println!(" Speedup: {:.2}x", speedup);
|
|
|
|
// Both should complete in reasonable time
|
|
assert!(
|
|
time_sequential.as_millis() < 1000,
|
|
"Sequential should be reasonably fast"
|
|
);
|
|
assert!(
|
|
time_ring.as_millis() < 1000,
|
|
"Ring should be reasonably fast"
|
|
);
|
|
}
|
|
|
|
Ok(())
|
|
}
|
|
|
|
/// Verify that the overlapped (ring AllReduce) gradient synchronization path
|
|
/// produces the same result as the sequential path.
|
|
///
|
|
/// Previously, ring_allreduce() contained its own simulation that bypassed the
|
|
/// ProcessGroup dispatch, meaning it would NOT use NCCL/RNCCL even when compiled in.
|
|
/// This test ensures both paths delegate to the same ProcessGroup allreduce logic.
|
|
#[tokio::test]
|
|
async fn test_ring_allreduce_matches_sequential_path() -> Result<()> {
|
|
let world_size = 4;
|
|
let input_value = 3.0_f32;
|
|
let shape_dims = vec![64usize];
|
|
let device = Device::cuda(0).unwrap_or(Device::default());
|
|
|
|
// Sequential path
|
|
let mut trainer_seq = MultiGpuTrainer::new(world_size, 0).await?;
|
|
trainer_seq.scaling_optimizer.communication_overlap_enabled = false;
|
|
let mut grad_seq = vec![Tensor::full(&shape_dims, input_value, &device)?];
|
|
trainer_seq.synchronize_gradients(&mut grad_seq).await?;
|
|
let seq_data = grad_seq[0].data()?;
|
|
|
|
// Ring AllReduce (overlapped) path
|
|
let mut trainer_ring = MultiGpuTrainer::new(world_size, 0).await?;
|
|
trainer_ring.scaling_optimizer.communication_overlap_enabled = true;
|
|
let mut grad_ring = vec![Tensor::full(&shape_dims, input_value, &device)?];
|
|
trainer_ring.synchronize_gradients(&mut grad_ring).await?;
|
|
let ring_data = grad_ring[0].data()?;
|
|
|
|
assert_eq!(seq_data.len(), ring_data.len(), "output lengths must match");
|
|
for (i, (&s, &r)) in seq_data.iter().zip(ring_data.iter()).enumerate() {
|
|
assert!(
|
|
(s - r).abs() < 1e-5,
|
|
"element {i}: sequential={s}, ring={r} -- overlapped path must match sequential"
|
|
);
|
|
}
|
|
Ok(())
|
|
}
|