fix(cuda): sync after async memcpy_dtoh in per_branch_grad_norms — direction-branch determinism

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) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-04-22 15:20:08 +02:00
parent 199feff4db
commit 5da434ab4b

View File

@@ -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];