fix: synchronize stream before DtoH readback after CUDA Graph launch

CUDA Graph capture disables cudarc event tracking, so device_ptr()
cannot rely on write events for ordering. Without explicit sync,
the DtoH memcpy can get CUDA_ERROR_INVALID_VALUE because the
kernel writes haven't completed yet.

Add stream.synchronize() between graph launch and scalar readback.
This ensures kernel output buffers (total_loss, grad_norm) are valid
before the host reads them. The sync cost is negligible — it's
1 sync per training step (already bottlenecked by kernel execution).

Fixes H100 CI failure: all DQN pipeline tests returned
"DtoH grad_norm: CUDA_ERROR_INVALID_VALUE" consistently.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-03-20 21:33:04 +01:00
parent a4f84a253e
commit 523139aac4

View File

@@ -864,8 +864,15 @@ impl GpuDqnTrainer {
self.submit_training_ops(online_dueling, online_branching)?;
}
// ── Synchronize stream before readback ──────────────────────
// CUDA Graph replay and direct kernel launches are async. The stream
// must be synchronized before DtoH to ensure kernel writes are visible.
// Also required because graph capture disables event tracking, so
// cudarc's device_ptr() can't rely on write events for ordering.
self.stream.synchronize()
.map_err(|e| MLError::ModelError(format!("stream sync before readback: {e}")))?;
// ── Gather 2 scalars to host: total_loss + grad_norm ────────
// Read each 1-element GPU buffer directly to host (2 × 4B DtoH).
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}")))?;