From 3e07bf41b87a8ed82186f0e70a7f17c713fca4bf Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Wed, 20 May 2026 23:31:40 +0200 Subject: [PATCH] test(crt-train): finite-diff validate output_smoothness kernel grad --- .../output_smoothness_grad_finite_diff.rs | 254 ++++++++++++++++++ 1 file changed, 254 insertions(+) create mode 100644 crates/ml-alpha/tests/output_smoothness_grad_finite_diff.rs diff --git a/crates/ml-alpha/tests/output_smoothness_grad_finite_diff.rs b/crates/ml-alpha/tests/output_smoothness_grad_finite_diff.rs new file mode 100644 index 000000000..1e1038ade --- /dev/null +++ b/crates/ml-alpha/tests/output_smoothness_grad_finite_diff.rs @@ -0,0 +1,254 @@ +//! Output smoothness kernel — forward + gradient finite-diff validation. +//! +//! Per `feedback_no_cpu_test_fallbacks.md`: GPU oracle only. The analytic +//! gradient is verified against finite-difference of the kernel's own +//! forward pass (the kernel is the reference; we only check that its +//! analytic backward matches its own numerical derivative). +//! +//! Per `pearl_tests_must_prove_not_lock_observations.md`: assertions +//! verify INVARIANTS (loss ≥ 0 for non-zero diffs; loss == 0 when probs +//! are constant in k; finite-diff match) — not specific numeric outputs +//! that would lock the kernel to a single implementation. + +use std::sync::Arc; + +use anyhow::{Context, Result}; +use approx::assert_relative_eq; +use cudarc::driver::{CudaSlice, CudaStream, DevicePtr, DevicePtrMut, LaunchConfig, PushKernelArg}; +use ml_alpha::heads::N_HORIZONS; +use ml_core::device::MlDevice; + +const CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/output_smoothness.cubin")); + +struct SmoothOut { + loss: f32, + raw_per_h: Vec, + grad: Vec, +} + +/// Launch the smoothness kernel against `probs [K*B*5]`, returning +/// (total λ-weighted loss, per-horizon raw mean-sq-diff, accumulated grad). +/// `grad_init` is the value `grad_probs` is pre-loaded with BEFORE the +/// kernel runs (since the kernel adds +=). Pass zeros to verify the +/// pure-smoothness gradient; non-zero values verify additivity. +fn run_smoothness_kernel( + dev: &MlDevice, + probs: &[f32], + lambda_per_h: &[f32; N_HORIZONS], + k: usize, + b: usize, + grad_init: &[f32], +) -> Result { + let total = k * b * N_HORIZONS; + assert_eq!(probs.len(), total); + assert_eq!(grad_init.len(), total); + + let stream: &Arc = dev.cuda_stream().context("smooth stream")?; + let ctx = dev.cuda_context().context("smooth ctx")?; + let module = ctx.load_cubin(CUBIN.to_vec()).context("load output_smoothness cubin")?; + let func = module.load_function("output_smoothness_loss_and_grad").context("load fn")?; + + let probs_d = upload(stream, probs)?; + let lambda_d = upload(stream, lambda_per_h.as_slice())?; + let mut grad_d = upload(stream, grad_init)?; + let mut loss_d = stream.alloc_zeros::(1).context("loss alloc")?; + let mut raw_d = stream.alloc_zeros::(N_HORIZONS).context("raw alloc")?; + + let k_i = k as i32; + let b_i = b as i32; + let cfg = LaunchConfig { + grid_dim: (1, 1, 1), + block_dim: (256, 1, 1), + shared_mem_bytes: 0, + }; + let mut launch = stream.launch_builder(&func); + launch + .arg(&probs_d).arg(&lambda_d) + .arg(&k_i).arg(&b_i) + .arg(&mut loss_d).arg(&mut raw_d).arg(&mut grad_d); + unsafe { launch.launch(cfg).context("smoothness launch")?; } + stream.synchronize().context("smoothness sync")?; + + Ok(SmoothOut { + loss: download(stream, &loss_d)?[0], + raw_per_h: download(stream, &raw_d)?, + grad: download(stream, &grad_d)?, + }) +} + +fn upload(stream: &Arc, host: &[f32]) -> Result> { + use ml_alpha::pinned_mem::MappedF32Buffer; + let n = host.len(); + let staging = unsafe { MappedF32Buffer::new(n) } + .map_err(|e| anyhow::anyhow!("smooth upload staging: {e}"))?; + staging.write_from_slice(host); + let mut dst = stream.alloc_zeros::(n).context("smooth upload alloc")?; + if n > 0 { + let nbytes = n * std::mem::size_of::(); + unsafe { + let (dst_ptr, _g) = dst.device_ptr_mut(stream); + cudarc::driver::result::memcpy_dtod_async(dst_ptr, staging.dev_ptr, nbytes, stream.cu_stream()) + .context("smooth upload DtoD")?; + } + } + Ok(dst) +} + +fn download(stream: &Arc, src: &CudaSlice) -> Result> { + use ml_alpha::pinned_mem::MappedF32Buffer; + let n = src.len(); + let staging = unsafe { MappedF32Buffer::new(n) } + .map_err(|e| anyhow::anyhow!("smooth download staging: {e}"))?; + let nbytes = n * std::mem::size_of::(); + unsafe { + let (src_ptr, _g) = src.device_ptr(stream); + cudarc::driver::result::memcpy_dtod_async(staging.dev_ptr, src_ptr, nbytes, stream.cu_stream()) + .context("smooth download DtoD")?; + } + stream.synchronize().context("smooth download sync")?; + Ok(staging.read_all()) +} + +fn test_device() -> MlDevice { + MlDevice::cuda(0).expect("CUDA 0 required for ml-alpha tests") +} + +#[test] +fn smoothness_zero_when_probs_constant_in_k() { + let dev = test_device(); + let k = 4usize; + let b = 2usize; + // probs[k, b, h] = 0.3 + 0.1*h + 0.05*b (depends on b, h — NOT k) + let mut probs = vec![0.0; k * b * N_HORIZONS]; + for kk in 0..k { + for bb in 0..b { + for h in 0..N_HORIZONS { + probs[kk * b * N_HORIZONS + bb * N_HORIZONS + h] = + 0.3 + 0.1 * (h as f32) + 0.05 * (bb as f32); + } + } + } + let lambda = [0.5_f32; N_HORIZONS]; + let grad_init = vec![0.0_f32; k * b * N_HORIZONS]; + let out = run_smoothness_kernel(&dev, &probs, &lambda, k, b, &grad_init).unwrap(); + // Constant-in-k probs ⇒ raw_h[h] == 0 for all h ⇒ total loss == 0 and grad all zero. + for (h, &raw) in out.raw_per_h.iter().enumerate() { + assert!(raw.abs() < 1e-7, "raw_per_h[{h}] = {raw} (expected ~0 for k-constant probs)"); + } + assert!(out.loss.abs() < 1e-7, "loss = {} (expected ~0)", out.loss); + for (i, &g) in out.grad.iter().enumerate() { + assert!(g.abs() < 1e-7, "grad[{i}] = {g} (expected ~0)"); + } +} + +#[test] +fn smoothness_loss_positive_for_jittery_probs() { + let dev = test_device(); + let k = 6usize; + let b = 1usize; + // probs alternate sharply in k for every (b, h). + let mut probs = vec![0.0; k * b * N_HORIZONS]; + for kk in 0..k { + for h in 0..N_HORIZONS { + probs[kk * b * N_HORIZONS + h] = if kk % 2 == 0 { 0.2 } else { 0.8 }; + } + } + let lambda = [1.0_f32; N_HORIZONS]; + let grad_init = vec![0.0_f32; probs.len()]; + let out = run_smoothness_kernel(&dev, &probs, &lambda, k, b, &grad_init).unwrap(); + // Each (p[k]-p[k-1])² = 0.36 for K-1=5 pairs per horizon ⇒ raw_h = 0.36 for all h. + for (h, &raw) in out.raw_per_h.iter().enumerate() { + assert!(raw > 0.3 && raw < 0.4, "raw_per_h[{h}] = {raw} (expected ≈ 0.36 for alternating 0.2/0.8)"); + } + assert!(out.loss > 0.0, "loss should be positive for jittery probs, got {}", out.loss); +} + +#[test] +fn smoothness_grad_is_additive_on_top_of_existing() { + let dev = test_device(); + let k = 4usize; + let b = 1usize; + let probs: Vec = (0..k * b * N_HORIZONS).map(|i| 0.1 + 0.07 * (i as f32 % 7.0)).collect(); + let lambda = [0.3_f32; N_HORIZONS]; + + // Run once with zero init, once with a known non-zero init. + let zero_init = vec![0.0_f32; probs.len()]; + let out_zero = run_smoothness_kernel(&dev, &probs, &lambda, k, b, &zero_init).unwrap(); + + let nz_init: Vec = (0..probs.len()).map(|i| 0.5 + 0.01 * (i as f32)).collect(); + let out_nz = run_smoothness_kernel(&dev, &probs, &lambda, k, b, &nz_init).unwrap(); + + // out_nz.grad[i] should equal nz_init[i] + out_zero.grad[i]. + for i in 0..probs.len() { + let expected = nz_init[i] + out_zero.grad[i]; + assert_relative_eq!(out_nz.grad[i], expected, epsilon = 5e-5, max_relative = 5e-4); + } +} + +#[test] +fn smoothness_grad_matches_finite_difference() { + let dev = test_device(); + let k = 5usize; + let b = 2usize; + let mut probs = vec![0.0; k * b * N_HORIZONS]; + for i in 0..probs.len() { + // Vary in k so derivatives are non-trivial. + probs[i] = 0.2 + 0.05 * ((i % 7) as f32) + 0.03 * ((i / 7) as f32 % 5.0); + } + let lambda: [f32; N_HORIZONS] = [0.1, 0.3, 1.0, 3.0, 10.0]; // mimic the production ratio + let zero_init = vec![0.0_f32; probs.len()]; + let base = run_smoothness_kernel(&dev, &probs, &lambda, k, b, &zero_init).unwrap(); + + let eps = 1e-3_f32; + // Probe a representative set of indices: corner (k=0), interior (k=2), + // far corner (k=K-1), across two batches and horizons. + let probe_indices: Vec = vec![ + // (k=0, b=0, h=0) + 0, + // (k=0, b=1, h=2) + 0 * b * N_HORIZONS + 1 * N_HORIZONS + 2, + // (k=2, b=0, h=1) + 2 * b * N_HORIZONS + 0 * N_HORIZONS + 1, + // (k=2, b=1, h=4) + 2 * b * N_HORIZONS + 1 * N_HORIZONS + 4, + // (k=K-1, b=0, h=3) + (k - 1) * b * N_HORIZONS + 0 * N_HORIZONS + 3, + // (k=K-1, b=1, h=0) + (k - 1) * b * N_HORIZONS + 1 * N_HORIZONS + 0, + ]; + for &i in &probe_indices { + let mut up = probs.clone(); + up[i] += eps; + let mut dn = probs.clone(); + dn[i] -= eps; + let lu = run_smoothness_kernel(&dev, &up, &lambda, k, b, &zero_init).unwrap().loss; + let ld = run_smoothness_kernel(&dev, &dn, &lambda, k, b, &zero_init).unwrap().loss; + let fd = (lu - ld) / (2.0 * eps); + assert_relative_eq!( + base.grad[i], + fd, + epsilon = 5e-4, + max_relative = 5e-2 + ); + } +} + +#[test] +fn smoothness_zero_lambda_gives_zero_grad_and_loss() { + // Default `smoothness_base_lambda = 0.0` ⇒ λ[h] = 0 ∀h ⇒ kernel + // produces neither loss nor gradient. Verifies the opt-in default + // is exactly behaviour-preserving. + let dev = test_device(); + let k = 4usize; + let b = 2usize; + let probs: Vec = (0..k * b * N_HORIZONS).map(|i| 0.1 + 0.05 * (i as f32 % 9.0)).collect(); + let lambda = [0.0_f32; N_HORIZONS]; + let grad_init = vec![0.0_f32; probs.len()]; + let out = run_smoothness_kernel(&dev, &probs, &lambda, k, b, &grad_init).unwrap(); + assert!(out.loss.abs() < 1e-7); + for (i, &g) in out.grad.iter().enumerate() { + assert!(g.abs() < 1e-7, "grad[{i}] = {g} (expected 0 for λ=0)"); + } + // raw_per_h is still populated (telemetry — independent of λ). + // We don't assert specific values; only that the loss is zero. +}