diff --git a/crates/ml/src/cuda_pipeline/gpu_training_guard.rs b/crates/ml/src/cuda_pipeline/gpu_training_guard.rs new file mode 100644 index 000000000..67e4c907c --- /dev/null +++ b/crates/ml/src/cuda_pipeline/gpu_training_guard.rs @@ -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> = OnceLock::new(); + +fn compile_training_guard_ptx() -> Result { + 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, + /// Host mirror of `guard_output_dev`. + pinned_guard: Vec, + + /// Output buffer for `qvalue_stats_reduce` (4 floats, device memory). + qstats_output_dev: CudaSlice, + /// Host mirror of `qstats_output_dev`. + pinned_qstats: Vec, + + /// Output buffer for `qvalue_divergence_check` (5 floats, device memory). + qdiv_output_dev: CudaSlice, + /// Host mirror of `qdiv_output_dev`. + pinned_qdiv: Vec, + + /// Accumulator buffer for `training_guard_accumulate` (3 floats: loss_sum, + /// grad_norm_sum, step_count — stored as float per kernel contract). + acc_buf: CudaSlice, + + 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 { + 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::(7).map_err(|e| { + MLError::ModelError(format!("alloc guard_output_dev: {e}")) + })?; + let qstats_output_dev = stream.alloc_zeros::(4).map_err(|e| { + MLError::ModelError(format!("alloc qstats_output_dev: {e}")) + })?; + let qdiv_output_dev = stream.alloc_zeros::(5).map_err(|e| { + MLError::ModelError(format!("alloc qdiv_output_dev: {e}")) + })?; + let acc_buf = stream.alloc_zeros::(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 { + 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::() + .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::() + .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 { + 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::() + .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 { + 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::() + .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}" + ); + } +} diff --git a/crates/ml/src/cuda_pipeline/mod.rs b/crates/ml/src/cuda_pipeline/mod.rs index 622bff6f3..817ca162b 100644 --- a/crates/ml/src/cuda_pipeline/mod.rs +++ b/crates/ml/src/cuda_pipeline/mod.rs @@ -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).