fix: raw DtoH for graph-modified scalar buffers (H100 stale event fix)

After CUDA Graph replay with event tracking re-enabled, the total_loss
and grad_norm CudaSlice buffers have stale write events from capture.
cudarc's memcpy_dtoh calls device_ptr() which waits on these stale events,
causing CUDA_ERROR_INVALID_VALUE on H100.

Fix: use raw_device_ptr (ManuallyDrop guard) + cuMemcpyDtoH_v2 for the
2 scalar readbacks (8 bytes total — legitimate GPU exit for metrics).
Stream is synchronized before readback, so event ordering is guaranteed.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-03-20 22:53:43 +01:00
parent 021de69235
commit 8c50e68fe7

View File

@@ -869,12 +869,26 @@ impl GpuDqnTrainer {
.map_err(|e| MLError::ModelError(format!("stream sync before readback: {e}")))?;
// ── Gather 2 scalars to host: total_loss + grad_norm ────────
// These buffers were written by the CUDA Graph (which runs with event
// tracking disabled). Their CudaSlice write events are stale, so
// cudarc's memcpy_dtoh would fail with CUDA_ERROR_INVALID_VALUE.
// Use raw_device_ptr (ManuallyDrop guard) + synchronous cuMemcpyDtoH.
let mut loss_host = [0.0_f32; 1];
self.stream.memcpy_dtoh(&self.total_loss_buf, &mut loss_host)
.map_err(|e| MLError::ModelError(format!("DtoH total_loss: {e}")))?;
let mut norm_host = [0.0_f32; 1];
self.stream.memcpy_dtoh(&self.grad_norm_buf, &mut norm_host)
.map_err(|e| MLError::ModelError(format!("DtoH grad_norm: {e}")))?;
let loss_ptr = raw_device_ptr(&self.total_loss_buf, &self.stream);
let norm_ptr = raw_device_ptr(&self.grad_norm_buf, &self.stream);
let r1 = unsafe { cudarc::driver::sys::cuMemcpyDtoH_v2(
loss_host.as_mut_ptr().cast(), loss_ptr, 4,
)};
if r1 != cudarc::driver::sys::CUresult::CUDA_SUCCESS {
return Err(MLError::ModelError(format!("DtoH total_loss: {r1:?}")));
}
let r2 = unsafe { cudarc::driver::sys::cuMemcpyDtoH_v2(
norm_host.as_mut_ptr().cast(), norm_ptr, 4,
)};
if r2 != cudarc::driver::sys::CUresult::CUDA_SUCCESS {
return Err(MLError::ModelError(format!("DtoH grad_norm: {r2:?}")));
}
self.scalars_readback_host = [loss_host[0], norm_host[0]];
Ok(FusedTrainScalars {