feat(sp18-v2): weight-drift HEALTH_DIAG diagnostic
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
238
crates/ml/tests/sp18_weight_drift_test.rs
Normal file
238
crates/ml/tests/sp18_weight_drift_test.rs
Normal 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(¶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<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(¶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<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, ¶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<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(¶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]);
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user