gitadmin/foxhunt
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

feat(sp18-v2): weight-drift HEALTH_DIAG diagnostic

Browse Source 
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)
This commit is contained in:
jgrusewski
2026-05-09 01:47:48 +02:00
parent 3e9cefdbd9
commit aab13a83f2
7 changed files with 718 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

14
crates/ml/build.rs
View File

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

145
crates/ml/src/cuda_pipeline/weight_drift_diag_kernel.cu Normal file
View File

@@ -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 <cuda_runtime.h>
#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 */
}
}

127
crates/ml/src/trainers/dqn/fused_training.rs
View File

@@ -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::<f32>() as u32;
unsafe {
use cudarc::driver::{LaunchConfig, PushKernelArg};
self.stream
.launch_builder(&self.sp18_weight_drift_diag_kernel)
.arg(&params_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 {

25
crates/ml/src/trainers/dqn/trainer/mod.rs
View File

@@ -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<cudarc::driver::CudaStream>) -> anyhow::Result<()> {
if let Some(ref fused) = self.fused_ctx {

50
crates/ml/src/trainers/dqn/trainer/training_loop.rs
View File

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

238
crates/ml/tests/sp18_weight_drift_test.rs Normal file
View File

@@ -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<CudaStream> {
let ctx = CudaContext::new(0).expect("CUDA context — is a GPU available?");
ctx.default_stream()
}
fn load_weight_drift_diag(stream: &Arc<CudaStream>) -> 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<CudaStream>,
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::<f32>() as u32;
unsafe {
stream
.launch_builder(kernel)
.arg(&params.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<f32> = (0..N).map(|i| (i as f32 * 0.001).sin() * 0.5).collect();
// params: best + drift (HD5-like decay).
let params: Vec<f32> = best
.iter()
.enumerate()
.map(|(i, &b)| b + 0.05 * (i as f32 * 0.013).cos())
.collect();
let (cpu_l2, cpu_rel) = cpu_oracle(&params, &best);
let params_buf = unsafe { MappedF32Buffer::new(N) }.expect("alloc params");
params_buf.write_from_slice(&params);
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, &params_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<f32> = (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, &params_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<f32> = vec![0.1; N];
let best: Vec<f32> = vec![0.0; N];
let params_buf = unsafe { MappedF32Buffer::new(N) }.expect("alloc params");
params_buf.write_from_slice(&params);
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, &params_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, &params_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]);
}
}

119
docs/dqn-wire-up-audit.md
View File

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