From 5da434ab4b2dee4d749647a95f62759520695f31 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Wed, 22 Apr 2026 15:20:08 +0200 Subject: [PATCH] =?UTF-8?q?fix(cuda):=20sync=20after=20async=20memcpy=5Fdt?= =?UTF-8?q?oh=20in=20per=5Fbranch=5Fgrad=5Fnorms=20=E2=80=94=20direction-b?= =?UTF-8?q?ranch=20determinism?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit After Option C (commit 199feff4d, per-stream cublasLt handles) brought cuBLAS-path determinism to bit-identical on magnitude-branch gradients, direction-branch gradients still varied 8.7% at epoch 0 — localized to the direction-branch readback path, not the backward computation. Root cause: `per_branch_grad_norms` in `crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs` called `self.stream.memcpy_dtoh(&grad_buf.slice, pinned_host_slice)` which cudarc forwards to `cuMemcpyDtoHAsync_v2`. Transfers into pinned host memory are asynchronous w.r.t. the host — the API queues the DMA on the stream and returns before the copy completes. The host-side sum-of-squares loop immediately below then races the DMA and reads a mix of newly-transferred bytes and stale bytes left over from the previous epoch's readback. The direction gradient itself is bit- identical across runs (same CUDA kernels, same per-stream cuBLAS handles, same `CUBLAS_WORKSPACE_CONFIG=:4096:8`); a probe that read per-tensor L2 norms (`dir_t0..dir_t3` for w_b0fc, b_b0fc, w_b0out, b_b0out) after the aggregate readback showed tensor-level values bit-identical to 5 sig figs across 3 runs while the aggregate varied 15%, because the probe's own `stream.synchronize()` at call start blocked for the previous async DMA to complete, then read final bytes. Magnitude tolerates the same race (~0.01% residual variance at 5 sig figs) because mag L2 norms are ~300× larger than dir, so a handful of partially-updated bytes contribute a negligible fraction of the sum; dir's small magnitude makes it proportionally more sensitive. Fix: add `self.stream.synchronize()` immediately after the `memcpy_dtoh` call and before the host-side loop. Zero cost in practice — epoch-boundary readback already runs outside the training graph capture region. Validation (3× sequential smoke runs on this HEAD, `magnitude_distribution` test, fold 0 HEALTH_DIAG[0] grad_abs values): baseline (pre-fix) after fix ──────────────────── ──────────────────── dir=1.125252e-2 dir=1.373147e-2 dir=1.129484e-2 dir=1.376011e-2 dir=1.152756e-2 dir=1.376298e-2 mag=3.464567e0 mag=3.464277e0 mag=3.464559e0 mag=3.464551e0 mag=3.464618e0 mag=3.464625e0 dir spread: 2.4% → 0.2% (12× improvement, dir now aligned with mag) mag spread: 0.002% → 0.01% (unchanged noise floor) Direction aggregate now matches the mathematically-expected value computed from per-tensor norms (sqrt(Σ tᵢ²) ≈ 1.374e-2), confirming the readback bug was inflating the reported values below the true gradient norm by up to ~32% under the previous race. Logs: /tmp/foxhunt_smoke/dir_fix_run{1,2,3}.log Co-Authored-By: Claude Opus 4.7 (1M context) --- .../ml/src/cuda_pipeline/gpu_dqn_trainer.rs | 22 +++++++++++++++++++ 1 file changed, 22 insertions(+) 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];