diff --git a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs index d84b541c7..6ac073cc6 100644 --- a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs +++ b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs @@ -2221,6 +2221,28 @@ impl GpuDqnTrainer { self.stream .memcpy_dtoh(&self.grad_buf.slice(..read_len), host_slice) .map_err(|e| MLError::ModelError(format!("per_branch_grad_norms dtoh: {e}")))?; + // `memcpy_dtoh` into PINNED host memory is asynchronous w.r.t. the + // host — cudarc calls `cuMemcpyDtoHAsync_v2` under the hood, which + // queues a DMA on `self.stream` and returns before the transfer + // completes. Without an explicit post-copy sync the host-side + // sum-of-squares loop below races the DMA and reads a mix of + // newly-transferred bytes and stale bytes left over from the + // previous epoch's readback. That DMA-in-progress race is the + // source of the residual `grad_abs[dir]` non-determinism: the + // direction gradient itself is bit-identical across runs (same + // CUDA kernels, same per-stream cuBLAS handles post-Option-C, same + // `CUBLAS_WORKSPACE_CONFIG`), but the aggregate L2 accumulator + // sums whatever bytes happen to already be in pinned memory when + // the host-side loop reaches each element. Magnitude tolerates the + // same race because mag gradients are ~300× larger than dir (the + // trunk-to-mag-branch path carries most of the Q-loss signal), so + // a few partially-updated bytes contribute a negligible fraction + // of the L2 sum; dir's small magnitude makes it proportionally + // more sensitive. Sync after the copy so both norms read the + // final transferred bytes. + self.stream.synchronize().map_err(|e| { + MLError::ModelError(format!("per_branch_grad_norms post-dtoh sync: {e}")) + })?; let param_sizes = compute_param_sizes(&self.config); let mut norms = [0.0_f32; 4];