distributed: NCCL tensor-parallel primitives (TpContext + AllReduce)
New xserv-distributed crate: hand-written NCCL FFI, TpContext (one rank per thread, bound to one GPU), and in-place BF16 AllReduce on the null stream so it orders naturally with the model's kernels. 2-GPU AllReduce test included. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
This commit is contained in:
65
crates/xserv-distributed/src/ffi.rs
Normal file
65
crates/xserv-distributed/src/ffi.rs
Normal file
@@ -0,0 +1,65 @@
|
||||
//! Minimal NCCL FFI bindings (hand-written, like the CUDA bindings).
|
||||
//! Only the collectives we need for tensor parallelism.
|
||||
|
||||
use std::ffi::c_void;
|
||||
use std::os::raw::c_char;
|
||||
use xserv_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 all ranks.
|
||||
#[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)
|
||||
pub const NCCL_FLOAT32: i32 = 7;
|
||||
pub const NCCL_BF16: i32 = 9;
|
||||
|
||||
// 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