test(crt-train): finite-diff validate output_smoothness kernel grad
This commit is contained in:
jgrusewski
254
crates/ml-alpha/tests/output_smoothness_grad_finite_diff.rs
Normal file
254
crates/ml-alpha/tests/output_smoothness_grad_finite_diff.rs
Normal file
@@ -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<f32>,
|
||||
grad: Vec<f32>,
|
||||
}
|
||||
|
||||
/// 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<SmoothOut> {
|
||||
let total = k * b * N_HORIZONS;
|
||||
assert_eq!(probs.len(), total);
|
||||
assert_eq!(grad_init.len(), total);
|
||||
|
||||
let stream: &Arc<CudaStream> = 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::<f32>(1).context("loss alloc")?;
|
||||
let mut raw_d = stream.alloc_zeros::<f32>(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<CudaStream>, host: &[f32]) -> Result<CudaSlice<f32>> {
|
||||
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::<f32>(n).context("smooth upload alloc")?;
|
||||
if n > 0 {
|
||||
let nbytes = n * std::mem::size_of::<f32>();
|
||||
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<CudaStream>, src: &CudaSlice<f32>) -> Result<Vec<f32>> {
|
||||
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::<f32>();
|
||||
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<f32> = (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<f32> = (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<usize> = 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<f32> = (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.
|
||||
}
|
||||
Reference in New Issue
Block a user