267 lines
8.1 KiB
Rust
267 lines
8.1 KiB
Rust
//! Example demonstrating NCCL backend usage for distributed training
|
|
//!
|
|
//! This example shows how to:
|
|
//! 1. Initialize a process group with NCCL backend
|
|
//! 2. Perform collective communication operations
|
|
//! 3. Handle different tensor operations
|
|
//! 4. Proper cleanup and error handling
|
|
//!
|
|
//! Usage:
|
|
//! ```bash
|
|
//! # Single GPU example
|
|
//! cargo run --example nccl_example --features nccl
|
|
//!
|
|
//! # Multi-GPU example (requires multiple GPUs)
|
|
//! CUDA_VISIBLE_DEVICES=0,1 mpirun -n 2 cargo run --example nccl_example --features nccl
|
|
//! ```
|
|
|
|
use anyhow::Result;
|
|
use rtx_distributed::{
|
|
Backend, BackendConfig, Device, ProcessGroup, Tensor, TensorShape, WorldInfo,
|
|
comm::{CommunicationPrimitive, ReduceOp},
|
|
};
|
|
use std::time::Instant;
|
|
use tracing::{info, warn};
|
|
|
|
#[tokio::main]
|
|
async fn main() -> Result<()> {
|
|
// Initialize tracing
|
|
tracing_subscriber::fmt()
|
|
.with_max_level(tracing::Level::INFO)
|
|
.init();
|
|
|
|
info!("Starting NCCL backend example");
|
|
|
|
// Check if NCCL feature is enabled
|
|
#[cfg(not(feature = "nccl"))]
|
|
{
|
|
warn!("NCCL feature not enabled. Compile with --features nccl");
|
|
return Ok(());
|
|
}
|
|
|
|
#[cfg(feature = "nccl")]
|
|
{
|
|
// Run the NCCL example
|
|
run_nccl_example().await?;
|
|
}
|
|
|
|
Ok(())
|
|
}
|
|
|
|
#[cfg(feature = "nccl")]
|
|
async fn run_nccl_example() -> Result<()> {
|
|
// Configuration
|
|
let world_size = 1; // Single process example
|
|
let rank = 0;
|
|
|
|
info!("Initializing NCCL backend for rank {}/{}", rank, world_size);
|
|
|
|
// Try NCCL backend first, fallback to CPU if CUDA not available
|
|
let (pg, actual_backend) = match initialize_process_group(world_size, rank).await {
|
|
Ok(pg) => {
|
|
info!("Successfully initialized NCCL backend");
|
|
(pg, "NCCL")
|
|
}
|
|
Err(e) => {
|
|
warn!(
|
|
"NCCL initialization failed: {}. Falling back to CPU backend",
|
|
e
|
|
);
|
|
let config = BackendConfig::cpu();
|
|
let pg =
|
|
ProcessGroup::new_with_config(Backend::Cpu, world_size as i32, rank as i32, config)
|
|
.await?;
|
|
(pg, "CPU")
|
|
}
|
|
};
|
|
|
|
info!("Using {} backend", actual_backend);
|
|
|
|
// Example 1: AllReduce operation
|
|
info!("=== Example 1: AllReduce Operation ===");
|
|
await_allreduce_example(&pg).await?;
|
|
|
|
// Example 2: Broadcast operation
|
|
info!("=== Example 2: Broadcast Operation ===");
|
|
await_broadcast_example(&pg).await?;
|
|
|
|
// Example 3: Performance comparison
|
|
info!("=== Example 3: Performance Comparison ===");
|
|
await_performance_comparison(&pg, actual_backend).await?;
|
|
|
|
// Example 4: Different tensor shapes and data types
|
|
info!("=== Example 4: Various Tensor Operations ===");
|
|
await_tensor_operations_example(&pg).await?;
|
|
|
|
// Example 5: Error handling
|
|
info!("=== Example 5: Error Handling ===");
|
|
await_error_handling_example(&pg).await?;
|
|
|
|
// Cleanup
|
|
info!("Cleaning up process group");
|
|
pg.cleanup().await?;
|
|
|
|
info!("NCCL backend example completed successfully");
|
|
Ok(())
|
|
}
|
|
|
|
#[cfg(feature = "nccl")]
|
|
async fn initialize_process_group(world_size: usize, rank: usize) -> Result<ProcessGroup> {
|
|
// Create NCCL backend configuration
|
|
let mut config = BackendConfig::nccl();
|
|
|
|
// Configure NCCL parameters for optimal performance
|
|
config.set_parameter("nccl_socket_ifname".to_string(), "^lo".to_string());
|
|
config.set_parameter("nccl_debug".to_string(), "INFO".to_string());
|
|
config.set_timeout(std::time::Duration::from_secs(60));
|
|
|
|
// Initialize process group
|
|
Ok(
|
|
ProcessGroup::new_with_config(Backend::Nccl, world_size as i32, rank as i32, config)
|
|
.await?,
|
|
)
|
|
}
|
|
|
|
#[cfg(feature = "nccl")]
|
|
async fn await_allreduce_example(pg: &ProcessGroup) -> Result<()> {
|
|
// Create test tensor
|
|
let shape = vec![1024];
|
|
let mut tensor = Tensor::ones(&shape, &Device::default())?;
|
|
|
|
info!(
|
|
"Original tensor sum: {:.2}",
|
|
tensor.data()?.iter().sum::<f32>()
|
|
);
|
|
|
|
// Perform AllReduce sum
|
|
let start = Instant::now();
|
|
pg.all_reduce(&mut tensor, ReduceOp::Sum).await?;
|
|
let duration = start.elapsed();
|
|
|
|
info!("AllReduce completed in {:?}", duration);
|
|
info!(
|
|
"Result tensor sum: {:.2}",
|
|
tensor.data()?.iter().sum::<f32>()
|
|
);
|
|
|
|
// Test different reduce operations
|
|
let mut max_tensor = Tensor::from_data(vec![1.0, 5.0, 3.0, 2.0], vec![4], &Device::default())?;
|
|
|
|
pg.all_reduce(&mut max_tensor, ReduceOp::Max).await?;
|
|
info!("Max reduce result: {:?}", max_tensor.data()?);
|
|
|
|
Ok(())
|
|
}
|
|
|
|
#[cfg(feature = "nccl")]
|
|
async fn await_broadcast_example(pg: &ProcessGroup) -> Result<()> {
|
|
// Create test tensor for broadcast
|
|
let mut broadcast_tensor = Tensor::from_data(vec![42.0; 100], vec![100], &Device::default())?;
|
|
|
|
info!("Broadcasting tensor with value 42.0");
|
|
|
|
let start = Instant::now();
|
|
pg.broadcast(&mut broadcast_tensor, 0).await?;
|
|
let duration = start.elapsed();
|
|
|
|
info!("Broadcast completed in {:?}", duration);
|
|
info!("First few values: {:?}", &broadcast_tensor.data()?[0..5]);
|
|
|
|
Ok(())
|
|
}
|
|
|
|
#[cfg(feature = "nccl")]
|
|
async fn await_performance_comparison(pg: &ProcessGroup, backend_name: &str) -> Result<()> {
|
|
let sizes = vec![1024, 10240, 102400];
|
|
|
|
info!("Performance comparison with {} backend:", backend_name);
|
|
|
|
for size in sizes {
|
|
let mut tensor = Tensor::ones(&vec![size], &Device::default())?;
|
|
|
|
let start = Instant::now();
|
|
pg.all_reduce(&mut tensor, ReduceOp::Sum).await?;
|
|
let duration = start.elapsed();
|
|
|
|
let bandwidth_gb = (size * 4) as f64 / (1024.0 * 1024.0 * 1024.0); // Approximate
|
|
let bandwidth_gbps = bandwidth_gb / duration.as_secs_f64();
|
|
|
|
info!(
|
|
"Size: {} elements, Time: {:?}, Bandwidth: {:.2} GB/s",
|
|
size, duration, bandwidth_gbps
|
|
);
|
|
}
|
|
|
|
Ok(())
|
|
}
|
|
|
|
#[cfg(feature = "nccl")]
|
|
async fn await_tensor_operations_example(pg: &ProcessGroup) -> Result<()> {
|
|
// Test different tensor shapes
|
|
let test_shapes = vec![
|
|
vec![1],
|
|
vec![100],
|
|
vec![10, 10],
|
|
vec![5, 5, 5],
|
|
vec![2, 2, 2, 2, 2],
|
|
];
|
|
|
|
for shape_dims in test_shapes {
|
|
info!("Testing shape: {:?}", shape_dims);
|
|
|
|
let mut tensor = Tensor::ones(&shape_dims, &Device::default())?;
|
|
let original_sum: f32 = tensor.data()?.iter().sum();
|
|
|
|
// AllReduce
|
|
pg.all_reduce(&mut tensor, ReduceOp::Sum).await?;
|
|
let new_sum: f32 = tensor.data()?.iter().sum();
|
|
|
|
info!(
|
|
" Original sum: {:.2}, After AllReduce: {:.2}",
|
|
original_sum, new_sum
|
|
);
|
|
|
|
// Broadcast
|
|
let mut broadcast_tensor =
|
|
Tensor::zeros(TensorShape::new(shape_dims.clone())?, &Device::default())?;
|
|
pg.broadcast(&mut broadcast_tensor, 0).await?;
|
|
|
|
info!(" Broadcast completed for shape {:?}", shape_dims);
|
|
}
|
|
|
|
Ok(())
|
|
}
|
|
|
|
#[cfg(feature = "nccl")]
|
|
async fn await_error_handling_example(pg: &ProcessGroup) -> Result<()> {
|
|
info!("Testing error conditions:");
|
|
|
|
// Test broadcast with invalid root
|
|
let mut tensor = Tensor::ones(&vec![10], &Device::default())?;
|
|
match pg.broadcast(&mut tensor, 999).await {
|
|
Ok(_) => warn!("Broadcast with invalid root should have failed"),
|
|
Err(e) => info!(" ✓ Correctly caught invalid root error: {}", e),
|
|
}
|
|
|
|
// Test with very large tensor (might succeed or fail depending on memory)
|
|
match Tensor::ones(&vec![100_000_000], &Device::default()) {
|
|
Ok(mut large_tensor) => {
|
|
info!(" Testing with large tensor (100M elements)");
|
|
match pg.all_reduce(&mut large_tensor, ReduceOp::Sum).await {
|
|
Ok(_) => info!(" ✓ Large tensor operation succeeded"),
|
|
Err(e) => info!(" ✓ Large tensor operation failed as expected: {}", e),
|
|
}
|
|
}
|
|
Err(e) => info!(" ✓ Large tensor creation failed as expected: {}", e),
|
|
}
|
|
|
|
Ok(())
|
|
}
|
|
|
|
#[cfg(not(feature = "nccl"))]
|
|
async fn run_nccl_example() -> Result<()> {
|
|
warn!("This example requires the 'nccl' feature to be enabled");
|
|
warn!("Please run with: cargo run --example nccl_example --features nccl");
|
|
Ok(())
|
|
}
|