Gahow Wang
e27df50ca9
New crate xtrain-distributed (mirrors xserv-distributed): hand-written NCCL
FFI (GetUniqueId / CommInitRank / AllReduce / CommDestroy / Group{Start,End},
ncclUniqueId passed by value per the NCCL ABI) and a safe DdpContext wrapper —
rank 0 mints the UniqueId, every rank inits its communicator under a group, and
all_reduce_average_grads in-place AllReduce(sum)s each param's .grad() device
buffer then scales by 1/world (reuses T7's scale_inplace kernel). AllReduce runs
on the null stream so it orders with the model's kernels (no extra barrier).
build.rs follows the per-crate convention: no nvcc -> no_cuda cfg (crate
compiles to empty, cargo check passes host-side); with nvcc, links -lnccl
-lcudart like xserv-distributed's build.rs.
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
77 lines
2.0 KiB
Rust
77 lines
2.0 KiB
Rust
//! Minimal NCCL FFI bindings (hand-written, like the CUDA bindings in
|
|
//! xtrain-cuda). Only the collectives data-parallel training needs:
|
|
//! unique-id creation, communicator init/destroy, and AllReduce. Mirrors
|
|
//! xserv-distributed's FFI.
|
|
|
|
use std::ffi::c_void;
|
|
use std::os::raw::c_char;
|
|
use xtrain_cuda::ffi::CudaStream;
|
|
|
|
/// Opaque NCCL communicator handle (`ncclComm_t`).
|
|
pub type NcclComm = *mut c_void;
|
|
|
|
/// `ncclUniqueId` is a 128-byte opaque blob shared from rank 0 to every rank.
|
|
#[repr(C)]
|
|
#[derive(Clone, Copy)]
|
|
pub struct NcclUniqueId {
|
|
pub internal: [c_char; 128],
|
|
}
|
|
|
|
impl Default for NcclUniqueId {
|
|
fn default() -> Self {
|
|
Self { internal: [0; 128] }
|
|
}
|
|
}
|
|
|
|
// ncclDataType_t (subset) — DDP all-reduces fp32 gradients.
|
|
pub const NCCL_FLOAT32: i32 = 7;
|
|
|
|
// ncclRedOp_t
|
|
pub const NCCL_SUM: i32 = 0;
|
|
|
|
// ncclResult_t
|
|
pub const NCCL_SUCCESS: i32 = 0;
|
|
|
|
unsafe extern "C" {
|
|
pub fn ncclGetUniqueId(uid: *mut NcclUniqueId) -> i32;
|
|
// ncclUniqueId is passed BY VALUE (a 128-byte struct) per the NCCL ABI.
|
|
pub fn ncclCommInitRank(
|
|
comm: *mut NcclComm,
|
|
nranks: i32,
|
|
commid: NcclUniqueId,
|
|
rank: i32,
|
|
) -> i32;
|
|
pub fn ncclCommDestroy(comm: NcclComm) -> i32;
|
|
pub fn ncclAllReduce(
|
|
sendbuff: *const c_void,
|
|
recvbuff: *mut c_void,
|
|
count: usize,
|
|
datatype: i32,
|
|
op: i32,
|
|
comm: NcclComm,
|
|
stream: CudaStream,
|
|
) -> i32;
|
|
pub fn ncclGroupStart() -> i32;
|
|
pub fn ncclGroupEnd() -> i32;
|
|
pub fn ncclGetErrorString(result: i32) -> *const c_char;
|
|
}
|
|
|
|
pub fn err_string(result: i32) -> String {
|
|
unsafe {
|
|
let p = ncclGetErrorString(result);
|
|
if p.is_null() {
|
|
return format!("nccl error {result}");
|
|
}
|
|
std::ffi::CStr::from_ptr(p).to_string_lossy().into_owned()
|
|
}
|
|
}
|
|
|
|
pub fn check(result: i32, what: &str) {
|
|
assert_eq!(
|
|
result,
|
|
NCCL_SUCCESS,
|
|
"{what} failed: {}",
|
|
err_string(result)
|
|
);
|
|
}
|