From d98ec8d343b3c9b73d7453417ff53315abbc8d7e Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Mon, 6 Apr 2026 01:16:12 +0200 Subject: [PATCH] fix: pinned host memory for async DtoH readback MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit cuMemcpyDtoHAsync with non-pinned host pointers blocks the CPU on CUDA 13 / H100 driver 580+ until ALL pending GPU work completes. With graph_forward + aux_ops + adam queued, this caused multi-minute hangs between training steps. Fix: allocate 12 bytes of pinned (page-locked) host memory via cuMemHostAlloc(DEVICEMAP) for the 3 scalar readbacks (loss, mse_loss, grad_norm). Pinned memory enables true async DtoH — CPU returns immediately, GPU copies when it reaches that point in the stream. Co-Authored-By: Claude Opus 4.6 (1M context) --- .../ml/src/cuda_pipeline/gpu_dqn_trainer.rs | 65 ++++++++++++------- 1 file changed, 40 insertions(+), 25 deletions(-) diff --git a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs index f59ca3817..e66eb596d 100644 --- a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs +++ b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs @@ -636,12 +636,10 @@ pub struct GpuDqnTrainer { /// CUDA event for async scalar readback (replaces per-step stream.synchronize) readback_event: Option, - /// Host-side buffer for deferred loss readback from previous step (blended MSE+C51) - pending_loss: f32, - /// Host-side buffer for deferred MSE loss readback from previous step - pending_mse_loss: f32, - /// Host-side buffer for deferred grad_norm readback from previous step (sum-of-squares) - pending_grad_norm: f32, + /// Pinned host buffer [3 × f32] for async DtoH scalar readback. + /// Layout: [0]=loss, [1]=mse_loss, [2]=grad_norm_sq. + /// Pinned memory enables true async cuMemcpyDtoHAsync without CPU blocking. + readback_pinned: *mut f32, /// Whether there's an in-flight readback to collect readback_pending: bool, @@ -918,6 +916,10 @@ impl Drop for GpuDqnTrainer { self.graph_forward_mse = None; self.graph_forward_ddqn = None; self.graph_adam = None; + // Free pinned host memory + if !self.readback_pinned.is_null() { + let _ = unsafe { cudarc::driver::result::free_host(self.readback_pinned.cast()) }; + } } } @@ -2804,9 +2806,14 @@ impl GpuDqnTrainer { scalars_readback_buf, scalars_readback_host, readback_event: None, - pending_loss: 0.0, - pending_mse_loss: 0.0, - pending_grad_norm: 0.0, + readback_pinned: { + let flags = cudarc::driver::sys::CU_MEMHOSTALLOC_DEVICEMAP; + unsafe { + cudarc::driver::result::malloc_host(3 * std::mem::size_of::(), flags) + .map_err(|e| MLError::ModelError(format!("pinned readback alloc: {e}")))? + as *mut f32 + } + }, readback_pending: false, bf16_states_buf, bf16_next_states_buf, @@ -3004,11 +3011,14 @@ impl GpuDqnTrainer { } } self.readback_pending = false; - // Blend MSE + C51 loss: guard against 0.0 * NaN = NaN (IEEE 754) - let blended = Self::blend_loss(self.c51_alpha, self.pending_mse_loss, self.pending_loss); + // Read from pinned buffer: [0]=loss, [1]=mse_loss, [2]=grad_norm_sq + let (loss, mse_loss, grad_norm_sq) = unsafe { + (*self.readback_pinned, *self.readback_pinned.add(1), *self.readback_pinned.add(2)) + }; + let blended = Self::blend_loss(self.c51_alpha, mse_loss, loss); FusedTrainScalars { total_loss: blended, - grad_norm: self.pending_grad_norm.sqrt(), + grad_norm: grad_norm_sq.sqrt(), } } else { // First step — no previous data @@ -3022,26 +3032,28 @@ impl GpuDqnTrainer { self.replay_adam()?; if adam_step < 3 { eprintln!("adam_readback: step {adam_step} — async readback start"); } - // 3. Sync stream then DtoH readback (synchronous — non-pinned host memory) - // cuMemcpyDtoHAsync with non-pinned host pointers degrades to sync on - // CUDA 13+ / H100 driver 580+, blocking for ALL pending GPU work. - // Explicit sync + sync copy is faster and predictable. - unsafe { cudarc::driver::sys::cuStreamSynchronize(self.stream.cu_stream()); } + // 3. Async DtoH into pinned host buffer (truly async — no CPU blocking) unsafe { - cudarc::driver::sys::cuMemcpyDtoH_v2( - (&mut self.pending_loss as *mut f32).cast(), + // [0] = loss + cudarc::driver::sys::cuMemcpyDtoHAsync_v2( + self.readback_pinned.cast(), self.total_loss_buf.raw_ptr(), std::mem::size_of::(), + self.stream.cu_stream(), ); - cudarc::driver::sys::cuMemcpyDtoH_v2( - (&mut self.pending_mse_loss as *mut f32).cast(), + // [1] = mse_loss + cudarc::driver::sys::cuMemcpyDtoHAsync_v2( + self.readback_pinned.add(1).cast(), self.mse_loss_buf.raw_ptr(), std::mem::size_of::(), + self.stream.cu_stream(), ); - cudarc::driver::sys::cuMemcpyDtoH_v2( - (&mut self.pending_grad_norm as *mut f32).cast(), + // [2] = grad_norm_sq + cudarc::driver::sys::cuMemcpyDtoHAsync_v2( + self.readback_pinned.add(2).cast(), self.grad_norm_f32_buf.raw_ptr(), std::mem::size_of::(), + self.stream.cu_stream(), ); } @@ -3075,10 +3087,13 @@ impl GpuDqnTrainer { } } self.readback_pending = false; - let blended = Self::blend_loss(self.c51_alpha, self.pending_mse_loss, self.pending_loss); + let (loss, mse_loss, grad_norm_sq) = unsafe { + (*self.readback_pinned, *self.readback_pinned.add(1), *self.readback_pinned.add(2)) + }; + let blended = Self::blend_loss(self.c51_alpha, mse_loss, loss); Ok(FusedTrainScalars { total_loss: blended, - grad_norm: self.pending_grad_norm.sqrt(), + grad_norm: grad_norm_sq.sqrt(), }) }