121 lines
4.4 KiB
Rust
121 lines
4.4 KiB
Rust
//! Standalone NCCL test to verify cudarc integration works
|
|
//!
|
|
//! This file tests NCCL functionality directly using cudarc
|
|
//! without depending on the rtx-distributed library.
|
|
//!
|
|
//! Run with: cargo run --bin nccl_test_standalone --features nccl
|
|
|
|
use std::sync::Arc;
|
|
|
|
#[cfg(feature = "nccl")]
|
|
fn main() -> Result<(), Box<dyn std::error::Error>> {
|
|
use cudarc::driver::CudaContext;
|
|
use cudarc::nccl::{Id, Comm, ReduceOp as NcclReduceOp};
|
|
|
|
println!("🚀 Starting NCCL integration test...");
|
|
|
|
// Test 1: CUDA Context Creation
|
|
println!("\n1️⃣ Testing CUDA context creation...");
|
|
let ctx = match CudaContext::new(0) {
|
|
Ok(ctx) => {
|
|
println!("✅ CUDA context created successfully on device {}", ctx.ordinal());
|
|
Arc::new(ctx)
|
|
}
|
|
Err(e) => {
|
|
println!("❌ CUDA context creation failed: {:?}", e);
|
|
println!("💡 This might be expected if you don't have CUDA installed");
|
|
return Ok(());
|
|
}
|
|
};
|
|
|
|
// Test 2: NCCL ID Generation
|
|
println!("\n2️⃣ Testing NCCL ID generation...");
|
|
let id = match Id::new() {
|
|
Ok(id) => {
|
|
println!("✅ NCCL ID generated successfully");
|
|
id
|
|
}
|
|
Err(e) => {
|
|
println!("❌ NCCL ID generation failed: {:?}", e);
|
|
println!("💡 This might be expected if you don't have NCCL installed");
|
|
return Ok(());
|
|
}
|
|
};
|
|
|
|
// Test 3: NCCL Communicator Creation
|
|
println!("\n3️⃣ Testing NCCL communicator creation...");
|
|
let stream = ctx.default_stream();
|
|
let comm = match Comm::from_rank(stream.clone(), 0, 1, id) {
|
|
Ok(comm) => {
|
|
println!("✅ NCCL communicator created successfully");
|
|
println!(" Rank: {}, World size: {}", comm.rank(), comm.world_size());
|
|
comm
|
|
}
|
|
Err(e) => {
|
|
println!("❌ NCCL communicator creation failed: {:?}", e);
|
|
return Err(e.into());
|
|
}
|
|
};
|
|
|
|
// Test 4: Memory Allocation and AllReduce
|
|
println!("\n4️⃣ Testing NCCL AllReduce operation...");
|
|
let input_data = vec![1.0f32, 2.0, 3.0, 4.0];
|
|
println!(" Input data: {:?}", input_data);
|
|
|
|
// Allocate device memory
|
|
let input_slice = stream.memcpy_stod(&input_data)?;
|
|
let mut output_slice = stream.alloc_zeros::<f32>(input_data.len())?;
|
|
|
|
// Perform AllReduce
|
|
let reduce_op = NcclReduceOp::Sum;
|
|
comm.all_reduce(&input_slice, &mut output_slice, &reduce_op)?;
|
|
|
|
// Synchronize and get results
|
|
stream.synchronize()?;
|
|
let result_data = stream.memcpy_dtov(&output_slice)?;
|
|
|
|
println!(" Output data: {:?}", result_data);
|
|
|
|
// Verify results (for single rank, output should equal input)
|
|
for (i, (&actual, &expected)) in result_data.iter().zip(input_data.iter()).enumerate() {
|
|
if (actual - expected).abs() > 1e-6 {
|
|
println!("❌ Result mismatch at index {}: expected {}, got {}", i, expected, actual);
|
|
return Err("AllReduce result verification failed".into());
|
|
}
|
|
}
|
|
println!("✅ AllReduce operation completed successfully");
|
|
|
|
// Test 5: Broadcast Operation
|
|
println!("\n5️⃣ Testing NCCL Broadcast operation...");
|
|
let broadcast_data = vec![10.0f32, 20.0, 30.0, 40.0];
|
|
println!(" Broadcast data: {:?}", broadcast_data);
|
|
|
|
let broadcast_input = stream.memcpy_stod(&broadcast_data)?;
|
|
let mut broadcast_output = stream.alloc_zeros::<f32>(broadcast_data.len())?;
|
|
|
|
// For single rank, we're both sender and receiver
|
|
comm.broadcast(Some(&broadcast_input), &mut broadcast_output, 0)?;
|
|
|
|
stream.synchronize()?;
|
|
let broadcast_result = stream.memcpy_dtov(&broadcast_output)?;
|
|
|
|
println!(" Broadcast result: {:?}", broadcast_result);
|
|
|
|
// Verify broadcast results
|
|
for (i, (&actual, &expected)) in broadcast_result.iter().zip(broadcast_data.iter()).enumerate() {
|
|
if (actual - expected).abs() > 1e-6 {
|
|
println!("❌ Broadcast result mismatch at index {}: expected {}, got {}", i, expected, actual);
|
|
return Err("Broadcast result verification failed".into());
|
|
}
|
|
}
|
|
println!("✅ Broadcast operation completed successfully");
|
|
|
|
println!("\n🎉 All NCCL tests passed! Integration is working correctly.");
|
|
Ok(())
|
|
}
|
|
|
|
#[cfg(not(feature = "nccl"))]
|
|
fn main() {
|
|
println!("⚠️ NCCL feature is not enabled.");
|
|
println!("💡 Run with: cargo run --bin nccl_test_standalone --features nccl");
|
|
} |