gitadmin/foxhunt
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

feat(cuda): add GpuTrainingGuard wrapper with pinned memory + accumulators

Browse Source 
Wraps 4 CUDA kernels (training_guard_check, training_guard_accumulate,
qvalue_stats_reduce, qvalue_divergence_check) with a Rust struct that
uses OnceLock PTX caching, pre-allocated device buffers, and host-side
Vec mirrors for zero-allocation readbacks per training step.
Accumulator (3-float acc_buf) stays on-device for epoch-boundary
averaging without CPU roundtrips.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-03-10 14:04:49 +01:00
parent fb1fac3e93
commit 966300aa49
2 changed files with 568 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

566
crates/ml/src/cuda_pipeline/gpu_training_guard.rs Normal file
View File

@@ -0,0 +1,566 @@
#![allow(unsafe_code)]
//! GPU-resident training guard and Q-value monitor.
//!
//! Wraps four CUDA kernels from `training_guard_kernel.cu` to perform
//! NaN/Inf detection, loss clipping, gradient collapse checks, and
//! Q-value statistics entirely on GPU.
//!
//! All results are read back via small DtoH copies (7 or 4 or 5 floats).
//! The accumulator buffer is 3 floats and stays on-device between steps;
//! a single readback at epoch boundary returns (mean_loss, mean_grad_norm).
use candle_core::cuda_backend::cudarc;
use candle_core::{Device, DType, Tensor};
use cudarc::driver::{CudaFunction, CudaSlice, LaunchConfig, PushKernelArg};
use cudarc::nvrtc::Ptx;
use std::sync::OnceLock;
use crate::MLError;
// ── PTX cache ──────────────────────────────────────────────────────────────
static TRAINING_GUARD_PTX: OnceLock<Result<Ptx, String>> = OnceLock::new();
fn compile_training_guard_ptx() -> Result<Ptx, String> {
let kernel_src = include_str!("training_guard_kernel.cu");
cudarc::nvrtc::compile_ptx(kernel_src)
.map_err(|e| format!("training_guard_kernel CUDA compilation failed: {e}"))
}
// ── Result types ───────────────────────────────────────────────────────────
/// Result from `training_guard_check`: per-step safety flags and scalar values.
#[derive(Debug, Clone)]
pub struct GuardResult {
/// True if loss or grad_norm is NaN or Inf.
pub halt_nan: bool,
/// True if loss exceeded the clip threshold.
pub halt_loss_clip: bool,
/// True if grad_norm fell below the collapse threshold (and not in warmup).
pub halt_grad_collapse: bool,
/// Loss after clipping (min(loss, clip_threshold)); NaN/Inf pass through.
pub clipped_loss: f32,
/// Raw loss as read from the GPU scalar.
pub raw_loss: f32,
/// Raw gradient norm as read from the GPU scalar.
pub raw_grad_norm: f32,
}
/// Per-batch Q-value statistics from `qvalue_stats_reduce`.
#[derive(Debug, Clone)]
pub struct QValueStats {
/// Min of per-sample max-Q across the batch.
pub q_min: f32,
/// Max of per-sample max-Q across the batch.
pub q_max: f32,
/// Mean of per-sample max-Q across the batch.
pub q_mean: f32,
/// Mean of all Q-values (all actions × all samples; collapse detection).
pub q_all_mean: f32,
}
/// Single-sample Q-value divergence result from `qvalue_divergence_check`.
#[derive(Debug, Clone)]
pub struct QValueDivergence {
/// Minimum Q-value.
pub q_min: f32,
/// Maximum Q-value.
pub q_max: f32,
/// Mean Q-value.
pub q_mean: f32,
/// Variance of Q-values.
pub q_variance: f32,
/// True if |q_min| or |q_max| exceeds the divergence threshold.
pub divergence_detected: bool,
}
// ── Main struct ────────────────────────────────────────────────────────────
/// GPU-resident training guard + Q-value monitor.
///
/// Holds pre-compiled CUDA functions, device output buffers, and pinned
/// host-side read buffers. Reusable across training steps without
/// re-allocation.
pub struct GpuTrainingGuard {
check_func: CudaFunction,
accumulate_func: CudaFunction,
qvalue_stats_func: CudaFunction,
qvalue_div_func: CudaFunction,
/// Output buffer for `training_guard_check` (7 floats, device memory).
guard_output_dev: CudaSlice<f32>,
/// Host mirror of `guard_output_dev`.
pinned_guard: Vec<f32>,
/// Output buffer for `qvalue_stats_reduce` (4 floats, device memory).
qstats_output_dev: CudaSlice<f32>,
/// Host mirror of `qstats_output_dev`.
pinned_qstats: Vec<f32>,
/// Output buffer for `qvalue_divergence_check` (5 floats, device memory).
qdiv_output_dev: CudaSlice<f32>,
/// Host mirror of `qdiv_output_dev`.
pinned_qdiv: Vec<f32>,
/// Accumulator buffer for `training_guard_accumulate` (3 floats: loss_sum,
/// grad_norm_sum, step_count — stored as float per kernel contract).
acc_buf: CudaSlice<f32>,
device: Device,
}
impl GpuTrainingGuard {
/// Create a new `GpuTrainingGuard` on the given CUDA device.
///
/// Compiles the PTX once per process (OnceLock cached), loads all four
/// kernel functions, and pre-allocates all device and host buffers.
pub fn new(device: &Device) -> Result<Self, MLError> {
let cuda_dev = match device {
Device::Cuda(ref dev) => dev,
_ => {
return Err(MLError::ModelError(
"GpuTrainingGuard requires a CUDA device".into(),
))
}
};
// Compile PTX (once per process)
let ptx_result = TRAINING_GUARD_PTX.get_or_init(compile_training_guard_ptx);
let ptx = ptx_result.as_ref().map_err(|e| {
MLError::ModelError(format!("training_guard PTX: {e}"))
})?;
// Load module and functions
let stream = cuda_dev.cuda_stream();
let context = stream.context();
let module = context.load_module(ptx.clone()).map_err(|e| {
MLError::ModelError(format!("training_guard module load: {e}"))
})?;
let check_func = module
.load_function("training_guard_check")
.map_err(|e| MLError::ModelError(format!("training_guard_check load: {e}")))?;
let accumulate_func = module
.load_function("training_guard_accumulate")
.map_err(|e| MLError::ModelError(format!("training_guard_accumulate load: {e}")))?;
let qvalue_stats_func = module
.load_function("qvalue_stats_reduce")
.map_err(|e| MLError::ModelError(format!("qvalue_stats_reduce load: {e}")))?;
let qvalue_div_func = module
.load_function("qvalue_divergence_check")
.map_err(|e| MLError::ModelError(format!("qvalue_divergence_check load: {e}")))?;
// Allocate device output buffers
let guard_output_dev = stream.alloc_zeros::<f32>(7).map_err(|e| {
MLError::ModelError(format!("alloc guard_output_dev: {e}"))
})?;
let qstats_output_dev = stream.alloc_zeros::<f32>(4).map_err(|e| {
MLError::ModelError(format!("alloc qstats_output_dev: {e}"))
})?;
let qdiv_output_dev = stream.alloc_zeros::<f32>(5).map_err(|e| {
MLError::ModelError(format!("alloc qdiv_output_dev: {e}"))
})?;
let acc_buf = stream.alloc_zeros::<f32>(3).map_err(|e| {
MLError::ModelError(format!("alloc acc_buf: {e}"))
})?;
Ok(Self {
check_func,
accumulate_func,
qvalue_stats_func,
qvalue_div_func,
guard_output_dev,
pinned_guard: vec![0.0_f32; 7],
qstats_output_dev,
pinned_qstats: vec![0.0_f32; 4],
qdiv_output_dev,
pinned_qdiv: vec![0.0_f32; 5],
acc_buf,
device: device.clone(),
})
}
/// Run the guard check and accumulate kernels for one training step.
///
/// `loss_gpu` and `grad_norm_gpu` must be F32 scalar tensors on CUDA (shape `[]` or `[1]`).
/// Returns the safety flags and scalar values from the check kernel.
pub fn check_and_accumulate(
&mut self,
loss_gpu: &Tensor,
grad_norm_gpu: &Tensor,
clip_threshold: f32,
collapse_threshold: f32,
warmup: bool,
) -> Result<GuardResult, MLError> {
let cuda_dev = match &self.device {
Device::Cuda(ref dev) => dev,
_ => return Err(MLError::ModelError("GpuTrainingGuard: not CUDA".into())),
};
let stream = cuda_dev.cuda_stream();
// Cast to F32 if needed
let loss_f32 = if loss_gpu.dtype() == DType::F32 {
loss_gpu.clone()
} else {
loss_gpu
.to_dtype(DType::F32)
.map_err(|e| MLError::ModelError(format!("loss cast to F32: {e}")))?
};
let grad_f32 = if grad_norm_gpu.dtype() == DType::F32 {
grad_norm_gpu.clone()
} else {
grad_norm_gpu
.to_dtype(DType::F32)
.map_err(|e| MLError::ModelError(format!("grad_norm cast to F32: {e}")))?
};
// Ensure contiguous
let loss_cont = loss_f32
.contiguous()
.map_err(|e| MLError::ModelError(format!("loss contiguous: {e}")))?;
let grad_cont = grad_f32
.contiguous()
.map_err(|e| MLError::ModelError(format!("grad_norm contiguous: {e}")))?;
// Extract CudaSlice views
let (loss_guard, loss_layout) = loss_cont.storage_and_layout();
let loss_slice = match &*loss_guard {
candle_core::Storage::Cuda(ref cs) => cs
.as_cuda_slice::<f32>()
.map_err(|e| MLError::ModelError(format!("loss as_cuda_slice: {e}")))?,
_ => return Err(MLError::ModelError("loss not on CUDA".into())),
};
let loss_view = loss_slice.slice(loss_layout.start_offset()..);
let (grad_guard, grad_layout) = grad_cont.storage_and_layout();
let grad_slice = match &*grad_guard {
candle_core::Storage::Cuda(ref cs) => cs
.as_cuda_slice::<f32>()
.map_err(|e| MLError::ModelError(format!("grad_norm as_cuda_slice: {e}")))?,
_ => return Err(MLError::ModelError("grad_norm not on CUDA".into())),
};
let grad_view = grad_slice.slice(grad_layout.start_offset()..);
// Zero the guard output buffer before writing
stream.memset_zeros(&mut self.guard_output_dev).map_err(|e| {
MLError::ModelError(format!("memset guard_output_dev: {e}"))
})?;
let warmup_int: i32 = if warmup { 1 } else { 0 };
// Launch kernel 1: training_guard_check — grid=(1,1,1), block=(1,1,1)
let single_cfg = LaunchConfig {
grid_dim: (1, 1, 1),
block_dim: (1, 1, 1),
shared_mem_bytes: 0,
};
// Safety: all GPU slices are valid, layout offsets applied, buffers sized correctly.
unsafe {
stream
.launch_builder(&self.check_func)
.arg(&loss_view)
.arg(&grad_view)
.arg(&mut self.guard_output_dev)
.arg(&clip_threshold)
.arg(&collapse_threshold)
.arg(&warmup_int)
.launch(single_cfg)
.map_err(|e| MLError::ModelError(format!("training_guard_check launch: {e}")))?;
}
// Launch kernel 2: training_guard_accumulate — grid=(1,1,1), block=(1,1,1)
// Safety: loss_view/grad_view valid; acc_buf is 3 f32 on same device.
unsafe {
stream
.launch_builder(&self.accumulate_func)
.arg(&loss_view)
.arg(&grad_view)
.arg(&mut self.acc_buf)
.launch(single_cfg)
.map_err(|e| {
MLError::ModelError(format!("training_guard_accumulate launch: {e}"))
})?;
}
// DtoH readback of check results (7 floats = 28 bytes)
stream
.memcpy_dtoh(&self.guard_output_dev, &mut self.pinned_guard)
.map_err(|e| MLError::ModelError(format!("guard readback: {e}")))?;
// Drop storage guards before any tensor creation
drop(loss_guard);
drop(grad_guard);
let g = &self.pinned_guard;
Ok(GuardResult {
halt_nan: g[0] != 0.0,
halt_loss_clip: g[1] != 0.0,
halt_grad_collapse: g[2] != 0.0,
clipped_loss: g[3],
raw_loss: g[4],
raw_grad_norm: g[5],
})
}
/// Read the epoch-boundary loss and grad_norm averages from the accumulator.
///
/// Returns `(mean_loss, mean_grad_norm)` as f64. The accumulator buffer
/// stays on-device; call `reset_accumulators()` after reading to clear it.
pub fn read_accumulators(&mut self) -> Result<(f64, f64), MLError> {
let cuda_dev = match &self.device {
Device::Cuda(ref dev) => dev,
_ => return Err(MLError::ModelError("GpuTrainingGuard: not CUDA".into())),
};
let stream = cuda_dev.cuda_stream();
// acc_buf layout: [0]=loss_sum, [1]=grad_norm_sum, [2]=step_count (float)
let mut host = [0.0_f32; 3];
stream.memcpy_dtoh(&self.acc_buf, &mut host).map_err(|e| {
MLError::ModelError(format!("acc_buf readback: {e}"))
})?;
let loss_sum = host[0] as f64;
let grad_sum = host[1] as f64;
let steps = host[2] as f64;
if steps <= 0.0 {
return Ok((0.0, 0.0));
}
Ok((loss_sum / steps, grad_sum / steps))
}
/// Zero the accumulator buffer on-device (call at epoch start/end).
pub fn reset_accumulators(&mut self) -> Result<(), MLError> {
let cuda_dev = match &self.device {
Device::Cuda(ref dev) => dev,
_ => return Err(MLError::ModelError("GpuTrainingGuard: not CUDA".into())),
};
let stream = cuda_dev.cuda_stream();
stream.memset_zeros(&mut self.acc_buf).map_err(|e| {
MLError::ModelError(format!("reset acc_buf: {e}"))
})?;
Ok(())
}
/// Compute Q-value statistics over a full batch via `qvalue_stats_reduce`.
///
/// `q_values` must be F32 with shape `[batch_size, num_actions]` on CUDA.
pub fn qvalue_stats(
&mut self,
q_values: &Tensor,
batch_size: usize,
num_actions: usize,
) -> Result<QValueStats, MLError> {
if batch_size == 0 || num_actions == 0 {
return Ok(QValueStats {
q_min: 0.0,
q_max: 0.0,
q_mean: 0.0,
q_all_mean: 0.0,
});
}
let cuda_dev = match &self.device {
Device::Cuda(ref dev) => dev,
_ => return Err(MLError::ModelError("GpuTrainingGuard: not CUDA".into())),
};
let stream = cuda_dev.cuda_stream();
let q_f32 = if q_values.dtype() == DType::F32 {
q_values.clone()
} else {
q_values
.to_dtype(DType::F32)
.map_err(|e| MLError::ModelError(format!("q_values cast to F32: {e}")))?
};
let q_cont = q_f32
.contiguous()
.map_err(|e| MLError::ModelError(format!("q_values contiguous: {e}")))?;
let (q_guard, q_layout) = q_cont.storage_and_layout();
let q_slice = match &*q_guard {
candle_core::Storage::Cuda(ref cs) => cs
.as_cuda_slice::<f32>()
.map_err(|e| MLError::ModelError(format!("q_values as_cuda_slice: {e}")))?,
_ => return Err(MLError::ModelError("q_values not on CUDA".into())),
};
let q_view = q_slice.slice(q_layout.start_offset()..);
// Zero output buffer
stream.memset_zeros(&mut self.qstats_output_dev).map_err(|e| {
MLError::ModelError(format!("memset qstats_output_dev: {e}"))
})?;
let bs_i32 = batch_size as i32;
let na_i32 = num_actions as i32;
// qvalue_stats_reduce: grid=(1,1,1), block=(256,1,1), grid-stride handles N>256
let cfg = LaunchConfig {
grid_dim: (1, 1, 1),
block_dim: (256, 1, 1),
shared_mem_bytes: 0,
};
// Safety: q_view has batch_size * num_actions elements. Output is 4 f32.
unsafe {
stream
.launch_builder(&self.qvalue_stats_func)
.arg(&q_view)
.arg(&mut self.qstats_output_dev)
.arg(&bs_i32)
.arg(&na_i32)
.launch(cfg)
.map_err(|e| MLError::ModelError(format!("qvalue_stats_reduce launch: {e}")))?;
}
stream
.memcpy_dtoh(&self.qstats_output_dev, &mut self.pinned_qstats)
.map_err(|e| MLError::ModelError(format!("qstats readback: {e}")))?;
drop(q_guard);
let s = &self.pinned_qstats;
Ok(QValueStats {
q_min: s[0],
q_max: s[1],
q_mean: s[2],
q_all_mean: s[3],
})
}
/// Check a single sample for Q-value divergence via `qvalue_divergence_check`.
///
/// `q_values` must be F32 with shape `[num_actions]` on CUDA.
pub fn qvalue_divergence(
&mut self,
q_values: &Tensor,
num_actions: usize,
threshold: f32,
) -> Result<QValueDivergence, MLError> {
if num_actions == 0 {
return Ok(QValueDivergence {
q_min: 0.0,
q_max: 0.0,
q_mean: 0.0,
q_variance: 0.0,
divergence_detected: false,
});
}
let cuda_dev = match &self.device {
Device::Cuda(ref dev) => dev,
_ => return Err(MLError::ModelError("GpuTrainingGuard: not CUDA".into())),
};
let stream = cuda_dev.cuda_stream();
let q_f32 = if q_values.dtype() == DType::F32 {
q_values.clone()
} else {
q_values
.to_dtype(DType::F32)
.map_err(|e| MLError::ModelError(format!("q_values cast to F32: {e}")))?
};
let q_cont = q_f32
.contiguous()
.map_err(|e| MLError::ModelError(format!("q_values contiguous: {e}")))?;
let (q_guard, q_layout) = q_cont.storage_and_layout();
let q_slice = match &*q_guard {
candle_core::Storage::Cuda(ref cs) => cs
.as_cuda_slice::<f32>()
.map_err(|e| MLError::ModelError(format!("q_values as_cuda_slice: {e}")))?,
_ => return Err(MLError::ModelError("q_values not on CUDA".into())),
};
let q_view = q_slice.slice(q_layout.start_offset()..);
// Zero output buffer
stream.memset_zeros(&mut self.qdiv_output_dev).map_err(|e| {
MLError::ModelError(format!("memset qdiv_output_dev: {e}"))
})?;
let na_i32 = num_actions as i32;
// qvalue_divergence_check: grid=(1,1,1), block=(1,1,1)
let single_cfg = LaunchConfig {
grid_dim: (1, 1, 1),
block_dim: (1, 1, 1),
shared_mem_bytes: 0,
};
// Safety: q_view has num_actions elements. Output is 5 f32.
unsafe {
stream
.launch_builder(&self.qvalue_div_func)
.arg(&q_view)
.arg(&mut self.qdiv_output_dev)
.arg(&na_i32)
.arg(&threshold)
.launch(single_cfg)
.map_err(|e| {
MLError::ModelError(format!("qvalue_divergence_check launch: {e}"))
})?;
}
stream
.memcpy_dtoh(&self.qdiv_output_dev, &mut self.pinned_qdiv)
.map_err(|e| MLError::ModelError(format!("qdiv readback: {e}")))?;
drop(q_guard);
let d = &self.pinned_qdiv;
Ok(QValueDivergence {
q_min: d[0],
q_max: d[1],
q_mean: d[2],
q_variance: d[3],
divergence_detected: d[4] != 0.0,
})
}
}
impl std::fmt::Debug for GpuTrainingGuard {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("GpuTrainingGuard")
.field("device", &self.device)
.finish()
}
}
#[cfg(test)]
mod tests {
use super::*;
/// Verify the PTX compiles without errors (skips gracefully when NVRTC is absent).
#[test]
fn test_ptx_compilation() {
let result = compile_training_guard_ptx();
if let Err(ref e) = result {
if e.contains("NVRTC")
|| e.contains("nvrtc")
|| e.contains("not found")
|| e.contains("No such file")
{
// NVRTC not installed — acceptable on CPU-only machines.
return;
}
panic!("training_guard PTX compilation failed: {e}");
}
}
/// Verify `GpuTrainingGuard::new` fails gracefully on a CPU device.
#[test]
fn test_cpu_device_rejected() {
let result = GpuTrainingGuard::new(&Device::Cpu);
assert!(result.is_err());
let err_msg = format!("{}", result.err().expect("should be error"));
assert!(
err_msg.contains("CUDA"),
"Error should mention CUDA: {err_msg}"
);
}
}

2
crates/ml/src/cuda_pipeline/mod.rs
View File

@@ -24,6 +24,8 @@ pub mod gpu_ppo_collector;
pub mod gpu_action_selector;
#[cfg(feature = "cuda")]
pub mod gpu_statistics;
#[cfg(feature = "cuda")]
pub mod gpu_training_guard;
// gpu_replay_buffer moved to ml-dqn crate
/// Maximum bytes allowed for a single GPU upload (2 GB safety limit).