Gahow Wang
453520d622
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>
98 lines
3.5 KiB
Rust
98 lines
3.5 KiB
Rust
//! Tensor-parallel primitives for xserv.
|
|
//!
|
|
//! Process model: one OS thread per TP rank, each bound to one GPU. NCCL is
|
|
//! used for the collective (AllReduce); a hand-rolled P2P AllReduce may replace
|
|
//! it later as a learning exercise (see docs/17-tensor-parallelism.md).
|
|
|
|
pub mod ffi;
|
|
|
|
use std::ffi::c_void;
|
|
|
|
use ffi::{NcclComm, NcclUniqueId};
|
|
use xserv_cuda::device;
|
|
use xserv_cuda::GpuBuffer;
|
|
|
|
pub use ffi::NcclUniqueId as UniqueId;
|
|
|
|
/// The CUDA "null" (default) stream. The model's kernels and cuBLAS calls run
|
|
/// on it, so issuing NCCL on the same stream keeps AllReduce correctly ordered
|
|
/// after the producing matmul and before the consuming kernel — no extra sync.
|
|
const NULL_STREAM: xserv_cuda::ffi::CudaStream = std::ptr::null_mut();
|
|
|
|
/// Generate a unique id on one rank (typically rank 0) and broadcast the bytes
|
|
/// to all ranks out-of-band (e.g. via a shared variable across threads).
|
|
pub fn get_unique_id() -> NcclUniqueId {
|
|
let mut id = NcclUniqueId::default();
|
|
ffi::check(unsafe { ffi::ncclGetUniqueId(&mut id) }, "ncclGetUniqueId");
|
|
id
|
|
}
|
|
|
|
/// Per-rank tensor-parallel context: NCCL communicator + a dedicated stream.
|
|
pub struct TpContext {
|
|
pub rank: usize,
|
|
pub world: usize,
|
|
pub device: u32,
|
|
comm: NcclComm,
|
|
}
|
|
|
|
// The NCCL communicator is owned by exactly one rank thread.
|
|
unsafe impl Send for TpContext {}
|
|
|
|
impl TpContext {
|
|
/// Initialize this rank. Must be called from the thread that will own this
|
|
/// rank's GPU work; binds the thread to `device` first. All ranks must call
|
|
/// this concurrently with the same `id` and `world`.
|
|
pub fn init(rank: usize, world: usize, id: NcclUniqueId, device: u32) -> Self {
|
|
device::set_device(device).expect("set_device");
|
|
let mut comm: NcclComm = std::ptr::null_mut();
|
|
// Wrap the concurrent inits in a group so they rendezvous without deadlock.
|
|
ffi::check(unsafe { ffi::ncclGroupStart() }, "ncclGroupStart(init)");
|
|
ffi::check(
|
|
unsafe { ffi::ncclCommInitRank(&mut comm, world as i32, id, rank as i32) },
|
|
"ncclCommInitRank",
|
|
);
|
|
ffi::check(unsafe { ffi::ncclGroupEnd() }, "ncclGroupEnd(init)");
|
|
Self { rank, world, device, comm }
|
|
}
|
|
|
|
/// In-place AllReduce(sum) over `count` BF16 elements in `buf`.
|
|
pub fn all_reduce_sum_bf16(&self, buf: &mut GpuBuffer, count: usize) {
|
|
self.all_reduce_sum_bf16_ptr(buf.as_mut_ptr() as *mut c_void, count);
|
|
}
|
|
|
|
/// In-place AllReduce(sum) directly on a device pointer (`count` BF16 elems),
|
|
/// issued on the null stream so it is ordered with the model's kernels.
|
|
/// Asynchronous: a later sync (e.g. the D2H logits copy) completes it.
|
|
///
|
|
/// # Safety
|
|
/// `ptr` must point to at least `count` BF16 elements of valid device memory
|
|
/// on this rank's device. The reduction is in-place (send == recv).
|
|
pub fn all_reduce_sum_bf16_ptr(&self, ptr: *mut c_void, count: usize) {
|
|
if self.world == 1 {
|
|
return; // nothing to reduce
|
|
}
|
|
ffi::check(
|
|
unsafe {
|
|
ffi::ncclAllReduce(
|
|
ptr as *const c_void,
|
|
ptr,
|
|
count,
|
|
ffi::NCCL_BF16,
|
|
ffi::NCCL_SUM,
|
|
self.comm,
|
|
NULL_STREAM,
|
|
)
|
|
},
|
|
"ncclAllReduce",
|
|
);
|
|
}
|
|
}
|
|
|
|
impl Drop for TpContext {
|
|
fn drop(&mut self) {
|
|
if !self.comm.is_null() {
|
|
unsafe { ffi::ncclCommDestroy(self.comm) };
|
|
}
|
|
}
|
|
}
|