345 lines
12 KiB
Rust
345 lines
12 KiB
Rust
//! Comprehensive tests for NCCL backend implementation
|
|
//!
|
|
//! This module contains tests that validate the NCCL backend functionality,
|
|
//! including communicator management, collective operations, and error handling.
|
|
#![cfg(feature = "disabled_tests")]
|
|
|
|
use crate::backend::{Backend, BackendConfig};
|
|
use crate::comm::ReduceOp;
|
|
use crate::error::Result;
|
|
use crate::group::{ProcessGroup, WorldInfo};
|
|
use crate::nccl::{NcclBackend, NcclConfig};
|
|
use crate::{Device, Tensor, TensorShape};
|
|
use proptest::prelude::*;
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
|
|
/// Test NCCL configuration creation and validation
|
|
#[test]
|
|
fn test_nccl_config() {
|
|
let config = NcclConfig::default();
|
|
assert_eq!(config.socket_ifname, "^lo");
|
|
assert_eq!(config.debug_level, "INFO");
|
|
assert_eq!(config.tree_threshold, 0);
|
|
assert_eq!(config.p2p_threshold, 8192);
|
|
assert_eq!(config.buffer_size, 1024 * 1024);
|
|
}
|
|
|
|
/// Test NCCL backend configuration creation
|
|
#[test]
|
|
fn test_nccl_backend_config() {
|
|
let config = BackendConfig::nccl();
|
|
assert_eq!(config.backend, Backend::Nccl);
|
|
assert!(config.nccl_config().is_some());
|
|
|
|
let nccl_config = config.nccl_config().unwrap();
|
|
assert_eq!(nccl_config.socket_ifname, "^lo");
|
|
}
|
|
|
|
/// Test NCCL backend creation
|
|
#[test]
|
|
fn test_nccl_backend_creation() {
|
|
let nccl_config = NcclConfig::default();
|
|
let backend = NcclBackend::new(nccl_config);
|
|
|
|
assert!(backend.world_comm().read().is_none());
|
|
}
|
|
|
|
/// Test process group creation with NCCL backend (CPU fallback mode)
|
|
#[tokio::test]
|
|
async fn test_process_group_nccl_fallback() {
|
|
let world_info = WorldInfo::new(2, 0, Backend::Nccl);
|
|
|
|
// This should work even without CUDA as it will fall back to CPU simulation
|
|
let result = ProcessGroup::new(Backend::Cpu, world_info);
|
|
|
|
match result {
|
|
Ok(pg) => {
|
|
assert_eq!(pg.world_size(), 2);
|
|
assert_eq!(pg.rank(), 0);
|
|
|
|
// Cleanup
|
|
let _ = pg.cleanup().await;
|
|
}
|
|
Err(_) => {
|
|
// Expected if CUDA/NCCL is not available
|
|
println!("NCCL not available, test skipped");
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Test tensor operations with various data types and sizes
|
|
#[tokio::test]
|
|
async fn test_tensor_operations() {
|
|
let world_info = WorldInfo::new(1, 0, Backend::Cpu);
|
|
let pg = ProcessGroup::new(Backend::Cpu, world_info).unwrap();
|
|
|
|
// Test various tensor shapes
|
|
let test_shapes = vec![
|
|
vec![1],
|
|
vec![10],
|
|
vec![128],
|
|
vec![1024],
|
|
vec![2, 3],
|
|
vec![4, 4, 4],
|
|
vec![8, 8, 8, 8],
|
|
];
|
|
|
|
for shape_dims in test_shapes {
|
|
let tensor = Tensor::ones(&shape_dims, &crate::Device::default()).unwrap();
|
|
let mut test_tensor = tensor.clone();
|
|
|
|
// Test AllReduce
|
|
let result = pg.all_reduce(&mut test_tensor, ReduceOp::Sum).await;
|
|
assert!(
|
|
result.is_ok(),
|
|
"AllReduce failed for shape {:?}",
|
|
shape_dims
|
|
);
|
|
|
|
// Test Broadcast
|
|
let mut broadcast_tensor = tensor.clone();
|
|
let result = pg.broadcast(&mut broadcast_tensor, 0).await;
|
|
assert!(
|
|
result.is_ok(),
|
|
"Broadcast failed for shape {:?}",
|
|
shape_dims
|
|
);
|
|
}
|
|
|
|
pg.cleanup().await.unwrap();
|
|
}
|
|
|
|
/// Test error conditions and edge cases
|
|
#[tokio::test]
|
|
async fn test_error_conditions() {
|
|
let world_info = WorldInfo::new(2, 0, Backend::Cpu);
|
|
let pg = ProcessGroup::new(Backend::Cpu, world_info).unwrap();
|
|
|
|
// Test with empty tensor
|
|
let empty_shape = TensorShape::new(vec![0]).unwrap();
|
|
let empty_tensor = Tensor::zeros(empty_shape, &crate::Device::default()).unwrap();
|
|
let mut test_tensor = empty_tensor.clone();
|
|
|
|
let result = pg.all_reduce(&mut test_tensor, ReduceOp::Sum).await;
|
|
// This might succeed or fail depending on implementation - both are acceptable
|
|
|
|
// Test broadcast with invalid root
|
|
let valid_tensor = Tensor::ones(&vec![3], &crate::Device::default()).unwrap();
|
|
let mut broadcast_tensor = valid_tensor.clone();
|
|
let result = pg.broadcast(&mut broadcast_tensor, 99).await; // Invalid root
|
|
assert!(result.is_err(), "Broadcast should fail with invalid root");
|
|
|
|
pg.cleanup().await.unwrap();
|
|
}
|
|
|
|
/// Test process group cleanup
|
|
#[tokio::test]
|
|
async fn test_process_group_cleanup() {
|
|
let world_info = WorldInfo::new(1, 0, Backend::Cpu);
|
|
let pg = ProcessGroup::new(Backend::Cpu, world_info).unwrap();
|
|
|
|
// Test cleanup
|
|
assert!(pg.cleanup().await.is_ok());
|
|
|
|
// Test double cleanup (should not error)
|
|
assert!(pg.cleanup().await.is_ok());
|
|
}
|
|
|
|
/// Test concurrent operations
|
|
#[tokio::test]
|
|
async fn test_concurrent_operations() {
|
|
let world_info = WorldInfo::new(1, 0, Backend::Cpu);
|
|
let pg = ProcessGroup::new(Backend::Cpu, world_info).unwrap();
|
|
|
|
// Create multiple tasks performing operations concurrently
|
|
let mut tasks = Vec::new();
|
|
|
|
for i in 0..5 {
|
|
let pg_clone = pg.clone();
|
|
let task = tokio::spawn(async move {
|
|
let tensor = Tensor::ones(&vec![10], &crate::Device::default()).unwrap();
|
|
let mut test_tensor = tensor.clone();
|
|
|
|
// Perform operation
|
|
let result = pg_clone.all_reduce(&mut test_tensor, ReduceOp::Sum).await;
|
|
assert!(result.is_ok(), "Concurrent operation {} failed", i);
|
|
});
|
|
tasks.push(task);
|
|
}
|
|
|
|
// Wait for all tasks to complete
|
|
for task in tasks {
|
|
task.await.unwrap();
|
|
}
|
|
|
|
pg.cleanup().await.unwrap();
|
|
}
|
|
|
|
/// Property-based test for tensor operations
|
|
proptest! {
|
|
#[test]
|
|
fn prop_test_tensor_shapes(
|
|
dims in prop::collection::vec(1usize..100, 1..4),
|
|
) {
|
|
let rt = tokio::runtime::Runtime::new().unwrap();
|
|
rt.block_on(async {
|
|
let config = BackendConfig::cpu();
|
|
let pg = ProcessGroup::new_with_config(Backend::Cpu, 1, 0, config).await.unwrap();
|
|
|
|
let tensor = Tensor::ones(&dims, &crate::Device::default()).unwrap();
|
|
let mut test_tensor = tensor.clone();
|
|
|
|
// Test that AllReduce works with any reasonable tensor shape
|
|
let result = pg.all_reduce(&mut test_tensor, ReduceOp::Sum).await;
|
|
prop_assert!(result.is_ok(), "AllReduce failed for dims {:?}", dims);
|
|
|
|
pg.cleanup().await.unwrap();
|
|
Ok(())
|
|
});
|
|
}
|
|
|
|
#[test]
|
|
fn prop_test_reduce_operations(
|
|
op in prop_oneof![
|
|
Just(ReduceOp::Sum),
|
|
Just(ReduceOp::Max),
|
|
Just(ReduceOp::Min),
|
|
Just(ReduceOp::Product),
|
|
],
|
|
) {
|
|
let rt = tokio::runtime::Runtime::new().unwrap();
|
|
rt.block_on(async {
|
|
let config = BackendConfig::cpu();
|
|
let pg = ProcessGroup::new_with_config(Backend::Cpu, 1, 0, config).await.unwrap();
|
|
|
|
let tensor = Tensor::ones(&vec![5], &crate::Device::default()).unwrap();
|
|
let mut test_tensor = tensor.clone();
|
|
|
|
// Test that different reduce operations work
|
|
let result = pg.all_reduce(&mut test_tensor, op).await;
|
|
prop_assert!(result.is_ok(), "AllReduce failed for op {:?}", op);
|
|
|
|
pg.cleanup().await.unwrap();
|
|
Ok(())
|
|
});
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Integration tests that require actual CUDA/NCCL environment
|
|
#[cfg(test)]
|
|
mod integration_tests {
|
|
use super::*;
|
|
|
|
/// Test that requires actual NCCL initialization
|
|
/// This test is only run when CUDA is available
|
|
#[tokio::test]
|
|
#[ignore = "requires CUDA environment"]
|
|
async fn test_real_nccl_initialization() {
|
|
// This test would only pass in an environment with CUDA and NCCL
|
|
let config = BackendConfig::nccl();
|
|
let world_info = WorldInfo::new(1, 0, Backend::Nccl);
|
|
|
|
let result = ProcessGroup::new(Backend::Nccl, world_info);
|
|
|
|
match result {
|
|
Ok(pg) => {
|
|
// Verify NCCL backend is actually being used
|
|
#[cfg(feature = "nccl")]
|
|
{
|
|
assert!(pg.has_nccl_backend());
|
|
assert!(pg.communication_primitive().unwrap().is_some());
|
|
}
|
|
|
|
// Test actual NCCL operations
|
|
let tensor = Tensor::ones(&vec![1000], &crate::Device::Cuda(0)).unwrap();
|
|
let mut test_tensor = tensor.clone();
|
|
|
|
let result = pg.all_reduce(&mut test_tensor, ReduceOp::Sum).await;
|
|
assert!(result.is_ok(), "NCCL AllReduce failed");
|
|
|
|
pg.cleanup().await.unwrap();
|
|
}
|
|
Err(e) => {
|
|
println!(
|
|
"NCCL initialization failed (expected in non-CUDA environment): {}",
|
|
e
|
|
);
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Test multi-GPU operations
|
|
#[tokio::test]
|
|
#[ignore = "requires multi-GPU CUDA environment"]
|
|
async fn test_multi_gpu_operations() {
|
|
// This would test actual multi-GPU NCCL operations
|
|
// Only runs in environments with multiple GPUs
|
|
|
|
for device_id in 0..2 {
|
|
// Test with 2 GPUs
|
|
let world_info = WorldInfo::new(2, device_id, Backend::Nccl);
|
|
let pg = ProcessGroup::new(Backend::Nccl, world_info);
|
|
|
|
if let Ok(pg) = pg {
|
|
let device = crate::Device::Cuda(device_id as usize);
|
|
let tensor = Tensor::ones(&vec![1024], &device).unwrap();
|
|
let mut test_tensor = tensor.clone();
|
|
|
|
// Test cross-GPU AllReduce
|
|
let result = pg.all_reduce(&mut test_tensor, ReduceOp::Sum).await;
|
|
assert!(
|
|
result.is_ok(),
|
|
"Multi-GPU AllReduce failed for device {}",
|
|
device_id
|
|
);
|
|
|
|
pg.cleanup().await.unwrap();
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Benchmark NCCL operations vs custom implementations
|
|
#[tokio::test]
|
|
#[ignore = "performance benchmark"]
|
|
async fn benchmark_nccl_vs_custom() {
|
|
use std::time::Instant;
|
|
|
|
let tensor_sizes = vec![1024, 10240, 102400, 1024000];
|
|
|
|
for size in tensor_sizes {
|
|
println!("Benchmarking tensor size: {}", size);
|
|
|
|
// Test NCCL backend (if available)
|
|
let world_info = WorldInfo::new(1, 0, Backend::Nccl);
|
|
if let Ok(nccl_pg) = ProcessGroup::new(Backend::Nccl, world_info) {
|
|
let tensor = Tensor::ones(&vec![size], &crate::Device::default()).unwrap();
|
|
let mut nccl_tensor = tensor.clone();
|
|
|
|
let start = Instant::now();
|
|
let _ = nccl_pg.all_reduce(&mut nccl_tensor, ReduceOp::Sum).await;
|
|
let nccl_duration = start.elapsed();
|
|
|
|
println!("NCCL AllReduce time: {:?}", nccl_duration);
|
|
nccl_pg.cleanup().await.unwrap();
|
|
}
|
|
|
|
// Test CPU backend for comparison
|
|
let world_info = WorldInfo::new(1, 0, Backend::Cpu);
|
|
let cpu_pg = ProcessGroup::new(Backend::Cpu, world_info).unwrap();
|
|
let tensor = Tensor::ones(&vec![size], &crate::Device::default()).unwrap();
|
|
let mut cpu_tensor = tensor.clone();
|
|
|
|
let start = Instant::now();
|
|
let _ = cpu_pg.all_reduce(&mut cpu_tensor, ReduceOp::Sum).await;
|
|
let cpu_duration = start.elapsed();
|
|
|
|
println!("CPU AllReduce time: {:?}", cpu_duration);
|
|
cpu_pg.cleanup().await.unwrap();
|
|
}
|
|
}
|
|
}
|