dist: nccl ffi + comm bootstrap
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>
This commit is contained in:
76
crates/xtrain-distributed/src/ffi.rs
Normal file
76
crates/xtrain-distributed/src/ffi.rs
Normal file
@@ -0,0 +1,76 @@
|
||||
//! 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)
|
||||
);
|
||||
}
|
||||
Reference in New Issue
Block a user