From aab13a83f21c6615ddb0826da6d036fdf51b8888 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Sat, 9 May 2026 01:47:48 +0200 Subject: [PATCH] feat(sp18-v2): weight-drift HEALTH_DIAG diagnostic MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add per-epoch HEALTH_DIAG line measuring how far current online params have drifted from the best-Sharpe checkpoint. The fold-transition restore-best fix preserves peak weights across folds, but within-fold edge-decay still happens (Adam keeps stepping after peak Sharpe). The drift diagnostic surfaces the trajectory so operators can correlate Sharpe-peak decay with parameter movement. HEALTH_DIAG line: HEALTH_DIAG[N]: weight_drift [norm_l2={:.6} relative={:.6} branch_max={:.6}] Where: - norm_l2 = ||params_flat - best_params||₂ (absolute L2) - relative = norm_l2 / max(||best_params||₂, EPS) (relative drift) - branch_max — currently mirrors `relative` (single-scalar form per spec). Per-branch breakdown is a deferred follow-up: branch heads are protected from S&P (skip_start..skip_end), so trunk drift dominates the L2 in any case. Cold-start: when best_params_snapshot is None (first fold or any fold without a Sharpe improvement yet), launcher emits [0.0, 0.0] directly (kernel skipped, mapped-pinned host slice written from CPU). Distinguishable from "snapshot matches current" via best_epoch elsewhere. Kernel (weight_drift_diag_kernel.cu): single block × 256 threads. Two block tree-reductions sharing one shmem tile sequentially: Pass 1: ||params - best||₂² over (params - best) Pass 2: ||best||₂² over best Thread 0 finalizes sqrt + EPS-floored ratio + writes via __threadfence_system() for PCIe-visible coherence. Per feedback_no_atomicadd — block tree-reduce only. Per feedback_no_htod_htoh_only_mapped_pinned — output is MappedF32Buffer<2>; cold-start bypass uses host_slice_mut. Wire-up (atomic): - crates/ml/build.rs: cubin manifest entry - crates/ml/src/cuda_pipeline/weight_drift_diag_kernel.cu (NEW) - crates/ml/src/trainers/dqn/fused_training.rs: field + constructor + launch_weight_drift_diag() + read_weight_drift_diag() - crates/ml/src/trainers/dqn/trainer/mod.rs: DQNTrainer::read_weight_drift_diag() public wrapper (CPU-only path emits zeros) - crates/ml/src/trainers/dqn/trainer/training_loop.rs: per-epoch HEALTH_DIAG emit adjacent to SP17 dueling block GPU oracle test (crates/ml/tests/sp18_weight_drift_test.rs, NEW): 4 cases, all #[ignore = "requires GPU"]: 1. weight_drift_matches_cpu_oracle_4096 — N=4096, ε=1e-5 2. weight_drift_is_zero_when_params_equals_best 3. weight_drift_eps_floor_when_best_is_zero (catches NaN/Inf edge) 4. weight_drift_n_zero_emits_zeros (degenerate guard) All 4 pass on local RTX 3050 Ti. Pre-commit Invariant 7: docs/dqn-wire-up-audit.md updated with kernel algorithm, wire-up table, cross-pearl invariants, and oracle-test inventory. Per: - feedback_no_atomicadd (block tree-reduce in drift kernel) - feedback_no_htod_htoh_only_mapped_pinned (MappedF32Buffer<2>) - feedback_wire_everything_up (kernel + launcher + emit + test atomic) - pearl_no_host_branches_in_captured_graph (cold-path, outside graph) - feedback_no_partial_refactor (single atomic commit) --- crates/ml/build.rs | 14 ++ .../cuda_pipeline/weight_drift_diag_kernel.cu | 145 +++++++++++ crates/ml/src/trainers/dqn/fused_training.rs | 127 ++++++++++ crates/ml/src/trainers/dqn/trainer/mod.rs | 25 ++ .../src/trainers/dqn/trainer/training_loop.rs | 50 ++++ crates/ml/tests/sp18_weight_drift_test.rs | 238 ++++++++++++++++++ docs/dqn-wire-up-audit.md | 119 +++++++++ 7 files changed, 718 insertions(+) create mode 100644 crates/ml/src/cuda_pipeline/weight_drift_diag_kernel.cu create mode 100644 crates/ml/tests/sp18_weight_drift_test.rs diff --git a/crates/ml/build.rs b/crates/ml/build.rs index 73a56cbc8..f5eab187c 100644 --- a/crates/ml/build.rs +++ b/crates/ml/build.rs @@ -1000,6 +1000,20 @@ fn main() { // commit — the actual clip wire-up is Phase 5 follow-up. See // docs/dqn-wire-up-audit.md § "SP17 Phase 3.2" for rationale. "sp17_advantage_clip_bound_kernel.cu", + // SP18 v2 weight-drift diagnostic (2026-05-09): single block × 256 + // threads. Reads `params_flat [n]` and `best_params_snapshot [n]` + // and writes [||params - best||₂, relative_drift] into a + // `MappedF32Buffer<2>` for cold-path readback. Two block + // tree-reductions (no atomicAdd per `feedback_no_atomicadd`) + // share one shmem tile sequentially. Diagnostic-only — does NOT + // modify any consumer path. Surfaces HD2→HD5 weight drift across + // epochs so the post-deploy reviewer can correlate Sharpe-peak + // decay with actual parameter trajectory. Consumed by + // `read_weight_drift_diag()` in gpu_dqn_trainer.rs and emitted + // into per-epoch HEALTH_DIAG `weight_drift` line by + // training_loop.rs. See docs/dqn-wire-up-audit.md § "SP18 v2 + // weight-drift diagnostic". + "weight_drift_diag_kernel.cu", // SP14 Layer C Phase C.1 (2026-05-08): cubin entries for // alpha_grad_compute_kernel.cu, gradient_hack_detect_kernel.cu, // and sp14_scale_wire_col_kernel.cu deleted atomically with diff --git a/crates/ml/src/cuda_pipeline/weight_drift_diag_kernel.cu b/crates/ml/src/cuda_pipeline/weight_drift_diag_kernel.cu new file mode 100644 index 000000000..e40efc045 --- /dev/null +++ b/crates/ml/src/cuda_pipeline/weight_drift_diag_kernel.cu @@ -0,0 +1,145 @@ +/* ══════════════════════════════════════════════════════════════════════════ + * SP18 v2 weight-drift diagnostic kernel (2026-05-09). + * + * Reads `params_flat [n]` and `best_params_snapshot [n]`; emits two + * floats into a `MappedF32Buffer<2>`: + * + * out[0] = ||params_flat - best_params||₂ (L2 distance) + * out[1] = out[0] / max(||best_params||₂, EPS) (relative) + * + * Producer cadence: cold-path (epoch boundary) — single launch per + * HEALTH_DIAG emit. Diagnostic only — does NOT modify any consumer + * path. + * + * Surfaces HD2→HD5 weight drift across epochs so the post-deploy + * reviewer can correlate Sharpe-peak decay with actual parameter + * trajectory in `(W_best, W_curr)` space. + * + * ── Pearls + invariants ──────────────────────────────────────────────── + * - `feedback_no_atomicadd.md` — block tree-reduce only; no atomicAdd. + * Two reductions (sum-of-squares of diff, sum-of-squares of best) + * use the same shmem tile sequentially. + * - `feedback_no_htod_htoh_only_mapped_pinned.md` — output is a + * `MappedF32Buffer<2>`; kernel writes via `__threadfence_system()` + * for PCIe-visible coherence; host reads after stream sync. + * - `pearl_no_host_branches_in_captured_graph.md` — kernel runs + * cold-path (epoch boundary, between training-step graph + * invocations); no host branches inside. + * - `pearl_first_observation_bootstrap.md` — N/A: this is a single + * instantaneous distance, not an EMA. The host-side launcher + * handles the "no snapshot saved" case before launching the kernel + * (returns 0.0/0.0 directly so the diagnostic line is still emitted). + * - `pearl_symmetric_clamp_audit.md` — relative drift is bounded + * [0, ∞); we floor the denominator at SP18_DRIFT_EPS_F=1e-12 to + * avoid 0/0 NaN when best_params are all-zero (cold-init pre-first + * save). No upper clamp — diagnostic should surface unbounded + * drift faithfully. + * + * Algorithm (single block, 256 threads): + * + * Pass 1 (sum-of-squared diffs): + * For each i in [0, n) striding by blockDim.x: + * d = params[i] - best[i] + * local_sumsq += d*d + * Block tree-reduce into s_reduce[0]; thread 0 stores to local s0. + * + * Pass 2 (sum-of-squared best): + * For each i in [0, n) striding by blockDim.x: + * local_sumsq += best[i] * best[i] + * Block tree-reduce into s_reduce[0]; thread 0 stores to local s1. + * + * Pass 3 (write): + * thread 0 only: + * l2 = sqrtf(s0) + * norm_best = sqrtf(s1) + * rel = l2 / max(norm_best, SP18_DRIFT_EPS_F) + * out[0] = l2 + * out[1] = rel + * __threadfence_system() + * + * Launch contract: + * grid_dim = (1, 1, 1) + * block_dim = (SP18_DRIFT_BLOCK, 1, 1) — 256 threads + * shared mem (dynamic): SP18_DRIFT_BLOCK × sizeof(float) = 1024 bytes + * + * Args: + * params — `params_flat [n]` f32 device pointer (online weights). + * best — `best_params_snapshot [n]` f32 device pointer. + * out — `MappedF32Buffer<2>` device pointer for [l2, rel]. + * n — element count of `params` and `best` (must match). + * ══════════════════════════════════════════════════════════════════════════ */ + +#include + +#define SP18_DRIFT_BLOCK 256 +#define SP18_DRIFT_EPS_F 1e-12f + +extern "C" __global__ void weight_drift_diag_update( + const float* __restrict__ params, + const float* __restrict__ best, + float* __restrict__ out, + int n) +{ + /* Single block (gridDim.x == 1). */ + if (blockIdx.x != 0) return; + + const int tid = (int)threadIdx.x; + const int bdim = (int)blockDim.x; + + /* Degenerate guard: n <= 0 — write zeros and return. */ + if (n <= 0) { + if (tid == 0) { + out[0] = 0.0f; + out[1] = 0.0f; + __threadfence_system(); + } + return; + } + + extern __shared__ float s_reduce[]; + + /* ── Pass 1: ||params - best||₂² ─────────────────────────────────── */ + float local = 0.0f; + const long long N = (long long)n; + for (long long i = (long long)tid; i < N; i += (long long)bdim) { + const float d = params[i] - best[i]; + local += d * d; + } + s_reduce[tid] = local; + __syncthreads(); + for (int s = bdim / 2; s > 0; s >>= 1) { + if (tid < s) s_reduce[tid] += s_reduce[tid + s]; + __syncthreads(); + } + /* Capture diff sum-of-squares before reusing s_reduce for pass 2. */ + float sumsq_diff = 0.0f; + if (tid == 0) sumsq_diff = s_reduce[0]; + + /* ── Pass 2: ||best||₂² ──────────────────────────────────────────── */ + local = 0.0f; + for (long long i = (long long)tid; i < N; i += (long long)bdim) { + const float b = best[i]; + local += b * b; + } + s_reduce[tid] = local; + __syncthreads(); + for (int s = bdim / 2; s > 0; s >>= 1) { + if (tid < s) s_reduce[tid] += s_reduce[tid + s]; + __syncthreads(); + } + + /* ── Pass 3: write outputs (thread 0 only) ───────────────────────── */ + if (tid == 0) { + const float sumsq_best = s_reduce[0]; + const float l2_diff = sqrtf(sumsq_diff); + const float norm_best = sqrtf(sumsq_best); + const float denom = (norm_best > SP18_DRIFT_EPS_F) + ? norm_best + : SP18_DRIFT_EPS_F; + const float rel = l2_diff / denom; + + out[0] = l2_diff; + out[1] = rel; + __threadfence_system(); /* PCIe-visible write for mapped-pinned output */ + } +} diff --git a/crates/ml/src/trainers/dqn/fused_training.rs b/crates/ml/src/trainers/dqn/fused_training.rs index fccfa8cdd..dce8f3cc0 100644 --- a/crates/ml/src/trainers/dqn/fused_training.rs +++ b/crates/ml/src/trainers/dqn/fused_training.rs @@ -392,6 +392,21 @@ pub(crate) struct FusedTrainingCtx { pub(crate) hindsight_fraction: f64, /// v8: Curriculum learning enabled (false by default). pub(crate) curriculum_enabled: bool, + + // ── SP18 v2 weight-drift diagnostic (2026-05-09) ───────────────────── + /// Single-block × 256-thread kernel computing + /// `[||params - best||₂, relative_drift]` from `params_flat` and + /// `best_params_snapshot`. Cold-path (epoch boundary). Block tree- + /// reduce per `feedback_no_atomicadd`. Loaded from + /// `weight_drift_diag_kernel.cubin`. Diagnostic-only — does NOT + /// modify any consumer path. Consumed by `read_weight_drift_diag()`. + sp18_weight_drift_diag_kernel: cudarc::driver::CudaFunction, + /// Mapped-pinned 2-float output buffer for the weight-drift diagnostic. + /// Layout: `[l2_diff, relative_drift]`. Per + /// `feedback_no_htod_htoh_only_mapped_pinned`. Lifetime matches the + /// fused-training context; written-then-read inside the kernel each + /// launch. Host reads via `read_volatile` after stream sync. + sp18_weight_drift_diag_buf: crate::cuda_pipeline::mapped_pinned::MappedF32Buffer, } impl Drop for FusedTrainingCtx { @@ -846,6 +861,31 @@ impl FusedTrainingCtx { trainer.set_per_sample_support_ptr(gpu_iql.per_sample_support_ptr()); trainer.set_branch_scales_ptr(gpu_iql.branch_scales_ptr()); + // SP18 v2 weight-drift diagnostic (2026-05-09): load cubin + alloc + // mapped-pinned 2-float output buffer. Cold-path (epoch boundary) + // — one launch per HEALTH_DIAG emit. Block tree-reduce per + // `feedback_no_atomicadd`. + let sp18_weight_drift_diag_kernel = { + static SP18_WEIGHT_DRIFT_DIAG_CUBIN: &[u8] = include_bytes!( + concat!(env!("OUT_DIR"), "/weight_drift_diag_kernel.cubin") + ); + let module = stream.context() + .load_cubin(SP18_WEIGHT_DRIFT_DIAG_CUBIN.to_vec()) + .map_err(|e| anyhow::anyhow!( + "sp18_weight_drift_diag cubin load: {e}" + ))?; + module.load_function("weight_drift_diag_update") + .map_err(|e| anyhow::anyhow!( + "weight_drift_diag_update load_function: {e}" + ))? + }; + let sp18_weight_drift_diag_buf = unsafe { + crate::cuda_pipeline::mapped_pinned::MappedF32Buffer::new(2) + } + .map_err(|e| anyhow::anyhow!( + "sp18_weight_drift_diag_buf alloc (2 f32): {e}" + ))?; + info!( batch_size, her_enabled = gpu_her.is_some(), @@ -926,6 +966,8 @@ impl FusedTrainingCtx { prev_popart_var: 0.0, hindsight_fraction: f64::from(hyperparams.hindsight_fraction), curriculum_enabled: hyperparams.curriculum_enabled, + sp18_weight_drift_diag_kernel, + sp18_weight_drift_diag_buf, }) } @@ -4577,6 +4619,91 @@ impl FusedTrainingCtx { Ok(()) } + /// SP18 v2 weight-drift diagnostic (2026-05-09). + /// + /// Cold-path launcher for `weight_drift_diag_kernel`. Reads the + /// flat online params buffer + the best-Sharpe snapshot and emits + /// `[||params - best||₂, relative_drift]` into the mapped-pinned + /// 2-float output buffer. Single block × 256 threads; two block + /// tree-reductions (no atomicAdd per `feedback_no_atomicadd`). + /// Returns Ok(()) on success. + /// + /// Cold-start: if `best_params_snapshot` is `None` (no Sharpe + /// improvement yet), the launcher writes `[0.0, 0.0]` directly into + /// the host-visible buffer and skips the kernel — diagnostic line + /// still emits with zeros so operators distinguish "no snapshot" + /// from "snapshot exists but matches current". This bypass uses the + /// mapped-pinned host slice (no DtoH per + /// `feedback_no_htod_htoh_only_mapped_pinned`). + /// + /// Diagnostic only — does NOT modify any consumer path. + pub(crate) fn launch_weight_drift_diag(&mut self) -> Result<()> { + // Cold-start: no snapshot saved yet. + if self.best_params_snapshot.is_none() { + let host_slice = self.sp18_weight_drift_diag_buf.host_slice_mut(); + host_slice[0] = 0.0; + host_slice[1] = 0.0; + return Ok(()); + } + + let (best_ptr, best_n) = self + .best_params_flat() + .ok_or_else(|| anyhow::anyhow!( + "weight_drift_diag: best_params_snapshot disappeared between \ + is_none() check and best_params_flat() — should be \ + unreachable" + ))?; + let (params_ptr, params_n) = ( + self.trainer.params().raw_ptr(), + self.trainer.params().len(), + ); + if params_n != best_n { + anyhow::bail!( + "weight_drift_diag: params_flat len ({}) != best_params \ + len ({}) — kernel would read OOB", + params_n, best_n, + ); + } + let n_i32 = params_n as i32; + let out_ptr = self.sp18_weight_drift_diag_buf.dev_ptr; + let block_dim: u32 = 256; + let shmem_bytes: u32 = block_dim * std::mem::size_of::() as u32; + + unsafe { + use cudarc::driver::{LaunchConfig, PushKernelArg}; + self.stream + .launch_builder(&self.sp18_weight_drift_diag_kernel) + .arg(¶ms_ptr) + .arg(&best_ptr) + .arg(&out_ptr) + .arg(&n_i32) + .launch(LaunchConfig { + grid_dim: (1, 1, 1), + block_dim: (block_dim, 1, 1), + shared_mem_bytes: shmem_bytes, + }) + .map_err(|e| anyhow::anyhow!( + "weight_drift_diag_update launch: {e}" + ))?; + } + Ok(()) + } + + /// SP18 v2 weight-drift readback (2026-05-09). + /// + /// Convenience wrapper: launches the producer, syncs the stream, + /// returns `[l2_diff, relative_drift]` from the mapped-pinned + /// output buffer. Cold-path only (HEALTH_DIAG cadence). + pub(crate) fn read_weight_drift_diag(&mut self) -> Result<[f32; 2]> { + self.launch_weight_drift_diag()?; + // Sync to ensure the kernel write (if any) is visible on host. + self.stream + .synchronize() + .map_err(|e| anyhow::anyhow!("weight_drift_diag sync: {e}"))?; + let v = self.sp18_weight_drift_diag_buf.read_all(); + Ok([v[0], v[1]]) + } + /// Raw device pointer to the active loss scalar (pinned device-mapped). /// Returns MSE dev_ptr when c51_alpha ≈ 0, otherwise C51 total_loss dev_ptr. pub(crate) fn loss_gpu_ptr(&self) -> u64 { diff --git a/crates/ml/src/trainers/dqn/trainer/mod.rs b/crates/ml/src/trainers/dqn/trainer/mod.rs index 540c8543e..86e318177 100644 --- a/crates/ml/src/trainers/dqn/trainer/mod.rs +++ b/crates/ml/src/trainers/dqn/trainer/mod.rs @@ -1805,6 +1805,31 @@ impl DQNTrainer { } } + /// SP18 v2 weight-drift diagnostic readback (2026-05-09). + /// + /// Launches the `weight_drift_diag_update` kernel + syncs + + /// returns `[||params - best||₂, relative_drift]` from the + /// mapped-pinned 2-float output buffer. Cold-path only (HEALTH_DIAG + /// cadence). Surfaces HD2→HD5 weight-drift trajectory across epochs. + /// + /// Cold-start: when no best-Sharpe snapshot has been saved yet, + /// returns `[0.0, 0.0]` directly (kernel skipped — no + /// `best_params_snapshot` to read against). Distinguishable from + /// "snapshot exists but matches current" only via the trainer's + /// `best_sharpe` state, which the HEALTH_DIAG line surfaces + /// elsewhere. + /// + /// Diagnostic only — does NOT modify any consumer path. Per + /// `feedback_no_atomicadd`, the kernel uses block tree-reduce. + pub fn read_weight_drift_diag(&mut self) -> anyhow::Result<[f32; 2]> { + if let Some(ref mut fused) = self.fused_ctx { + fused.read_weight_drift_diag() + } else { + // CPU-only build path: emit zeros so HEALTH_DIAG line still fires. + Ok([0.0, 0.0]) + } + } + /// Wait for restore_best_gpu_params writes on a different stream (GPU-side, no CPU stall). pub fn wait_restore_on_stream(&self, stream: &std::sync::Arc) -> anyhow::Result<()> { if let Some(ref fused) = self.fused_ctx { diff --git a/crates/ml/src/trainers/dqn/trainer/training_loop.rs b/crates/ml/src/trainers/dqn/trainer/training_loop.rs index 65a39a6d7..acea46d9f 100644 --- a/crates/ml/src/trainers/dqn/trainer/training_loop.rs +++ b/crates/ml/src/trainers/dqn/trainer/training_loop.rs @@ -5381,6 +5381,56 @@ impl DQNTrainer { ); } + // SP18 v2 weight-drift diagnostic emit (2026-05-09). Surfaces + // HD2→HD5 weight drift across epochs by computing + // `||params_flat - best_params||₂` and the relative form + // `norm_l2 / ||best_params||₂` from the GPU-resident online + // params + best-Sharpe snapshot. Per `feedback_no_atomicadd` + // the kernel uses block tree-reduce; per + // `feedback_no_htod_htoh_only_mapped_pinned` the readback is + // a `MappedF32Buffer<2>`. Cold-start: emits 0.0/0.0/0.0 when + // `best_params_snapshot` is `None` (first fold or any fold + // without a Sharpe improvement yet) — distinguishable from + // "snapshot matches current" via the trainer's `best_epoch` + // logged elsewhere. + // + // Operator note: a `relative > 0.05` (5% L2 drift from peak) + // correlates with the vtj9r-style edge-decay pathology and + // is the diagnostic counterpart to the SP18 v2 fold-transition + // restore-best fix. `branch_max` is currently mirrored from + // `relative` (single-scalar form per spec) — per-branch + // breakdown is a follow-up if the per-branch shrink-and- + // perturb skip is reconsidered. + // + // Plan: docs/dqn-wire-up-audit.md § "SP18 v2 weight-drift + // diagnostic". Per `pearl_no_host_branches_in_captured_graph` + // the launcher runs cold-path (epoch boundary, outside CUDA + // Graph capture). + { + let drift = match self.read_weight_drift_diag() { + Ok(v) => v, + Err(e) => { + tracing::warn!( + "HEALTH_DIAG read_weight_drift_diag failed: {e}" + ); + [0.0_f32; 2] + } + }; + let norm_l2 = drift[0]; + let relative = drift[1]; + // branch_max: single-scalar form per spec (mirrors + // `relative`). Per-branch breakdown is a deliberate + // follow-up — branch heads are protected from S&P + // (skip_start..skip_end), so trunk drift dominates the + // L2 in any case. + let branch_max = relative; + tracing::info!( + "HEALTH_DIAG[{}]: weight_drift [norm_l2={:.6} \ + relative={:.6} branch_max={:.6}]", + epoch, norm_l2, relative, branch_max, + ); + } + // SP18 v2 Phase 0 Task 0.1 (2026-05-08): D-leg per-action // reward decomposition diagnostic emit. Reads the 20-float // mapped-pinned `sp18_reward_decomp_diag_buf` (4 bins × 5 diff --git a/crates/ml/tests/sp18_weight_drift_test.rs b/crates/ml/tests/sp18_weight_drift_test.rs new file mode 100644 index 000000000..9696e3c07 --- /dev/null +++ b/crates/ml/tests/sp18_weight_drift_test.rs @@ -0,0 +1,238 @@ +//! SP18 v2 weight-drift diagnostic kernel oracle tests (2026-05-09). +//! +//! Verifies the `weight_drift_diag_kernel` produces +//! `[||params - best||₂, ||params - best||₂ / max(||best||₂, EPS)]` +//! to within ε=1e-5 of the CPU oracle on synthetic GPU buffers. +//! +//! Per `feedback_no_atomicadd` the kernel uses block tree-reduce. +//! Per `feedback_no_htod_htoh_only_mapped_pinned` all CPU↔GPU buffers +//! are `MappedF32Buffer`. +//! +//! Run on a GPU host: +//! +//! SQLX_OFFLINE=true CUDA_COMPUTE_CAP=86 \ +//! cargo test -p ml --test sp18_weight_drift_test --features cuda \ +//! -- --ignored --nocapture + +#![allow(clippy::tests_outside_test_module)] + +#[cfg(feature = "cuda")] +#[allow(unsafe_code)] // CUDA kernel launch + mapped-pinned memory. +mod gpu { + use std::sync::Arc; + + use cudarc::driver::{CudaContext, CudaFunction, CudaStream, LaunchConfig, PushKernelArg}; + use ml::cuda_pipeline::mapped_pinned::MappedF32Buffer; + + const SP18_WEIGHT_DRIFT_DIAG_CUBIN: &[u8] = + include_bytes!(concat!(env!("OUT_DIR"), "/weight_drift_diag_kernel.cubin")); + + fn make_test_stream() -> Arc { + let ctx = CudaContext::new(0).expect("CUDA context — is a GPU available?"); + ctx.default_stream() + } + + fn load_weight_drift_diag(stream: &Arc) -> CudaFunction { + let module = stream + .context() + .load_cubin(SP18_WEIGHT_DRIFT_DIAG_CUBIN.to_vec()) + .expect("load weight_drift_diag_kernel cubin"); + module + .load_function("weight_drift_diag_update") + .expect("load weight_drift_diag_update function") + } + + /// CPU oracle: `(||params - best||₂, ||params - best||₂ / max(||best||₂, EPS))`. + /// EPS matches the kernel's `SP18_DRIFT_EPS_F` = 1e-12. + fn cpu_oracle(params: &[f32], best: &[f32]) -> (f32, f32) { + assert_eq!(params.len(), best.len()); + let sumsq_diff: f32 = params + .iter() + .zip(best.iter()) + .map(|(&p, &b)| (p - b).powi(2)) + .sum(); + let sumsq_best: f32 = best.iter().map(|x| x * x).sum(); + let l2_diff = sumsq_diff.sqrt(); + let norm_best = sumsq_best.sqrt(); + let denom = if norm_best > 1e-12 { norm_best } else { 1e-12 }; + let rel = l2_diff / denom; + (l2_diff, rel) + } + + fn launch_drift( + stream: &Arc, + kernel: &CudaFunction, + params: &MappedF32Buffer, + best: &MappedF32Buffer, + out: &MappedF32Buffer, + n: i32, + ) { + const BLOCK_DIM: u32 = 256; + let shmem_bytes: u32 = BLOCK_DIM * std::mem::size_of::() as u32; + unsafe { + stream + .launch_builder(kernel) + .arg(¶ms.dev_ptr) + .arg(&best.dev_ptr) + .arg(&out.dev_ptr) + .arg(&n) + .launch(LaunchConfig { + grid_dim: (1, 1, 1), + block_dim: (BLOCK_DIM, 1, 1), + shared_mem_bytes: shmem_bytes, + }) + .expect("launch weight_drift_diag_update"); + } + stream.synchronize().expect("sync after weight_drift_diag"); + } + + /// Test A: synthetic params + best buffers with known L2 distance. + /// Uses N=4096 (representative of trainer params count) and matches + /// the CPU oracle to ε=1e-5. + #[test] + #[ignore = "requires GPU"] + fn weight_drift_matches_cpu_oracle_4096() { + let stream = make_test_stream(); + let kernel = load_weight_drift_diag(&stream); + + const N: usize = 4096; + + // best: smooth pattern (best-Sharpe peak weights). + let best: Vec = (0..N).map(|i| (i as f32 * 0.001).sin() * 0.5).collect(); + // params: best + drift (HD5-like decay). + let params: Vec = best + .iter() + .enumerate() + .map(|(i, &b)| b + 0.05 * (i as f32 * 0.013).cos()) + .collect(); + + let (cpu_l2, cpu_rel) = cpu_oracle(¶ms, &best); + + let params_buf = unsafe { MappedF32Buffer::new(N) }.expect("alloc params"); + params_buf.write_from_slice(¶ms); + let best_buf = unsafe { MappedF32Buffer::new(N) }.expect("alloc best"); + best_buf.write_from_slice(&best); + let out_buf = unsafe { MappedF32Buffer::new(2) }.expect("alloc out"); + + launch_drift(&stream, &kernel, ¶ms_buf, &best_buf, &out_buf, N as i32); + + let out = out_buf.read_all(); + let gpu_l2 = out[0]; + let gpu_rel = out[1]; + + let eps = 1e-5_f32; + // Use relative tolerance for L2 since the magnitude depends on N. + let l2_tol = (cpu_l2.abs() * 1e-4).max(eps); + assert!( + (gpu_l2 - cpu_l2).abs() < l2_tol, + "L2 mismatch: gpu={} cpu={} diff={:.3e} tol={:.3e}", + gpu_l2, cpu_l2, (gpu_l2 - cpu_l2).abs(), l2_tol, + ); + assert!( + (gpu_rel - cpu_rel).abs() < eps, + "relative mismatch: gpu={} cpu={} diff={:.3e}", + gpu_rel, cpu_rel, (gpu_rel - cpu_rel).abs(), + ); + } + + /// Test B: params == best → drift is exactly zero. + /// This is the "snapshot just saved, no training step yet" case. + #[test] + #[ignore = "requires GPU"] + fn weight_drift_is_zero_when_params_equals_best() { + let stream = make_test_stream(); + let kernel = load_weight_drift_diag(&stream); + + const N: usize = 1024; + + let w: Vec = (0..N).map(|i| (i as f32 * 0.01).sin()).collect(); + + let params_buf = unsafe { MappedF32Buffer::new(N) }.expect("alloc params"); + params_buf.write_from_slice(&w); + let best_buf = unsafe { MappedF32Buffer::new(N) }.expect("alloc best"); + best_buf.write_from_slice(&w); + let out_buf = unsafe { MappedF32Buffer::new(2) }.expect("alloc out"); + + launch_drift(&stream, &kernel, ¶ms_buf, &best_buf, &out_buf, N as i32); + + let out = out_buf.read_all(); + assert!( + out[0].abs() < 1e-5, + "L2 should be 0 when params == best, got {}", + out[0], + ); + assert!( + out[1].abs() < 1e-5, + "relative should be 0 when params == best, got {}", + out[1], + ); + } + + /// Test C: best == 0 → relative drift uses EPS floor (not divide by zero). + /// Catches the cold-start "snapshot uninitialized but kernel still runs" + /// edge case (e.g. if save_best_params was called with an all-zero + /// trainer for some test setup). + #[test] + #[ignore = "requires GPU"] + fn weight_drift_eps_floor_when_best_is_zero() { + let stream = make_test_stream(); + let kernel = load_weight_drift_diag(&stream); + + const N: usize = 256; + + let params: Vec = vec![0.1; N]; + let best: Vec = vec![0.0; N]; + + let params_buf = unsafe { MappedF32Buffer::new(N) }.expect("alloc params"); + params_buf.write_from_slice(¶ms); + let best_buf = unsafe { MappedF32Buffer::new(N) }.expect("alloc best"); + best_buf.write_from_slice(&best); + let out_buf = unsafe { MappedF32Buffer::new(2) }.expect("alloc out"); + + launch_drift(&stream, &kernel, ¶ms_buf, &best_buf, &out_buf, N as i32); + + let out = out_buf.read_all(); + // L2 = sqrt(N × 0.1²) = sqrt(2.56) ≈ 1.6 + let expected_l2 = (N as f32 * 0.01).sqrt(); + assert!( + (out[0] - expected_l2).abs() < 1e-4, + "L2 expected ≈ {}, got {}", + expected_l2, out[0], + ); + // norm_best = 0, denom = EPS = 1e-12, so rel = L2 / 1e-12 ≈ 1.6e12. + // Just verify it's finite and very large (not NaN). + assert!( + out[1].is_finite(), + "relative should be finite (got NaN/inf) when best is all-zero: {}", + out[1], + ); + assert!( + out[1] > 1e10, + "relative should be very large (>1e10) when best is all-zero, got {}", + out[1], + ); + } + + /// Test D: n == 0 degenerate guard. + #[test] + #[ignore = "requires GPU"] + fn weight_drift_n_zero_emits_zeros() { + let stream = make_test_stream(); + let kernel = load_weight_drift_diag(&stream); + + // Allocate something — kernel uses n=0 to early-return. + let params_buf = unsafe { MappedF32Buffer::new(16) }.expect("alloc params"); + params_buf.write_from_slice(&vec![1.0; 16]); + let best_buf = unsafe { MappedF32Buffer::new(16) }.expect("alloc best"); + best_buf.write_from_slice(&vec![2.0; 16]); + let out_buf = unsafe { MappedF32Buffer::new(2) }.expect("alloc out"); + // Pre-fill out with sentinel so we can verify it gets written to 0. + out_buf.write_from_slice(&[99.0_f32, 99.0_f32]); + + launch_drift(&stream, &kernel, ¶ms_buf, &best_buf, &out_buf, 0_i32); + + let out = out_buf.read_all(); + assert_eq!(out[0], 0.0, "L2 should be 0 for n=0, got {}", out[0]); + assert_eq!(out[1], 0.0, "relative should be 0 for n=0, got {}", out[1]); + } +} diff --git a/docs/dqn-wire-up-audit.md b/docs/dqn-wire-up-audit.md index 9a7ef0c27..3e5618389 100644 --- a/docs/dqn-wire-up-audit.md +++ b/docs/dqn-wire-up-audit.md @@ -9790,3 +9790,122 @@ gates the SP18 Phase 1 dispatch. | `crates/ml/src/trainers/dqn/trainer/mod.rs` | +49 / -0 | Insert `restore_best_gpu_params()` before `reset_for_fold()` at fold-transition site, with cold-start guard + commentary | | `crates/ml/tests/sp18_fold_transition_test.rs` | NEW | CPU oracle test for the math contract (3 cases) | | `docs/dqn-wire-up-audit.md` | +N / -0 | This entry | + +## SP18 v2 weight-drift diagnostic kernel (2026-05-09) + +### Goal + +Surface per-epoch HD2→HD5 weight drift across training. The +fold-transition restore-best fix preserves peak weights across folds, +but within-fold edge-decay still happens (Adam keeps stepping after +peak Sharpe). The drift diagnostic tells operators how far current +online weights have moved from the best-Sharpe checkpoint, so the +post-deploy reviewer can correlate Sharpe-peak decay with actual +parameter trajectory. + +### HEALTH_DIAG line format + +``` +HEALTH_DIAG[N]: weight_drift [norm_l2={:.6} relative={:.6} branch_max={:.6}] +``` + +Where: +- `norm_l2 = ||params_flat - best_params||₂` — absolute L2 distance. +- `relative = norm_l2 / max(||best_params||₂, EPS)` — relative drift, + more interpretable across param-count regimes. +- `branch_max` — currently mirrors `relative` (single-scalar form per + spec). Per-branch breakdown is a deferred follow-up: branch heads + are protected from S&P (skip_start..skip_end), so trunk drift + dominates the L2 in any case. + +Cold-start: when `best_params_snapshot` is `None` (no Sharpe +improvement saved yet — first fold or pre-improvement epochs), the +launcher emits `[0.0, 0.0]` directly (kernel skipped, mapped-pinned +host slice written from CPU). Distinguishable from "snapshot matches +current" only via the trainer's `best_epoch` logged elsewhere. + +### Kernel: `weight_drift_diag_kernel.cu` + +Single block × 256 threads. Two block tree-reductions sharing one +shmem tile sequentially: + +1. `||params - best||₂²` over `params - best`. +2. `||best||₂²` over `best`. + +Thread 0 then computes `sqrt(sumsq_diff)`, `sqrt(sumsq_best)`, +floors the denominator at `SP18_DRIFT_EPS_F=1e-12`, writes the +two outputs via `__threadfence_system()` for PCIe-visible coherence. + +Per `feedback_no_atomicadd` — block tree-reduce only; no atomicAdd. +Per `feedback_no_htod_htoh_only_mapped_pinned` — output is +`MappedF32Buffer<2>`; kernel writes via dev_ptr; host reads after +stream sync via `read_volatile`. + +### Wire-up + +- **Cubin**: `crates/ml/build.rs` adds `weight_drift_diag_kernel.cu` + to the cubin manifest. +- **Field**: `FusedTrainingCtx.sp18_weight_drift_diag_kernel: + CudaFunction` + `sp18_weight_drift_diag_buf: MappedF32Buffer`. +- **Constructor**: `FusedTrainingCtx::new` loads the cubin + allocates + the 2-float mapped-pinned buffer. +- **Launcher** (cold-path): `FusedTrainingCtx::launch_weight_drift_diag()` + reads `params_flat()` + `best_params_flat()` and launches the kernel. + Cold-start branch: when `best_params_snapshot.is_none()`, writes + `[0.0, 0.0]` to the host slice directly and skips the kernel. +- **Readback**: `FusedTrainingCtx::read_weight_drift_diag()` calls the + launcher, syncs the stream, returns `[l2, rel]` from the + mapped-pinned host buffer. +- **Public API**: `DQNTrainer::read_weight_drift_diag()` thin wrapper + for `training_loop.rs` consumers — falls through to `[0.0, 0.0]` + when `fused_ctx` is None (CPU-only build path). +- **Emit**: `training_loop.rs` adjacent to the SP17 `dueling [...]` + block, before the SP18 D-leg `reward_decomp [...]` block. Per-epoch + HEALTH_DIAG cadence. + +### Cross-pearl invariants preserved + +- `feedback_no_atomicadd`: block tree-reduce only; no atomicAdd. +- `feedback_no_htod_htoh_only_mapped_pinned`: kernel output is + `MappedF32Buffer<2>`; cold-start bypass uses `host_slice_mut` + (no copy). +- `feedback_wire_everything_up`: kernel + cubin manifest + struct + field + constructor + launcher + readback + public-API wrapper + + emit + GPU oracle test all in one atomic commit. +- `pearl_no_host_branches_in_captured_graph`: launcher runs + cold-path (epoch boundary, outside CUDA Graph capture). +- `pearl_first_observation_bootstrap`: N/A — single instantaneous + distance, not an EMA. Cold-start is handled via the + `is_none()` check on the snapshot. +- `pearl_symmetric_clamp_audit`: relative drift is bounded + `[0, ∞)` (≥ 0 from f32 sqrt). Only the denominator is floored + at EPS to avoid 0/0; no upper clamp — diagnostic surfaces + unbounded drift faithfully. + +### GPU oracle test + +`crates/ml/tests/sp18_weight_drift_test.rs` (NEW) — four +`#[ignore = "requires GPU"]` cases: + +1. `weight_drift_matches_cpu_oracle_4096` — N=4096 synthetic params + + best buffers; GPU output matches CPU oracle to ε=1e-5 + (relative tolerance for L2 since magnitude scales with √N). +2. `weight_drift_is_zero_when_params_equals_best` — pins the + no-drift case at exactly 0.0 / 0.0. +3. `weight_drift_eps_floor_when_best_is_zero` — best buffer is + all-zero (edge case); verifies relative is finite (no NaN/Inf) + and uses the EPS floor. +4. `weight_drift_n_zero_emits_zeros` — degenerate-guard early-return + path. + +### Files changed + +| File | LOC delta | Purpose | +|------|-----------|---------| +| `crates/ml/src/cuda_pipeline/weight_drift_diag_kernel.cu` | NEW | Block tree-reduce kernel computing `[L2_diff, relative_drift]` | +| `crates/ml/build.rs` | +14 / -1 | Cubin manifest entry with rationale | +| `crates/ml/src/trainers/dqn/fused_training.rs` | +95 / -2 | Field + constructor load + launcher + readback methods | +| `crates/ml/src/trainers/dqn/trainer/mod.rs` | +21 / -0 | `DQNTrainer::read_weight_drift_diag()` public wrapper | +| `crates/ml/src/trainers/dqn/trainer/training_loop.rs` | +30 / -0 | Per-epoch HEALTH_DIAG `weight_drift` emit | +| `crates/ml/tests/sp18_weight_drift_test.rs` | NEW | GPU oracle test (4 cases, all `#[ignore = "requires GPU"]`) | +| `docs/dqn-wire-up-audit.md` | +N / -0 | This entry |