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

perf: dual graph forward (MSE-only) + NaN-safe scale kernel

Browse Source 
Two optimizations:

1. Dual CUDA graph: graph_forward_mse skips C51 loss/grad/SAXPY blend
   during MSE warmup (c51_alpha=0). Eliminates 7 kernel invocations
   from the graph replay. replay_forward() routes by alpha automatically.

2. NaN-safe dqn_scale_f32_kernel: writes 0.0 directly when alpha==0
   instead of multiplying (IEEE 754: 0*NaN=NaN). Defense-in-depth for
   the C51 gradient buffer which can contain NaN from random logits.

Fold 3 epoch: 5.08s → 4.92s (under 5s target).

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-04-04 20:58:54 +02:00
parent 137480c887
commit 657d7b49ec
2 changed files with 85 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
crates/ml/src/cuda_pipeline/dqn_utility_kernels.cu
View File

@@ -184,7 +184,11 @@ extern "C" __global__ void dqn_scale_f32_kernel(
int n
) {
int i = blockIdx.x * blockDim.x + threadIdx.x;
if (i < n) y[i] *= alpha;
if (i < n) {
/* NaN-safe: IEEE 754 says 0*NaN=NaN. When c51_alpha=0 (MSE warmup),
* C51 grad buffer can contain NaN. Scaling by 0 must yield 0. */
y[i] = (alpha == 0.0f) ? 0.0f : y[i] * alpha;
}
}
/* ══════════════════════════════════════════════════════════════════════

85
crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs
View File

@@ -599,6 +599,8 @@ pub struct GpuDqnTrainer {
// Between A and B: external code can ADD auxiliary gradients to grad_buf
// (IQN trunk, ensemble heads) — single Adam sees combined gradient.
pub(crate) graph_forward: Option<SendSyncGraph>,
/// MSE-only forward graph (no C51 — faster, avoids NaN during warmup).
pub(crate) graph_forward_mse: Option<SendSyncGraph>,
pub(crate) graph_adam: Option<SendSyncGraph>,
// ── Consolidated transfer buffers ─────────────────────────────
@@ -902,6 +904,7 @@ impl Drop for GpuDqnTrainer {
// Synchronize stream and destroy graphs BEFORE CudaSlice fields drop.
unsafe { cudarc::driver::sys::cuStreamSynchronize(self.stream.cu_stream()); }
self.graph_forward = None;
self.graph_forward_mse = None;
self.graph_adam = None;
}
}
@@ -2766,6 +2769,7 @@ impl GpuDqnTrainer {
target_params_initialized: false,
attention_initialized: false,
graph_forward: None,
graph_forward_mse: None,
graph_adam: None,
upload_staging_buf,
upload_staging_len,
@@ -3487,8 +3491,14 @@ impl GpuDqnTrainer {
}
pub fn replay_forward(&self) -> Result<(), MLError> {
if let Some(ref graph) = self.graph_forward {
graph.0.launch().map_err(|e| {
// Route to MSE-only graph during warmup (faster, avoids C51 NaN)
let graph = if self.c51_alpha < 1e-6 {
self.graph_forward_mse.as_ref().or(self.graph_forward.as_ref())
} else {
self.graph_forward.as_ref()
};
if let Some(g) = graph {
g.0.launch().map_err(|e| {
MLError::ModelError(format!("CUDA graph_forward replay: {e}"))
})?;
}
@@ -4018,6 +4028,36 @@ impl GpuDqnTrainer {
)
})?;
// ── Capture graph_forward_mse (MSE-only, no C51) ─────────────────
let begin_mse = self.stream.begin_capture(
cudarc::driver::sys::CUstreamCaptureMode::CU_STREAM_CAPTURE_MODE_THREAD_LOCAL,
);
if let Err(e) = begin_mse {
unsafe { self.stream.context().enable_event_tracking(); }
let _ = self.stream.context().check_err();
return Err(MLError::ModelError(format!("graph_forward_mse begin_capture: {e}")));
}
let submit_mse_result = self.submit_forward_ops_mse_only();
let graph_mse_result = self.stream.end_capture(
cudarc::driver::sys::CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH,
);
if let Err(e) = submit_mse_result {
unsafe { self.stream.context().enable_event_tracking(); }
let _ = self.stream.context().check_err();
return Err(e);
}
let graph_mse = graph_mse_result
.map_err(|e| {
unsafe { self.stream.context().enable_event_tracking(); }
let _ = self.stream.context().check_err();
MLError::ModelError(format!("graph_forward_mse end_capture: {e}"))
})?
.ok_or_else(|| {
unsafe { self.stream.context().enable_event_tracking(); }
let _ = self.stream.context().check_err();
MLError::ModelError("graph_forward_mse capture returned None".into())
})?;
// ── Capture graph_adam ──────────────────────────────────────────
let begin_result = self.stream.begin_capture(
cudarc::driver::sys::CUstreamCaptureMode::CU_STREAM_CAPTURE_MODE_THREAD_LOCAL,
@@ -4053,16 +4093,20 @@ impl GpuDqnTrainer {
graph_fwd.launch().map_err(|e| {
MLError::ModelError(format!("CUDA graph_forward first launch: {e}"))
})?;
graph_mse.launch().map_err(|e| {
MLError::ModelError(format!("CUDA graph_forward_mse first launch: {e}"))
})?;
graph_adam.launch().map_err(|e| {
MLError::ModelError(format!("CUDA graph_adam first launch: {e}"))
})?;
info!(
"GpuDqnTrainer: 2 CUDA graphs captured and launched \
(graph_forward: 5 memsets + forward + loss + grad + backward; \
graph_adam: grad_norm + adam + 20 d2d unflatten)"
"GpuDqnTrainer: 3 CUDA graphs captured and launched \
(graph_forward: MSE+C51; graph_forward_mse: MSE-only; \
graph_adam: grad_norm + adam + unflatten)"
);
self.graph_forward = Some(SendSyncGraph(graph_fwd));
self.graph_forward_mse = Some(SendSyncGraph(graph_mse));
self.graph_adam = Some(SendSyncGraph(graph_adam));
self.last_captured_loss_mode = Some(self.loss_mode);
Ok(())
@@ -4209,6 +4253,36 @@ impl GpuDqnTrainer {
Ok(())
}
/// Submit MSE-only forward ops — no C51 loss/grad/blend.
///
/// Used during MSE warmup (c51_alpha ≈ 0). Eliminates C51 loss, C51 grad,
/// mixup barrier memset, and 4× SAXPY blend kernels. MSE grad writes
/// directly to main gradient buffers (no scratch indirection).
pub(crate) fn submit_forward_ops_mse_only(&mut self) -> Result<(), MLError> {
self.stream.memset_zeros(&mut self.total_loss_buf)
.map_err(|e| MLError::ModelError(format!("zero total_loss: {e}")))?;
self.stream.memset_zeros(&mut self.mse_loss_buf)
.map_err(|e| MLError::ModelError(format!("zero mse_loss: {e}")))?;
self.stream.memset_zeros(&mut self.grad_buf)
.map_err(|e| MLError::ModelError(format!("zero grad_buf: {e}")))?;
self.stream.memset_zeros(&mut self.d_value_logits_buf)
.map_err(|e| MLError::ModelError(format!("zero d_value_logits: {e}")))?;
self.stream.memset_zeros(&mut self.d_adv_logits_buf)
.map_err(|e| MLError::ModelError(format!("zero d_adv_logits: {e}")))?;
self.launch_cublas_forward()?;
self.launch_curiosity_inference()?;
// MSE loss + grad → MAIN buffers directly (no scratch, no SAXPY blend)
self.launch_mse_loss()?;
self.launch_mse_grad_inner(&self.d_value_logits_buf, &self.d_adv_logits_buf)?;
self.cast_d_logits_to_bf16()?;
self.launch_cublas_backward()?;
Ok(())
}
/// Submit the optimizer phase ops to the stream (captured into graph_adam).
///
/// Steps: zero grad_norm → grad_norm → Adam → unflatten.
@@ -4244,6 +4318,7 @@ impl GpuDqnTrainer {
/// The next `train_step()` will re-capture fresh graphs.
pub fn invalidate_training_graph(&mut self) {
self.graph_forward = None;
self.graph_forward_mse = None;
self.graph_adam = None;
self.last_captured_loss_mode = None;
self.params_initialized = false;