test(rl): GPU oracle tests for IQN + NoisyNet + PRNG advance fix
Six GPU-oracle tests validating Phase 1 device-side PRNG kernels and the IQN/NoisyNet heads (mapped-pinned data transfer throughout): IQN: tau U(0,1) range + PRNG advance, forward finite, expected_q mean NoisyNet: factored transform invariant, zero-noise = mu, resample δ Fix: PRNG advance used prng_state[b] (zero on first call) instead of the self-seeded `seed` — xorshift32(0)=0 made second call identical. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
This commit is contained in:
jgrusewski
@@ -76,9 +76,10 @@ extern "C" __global__ void rl_sample_noise(
|
||||
eps_out[tid] = factored_noise_f(raw);
|
||||
}
|
||||
|
||||
// Thread 0 advances the persistent global seed.
|
||||
// Thread 0 advances the persistent global seed. Uses `seed`
|
||||
// (includes self-seed) not `prng_state[0]` (may be zero on first call).
|
||||
if (tid == 0) {
|
||||
uint32_t adv = prng_state[0];
|
||||
uint32_t adv = seed;
|
||||
#pragma unroll
|
||||
for (int w = 0; w < 8; ++w) xorshift32(&adv);
|
||||
prng_state[0] = adv;
|
||||
|
||||
@@ -58,9 +58,11 @@ extern "C" __global__ void rl_sample_tau(
|
||||
tau[b * n_tau + tid] = u;
|
||||
|
||||
// Thread 0 advances the persistent per-batch seed so the next
|
||||
// step produces different tau samples.
|
||||
// step produces different tau samples. Uses `seed` (which
|
||||
// includes the self-seed value) not `prng_state[b]` (which may
|
||||
// still be zero on the first call).
|
||||
if (tid == 0) {
|
||||
uint32_t adv = prng_state[b];
|
||||
uint32_t adv = seed;
|
||||
#pragma unroll
|
||||
for (int w = 0; w < 8; ++w) xorshift32(&adv);
|
||||
prng_state[b] = adv;
|
||||
|
||||
305
crates/ml-alpha/tests/iqn_oracle.rs
Normal file
305
crates/ml-alpha/tests/iqn_oracle.rs
Normal file
@@ -0,0 +1,305 @@
|
||||
//! GPU-oracle tests for the IQN head and the `rl_sample_tau` device-side
|
||||
//! PRNG kernel.
|
||||
//!
|
||||
//! Three analytical invariants, zero CPU reference implementations:
|
||||
//!
|
||||
//! 1. `rl_sample_tau_produces_uniform_01` — xorshift32 kernel outputs
|
||||
//! lie in [0, 1), and a second call with the same PRNG state buffer
|
||||
//! produces different values (seed advances).
|
||||
//! 2. `iqn_forward_output_shape` — `IqnHead::forward` produces finite
|
||||
//! Q-values of shape `[B, N_TAU, N_ACTIONS]` given valid h_t + tau.
|
||||
//! 3. `iqn_expected_q_is_mean_over_tau` — `IqnHead::expected_q` output
|
||||
//! equals the arithmetic mean of `q_values` over the tau dimension,
|
||||
//! within 1e-4. This is a definitional identity.
|
||||
//!
|
||||
//! Per `feedback_no_cpu_test_fallbacks`: GPU oracle only, no CPU
|
||||
//! reference implementations.
|
||||
//! Per `feedback_no_atomicadd`: no atomics.
|
||||
//!
|
||||
//! Run with:
|
||||
//! `cargo test -p ml-alpha --test iqn_oracle -- --ignored --nocapture`
|
||||
|
||||
use std::sync::Arc;
|
||||
|
||||
use cudarc::driver::{CudaSlice, CudaStream, DevicePtr, DevicePtrMut, LaunchConfig, PushKernelArg};
|
||||
use ml_alpha::heads::HIDDEN_DIM;
|
||||
use ml_alpha::pinned_mem::MappedF32Buffer;
|
||||
use ml_alpha::rl::common::N_ACTIONS;
|
||||
use ml_alpha::rl::iqn::{IqnHead, IqnHeadConfig};
|
||||
use ml_core::device::MlDevice;
|
||||
|
||||
const B: usize = 2;
|
||||
const N_TAU: usize = 32;
|
||||
|
||||
const RL_SAMPLE_TAU_CUBIN: &[u8] =
|
||||
include_bytes!(concat!(env!("OUT_DIR"), "/rl_sample_tau.cubin"));
|
||||
|
||||
fn get_device() -> Option<(MlDevice, Arc<CudaStream>)> {
|
||||
let dev = match MlDevice::cuda(0) {
|
||||
Ok(d) => d,
|
||||
Err(e) => {
|
||||
eprintln!("CUDA 0 not available — skipping ({e})");
|
||||
return None;
|
||||
}
|
||||
};
|
||||
let stream = dev.cuda_stream().expect("cuda_stream").clone();
|
||||
Some((dev, stream))
|
||||
}
|
||||
|
||||
fn upload_f32(stream: &Arc<CudaStream>, host: &[f32]) -> CudaSlice<f32> {
|
||||
let n = host.len();
|
||||
let staging = unsafe { MappedF32Buffer::new(n) }.expect("staging alloc");
|
||||
staging.write_from_slice(host);
|
||||
let mut d = stream.alloc_zeros::<f32>(n).expect("alloc f32");
|
||||
unsafe {
|
||||
let (dst, _g) = d.device_ptr_mut(stream);
|
||||
cudarc::driver::result::memcpy_dtod_async(
|
||||
dst, staging.dev_ptr, n * 4, stream.cu_stream(),
|
||||
).expect("upload DtoD");
|
||||
}
|
||||
stream.synchronize().expect("upload sync");
|
||||
d
|
||||
}
|
||||
|
||||
fn download_f32(stream: &Arc<CudaStream>, d: &CudaSlice<f32>, n: usize) -> Vec<f32> {
|
||||
let staging = unsafe { MappedF32Buffer::new(n) }.expect("staging alloc");
|
||||
unsafe {
|
||||
let (src, _g) = d.device_ptr(stream);
|
||||
cudarc::driver::result::memcpy_dtod_async(
|
||||
staging.dev_ptr, src, n * 4, stream.cu_stream(),
|
||||
).expect("download DtoD");
|
||||
}
|
||||
stream.synchronize().expect("download sync");
|
||||
let mut out = vec![0.0_f32; n];
|
||||
for i in 0..n {
|
||||
out[i] = unsafe { std::ptr::read_volatile(staging.host_ptr.add(i)) };
|
||||
}
|
||||
out
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Test 1: rl_sample_tau kernel produces U(0,1) and advances PRNG
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
#[test]
|
||||
#[ignore = "requires CUDA (MlDevice::cuda(0))"]
|
||||
fn rl_sample_tau_produces_uniform_01() {
|
||||
let Some((dev, stream)) = get_device() else { return };
|
||||
let ctx = dev.cuda_context().expect("cuda_context");
|
||||
|
||||
let module = ctx
|
||||
.load_cubin(RL_SAMPLE_TAU_CUBIN.to_vec())
|
||||
.expect("load rl_sample_tau cubin");
|
||||
let sample_tau_fn = module
|
||||
.load_function("rl_sample_tau")
|
||||
.expect("load rl_sample_tau fn");
|
||||
|
||||
// alloc_zeros self-seeds the xorshift32 on first call (seed==0 →
|
||||
// Knuth golden-ratio hash from batch index).
|
||||
let mut prng_state_d = stream
|
||||
.alloc_zeros::<u32>(B)
|
||||
.expect("alloc prng_state_d");
|
||||
let mut tau_d = stream
|
||||
.alloc_zeros::<f32>(B * N_TAU)
|
||||
.expect("alloc tau_d");
|
||||
|
||||
let b_i = B as i32;
|
||||
let n_tau_i = N_TAU as i32;
|
||||
let tau_cfg = LaunchConfig {
|
||||
grid_dim: (B as u32, 1, 1),
|
||||
block_dim: (N_TAU as u32, 1, 1),
|
||||
shared_mem_bytes: 0,
|
||||
};
|
||||
|
||||
// --- First launch ---
|
||||
unsafe {
|
||||
stream
|
||||
.launch_builder(&sample_tau_fn)
|
||||
.arg(&mut prng_state_d)
|
||||
.arg(&mut tau_d)
|
||||
.arg(&b_i)
|
||||
.arg(&n_tau_i)
|
||||
.launch(tau_cfg)
|
||||
.expect("rl_sample_tau launch 1");
|
||||
}
|
||||
stream.synchronize().expect("sync 1");
|
||||
|
||||
let tau_host_1 = download_f32(&stream, &tau_d, B * N_TAU);
|
||||
|
||||
for (i, &t) in tau_host_1.iter().enumerate() {
|
||||
assert!(
|
||||
t >= 0.0 && t < 1.0,
|
||||
"tau[{i}] = {t} outside [0, 1) on first call"
|
||||
);
|
||||
}
|
||||
|
||||
// --- Second launch (PRNG advances) ---
|
||||
unsafe {
|
||||
stream
|
||||
.launch_builder(&sample_tau_fn)
|
||||
.arg(&mut prng_state_d)
|
||||
.arg(&mut tau_d)
|
||||
.arg(&b_i)
|
||||
.arg(&n_tau_i)
|
||||
.launch(tau_cfg)
|
||||
.expect("rl_sample_tau launch 2");
|
||||
}
|
||||
stream.synchronize().expect("sync 2");
|
||||
|
||||
let tau_host_2 = download_f32(&stream, &tau_d, B * N_TAU);
|
||||
|
||||
for (i, &t) in tau_host_2.iter().enumerate() {
|
||||
assert!(
|
||||
t >= 0.0 && t < 1.0,
|
||||
"tau[{i}] = {t} outside [0, 1) on second call"
|
||||
);
|
||||
}
|
||||
|
||||
// At least one value must differ between calls — PRNG must advance.
|
||||
let any_differ = tau_host_1
|
||||
.iter()
|
||||
.zip(tau_host_2.iter())
|
||||
.any(|(a, b)| (a - b).abs() > f32::EPSILON);
|
||||
assert!(
|
||||
any_differ,
|
||||
"Second rl_sample_tau call produced identical values — PRNG did not advance"
|
||||
);
|
||||
|
||||
eprintln!(
|
||||
"IQN.1 OK — rl_sample_tau: all {} values in [0,1), second call differs",
|
||||
B * N_TAU
|
||||
);
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Test 2: IqnHead::forward produces finite Q-values of correct shape
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
#[test]
|
||||
#[ignore = "requires CUDA (MlDevice::cuda(0))"]
|
||||
fn iqn_forward_output_shape() {
|
||||
let Some((dev, stream)) = get_device() else { return };
|
||||
|
||||
let cfg = IqnHeadConfig {
|
||||
hidden_dim: HIDDEN_DIM,
|
||||
n_tau: N_TAU,
|
||||
..IqnHeadConfig::default()
|
||||
};
|
||||
let iqn = IqnHead::new(&dev, cfg).expect("IqnHead::new");
|
||||
|
||||
// h_t [B, HIDDEN_DIM] — small deterministic values
|
||||
let h_t_host: Vec<f32> = (0..B * HIDDEN_DIM)
|
||||
.map(|i| ((i as f32) * 0.001 - 0.064).sin() * 0.1)
|
||||
.collect();
|
||||
let h_t_d = upload_f32(&stream, &h_t_host);
|
||||
|
||||
// tau [B, N_TAU] — evenly spaced quantile fractions
|
||||
let tau_host: Vec<f32> = (0..B * N_TAU)
|
||||
.map(|i| {
|
||||
let within_row = i % N_TAU;
|
||||
(within_row as f32 + 0.5) / N_TAU as f32
|
||||
})
|
||||
.collect();
|
||||
let tau_d = upload_f32(&stream, &tau_host);
|
||||
|
||||
let mut q_values_d = stream
|
||||
.alloc_zeros::<f32>(B * N_TAU * N_ACTIONS)
|
||||
.expect("alloc q_values");
|
||||
|
||||
iqn.forward(&h_t_d, &tau_d, B, N_TAU, &mut q_values_d)
|
||||
.expect("IqnHead::forward");
|
||||
stream.synchronize().expect("sync forward");
|
||||
|
||||
let q_host = download_f32(&stream, &q_values_d, B * N_TAU * N_ACTIONS);
|
||||
|
||||
for (i, &q) in q_host.iter().enumerate() {
|
||||
assert!(
|
||||
q.is_finite(),
|
||||
"q_values[{i}] = {q} is not finite"
|
||||
);
|
||||
}
|
||||
|
||||
eprintln!(
|
||||
"IQN.2 OK — IqnHead::forward: all {} values finite (B={B}, N_TAU={N_TAU}, N_ACTIONS={N_ACTIONS})",
|
||||
q_host.len()
|
||||
);
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Test 3: expected_q == mean over tau dimension (analytical invariant)
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
#[test]
|
||||
#[ignore = "requires CUDA (MlDevice::cuda(0))"]
|
||||
fn iqn_expected_q_is_mean_over_tau() {
|
||||
let Some((dev, stream)) = get_device() else { return };
|
||||
|
||||
let cfg = IqnHeadConfig {
|
||||
hidden_dim: HIDDEN_DIM,
|
||||
n_tau: N_TAU,
|
||||
..IqnHeadConfig::default()
|
||||
};
|
||||
let iqn = IqnHead::new(&dev, cfg).expect("IqnHead::new");
|
||||
|
||||
// h_t [B, HIDDEN_DIM] — varied values to avoid degenerate Q output
|
||||
let h_t_host: Vec<f32> = (0..B * HIDDEN_DIM)
|
||||
.map(|i| ((i as f32) * 0.0073 + 1.0).sin() * 0.2)
|
||||
.collect();
|
||||
let h_t_d = upload_f32(&stream, &h_t_host);
|
||||
|
||||
// tau [B, N_TAU] — spread across the unit interval
|
||||
let tau_host: Vec<f32> = (0..B * N_TAU)
|
||||
.map(|i| {
|
||||
let within_row = i % N_TAU;
|
||||
(within_row as f32 + 0.5) / N_TAU as f32
|
||||
})
|
||||
.collect();
|
||||
let tau_d = upload_f32(&stream, &tau_host);
|
||||
|
||||
// Forward pass
|
||||
let mut q_values_d = stream
|
||||
.alloc_zeros::<f32>(B * N_TAU * N_ACTIONS)
|
||||
.expect("alloc q_values");
|
||||
iqn.forward(&h_t_d, &tau_d, B, N_TAU, &mut q_values_d)
|
||||
.expect("IqnHead::forward");
|
||||
|
||||
// Expected Q via the kernel
|
||||
let mut expected_q_d = stream
|
||||
.alloc_zeros::<f32>(B * N_ACTIONS)
|
||||
.expect("alloc expected_q");
|
||||
iqn.expected_q(&q_values_d, B, N_TAU, &mut expected_q_d)
|
||||
.expect("IqnHead::expected_q");
|
||||
stream.synchronize().expect("sync expected_q");
|
||||
|
||||
// Read back both
|
||||
let q_host = download_f32(&stream, &q_values_d, B * N_TAU * N_ACTIONS);
|
||||
let eq_host = download_f32(&stream, &expected_q_d, B * N_ACTIONS);
|
||||
|
||||
// Verify: expected_q[b, a] == mean_tau(q_values[b, :, a])
|
||||
// q_values layout: [B, N_TAU, N_ACTIONS] row-major
|
||||
let mut max_err: f32 = 0.0;
|
||||
for b in 0..B {
|
||||
for a in 0..N_ACTIONS {
|
||||
let mut sum = 0.0_f32;
|
||||
for t in 0..N_TAU {
|
||||
sum += q_host[b * N_TAU * N_ACTIONS + t * N_ACTIONS + a];
|
||||
}
|
||||
let mean = sum / N_TAU as f32;
|
||||
let kernel_val = eq_host[b * N_ACTIONS + a];
|
||||
let err = (mean - kernel_val).abs();
|
||||
if err > max_err {
|
||||
max_err = err;
|
||||
}
|
||||
assert!(
|
||||
err < 1e-4,
|
||||
"expected_q[{b},{a}] mismatch: kernel={kernel_val}, \
|
||||
mean_tau={mean}, err={err}"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
eprintln!(
|
||||
"IQN.3 OK — expected_q matches mean over tau for all (b,a) pairs \
|
||||
(max_err={max_err:.2e}, B={B}, N_TAU={N_TAU}, N_ACTIONS={N_ACTIONS})"
|
||||
);
|
||||
}
|
||||
292
crates/ml-alpha/tests/noisy_oracle.rs
Normal file
292
crates/ml-alpha/tests/noisy_oracle.rs
Normal file
@@ -0,0 +1,292 @@
|
||||
//! GPU oracle tests for `NoisyLinear` and the `rl_sample_noise` kernel.
|
||||
//!
|
||||
//! Three tests verify the factored noise generation, the zero-noise
|
||||
//! forward pass (mu-only), and that resampling noise changes the output.
|
||||
//!
|
||||
//! Per `feedback_no_cpu_test_fallbacks`: GPU oracle only, no CPU
|
||||
//! reference implementations. All assertions are analytical invariants
|
||||
//! derived from the kernel contracts.
|
||||
//!
|
||||
//! Per `feedback_no_atomicadd`: no atomics in any tested kernel path.
|
||||
//!
|
||||
//! Run with:
|
||||
//! `SQLX_OFFLINE=true cargo test -p ml-alpha --test noisy_oracle -- --ignored --nocapture`
|
||||
|
||||
#![allow(unsafe_code)] // CUDA kernel launch + mapped-pinned memory.
|
||||
|
||||
use cudarc::driver::{CudaStream, DevicePtr, DevicePtrMut, LaunchConfig, PushKernelArg};
|
||||
use ml_alpha::pinned_mem::MappedF32Buffer;
|
||||
use ml_alpha::rl::noisy::{NoisyLinear, NoisyLinearConfig};
|
||||
use ml_core::device::MlDevice;
|
||||
use std::sync::Arc;
|
||||
|
||||
const IN_DIM: usize = 128;
|
||||
const OUT_DIM: usize = 11;
|
||||
const B: usize = 2;
|
||||
|
||||
const SAMPLE_NOISE_CUBIN: &[u8] = include_bytes!(concat!(
|
||||
env!("OUT_DIR"),
|
||||
"/rl_sample_noise.cubin"
|
||||
));
|
||||
|
||||
fn make_device() -> Option<MlDevice> {
|
||||
match MlDevice::cuda(0) {
|
||||
Ok(d) => Some(d),
|
||||
Err(e) => {
|
||||
eprintln!("CUDA 0 not available — skipping ({e})");
|
||||
None
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn upload_f32(stream: &Arc<CudaStream>, host: &[f32]) -> cudarc::driver::CudaSlice<f32> {
|
||||
let n = host.len();
|
||||
let staging = unsafe { MappedF32Buffer::new(n) }.expect("staging alloc");
|
||||
staging.write_from_slice(host);
|
||||
let mut d = stream.alloc_zeros::<f32>(n).expect("alloc f32");
|
||||
unsafe {
|
||||
let (dst, _g) = d.device_ptr_mut(stream);
|
||||
cudarc::driver::result::memcpy_dtod_async(
|
||||
dst, staging.dev_ptr, n * 4, stream.cu_stream(),
|
||||
).expect("upload DtoD");
|
||||
}
|
||||
stream.synchronize().expect("upload sync");
|
||||
d
|
||||
}
|
||||
|
||||
fn download_f32(stream: &Arc<CudaStream>, d: &cudarc::driver::CudaSlice<f32>) -> Vec<f32> {
|
||||
let n = d.len();
|
||||
let staging = unsafe { MappedF32Buffer::new(n) }.expect("staging alloc");
|
||||
unsafe {
|
||||
let (src, _g) = d.device_ptr(stream);
|
||||
cudarc::driver::result::memcpy_dtod_async(
|
||||
staging.dev_ptr, src, n * 4, stream.cu_stream(),
|
||||
).expect("download DtoD");
|
||||
}
|
||||
stream.synchronize().expect("download sync");
|
||||
let mut out = vec![0.0_f32; n];
|
||||
for i in 0..n {
|
||||
out[i] = unsafe { std::ptr::read_volatile(staging.host_ptr.add(i)) };
|
||||
}
|
||||
out
|
||||
}
|
||||
|
||||
// ── Test 1: rl_sample_noise factored transform ─────────────────────────
|
||||
|
||||
/// Verify the `rl_sample_noise` kernel by launching it directly.
|
||||
///
|
||||
/// Invariants checked (all analytical — no CPU reference):
|
||||
/// 1. Every element of eps_in and eps_out is finite and nonzero
|
||||
/// (noise was actually generated, not left at zero).
|
||||
/// 2. The factored noise transform f(x) = sign(x) * sqrt(|x|) maps
|
||||
/// raw values from [-1, +1) into a range where, for any output
|
||||
/// `v = f(raw)`, the identity `v.signum() * v.powi(2)` recovers
|
||||
/// a value in [-1, +1). This holds because `|f(raw)| = sqrt(|raw|)`
|
||||
/// and `|raw| < 1` implies `|f(raw)| < 1`, so `f(raw)^2 = |raw| < 1`.
|
||||
#[test]
|
||||
#[ignore = "requires CUDA"]
|
||||
fn rl_sample_noise_factored_transform() {
|
||||
let Some(dev) = make_device() else { return };
|
||||
let stream = dev.cuda_stream().expect("cuda_stream").clone();
|
||||
let ctx = dev.cuda_context().expect("cuda_context");
|
||||
|
||||
// Load the cubin and extract the kernel function.
|
||||
let module = ctx
|
||||
.load_cubin(SAMPLE_NOISE_CUBIN.to_vec())
|
||||
.expect("load rl_sample_noise cubin");
|
||||
let kernel_fn = module
|
||||
.load_function("rl_sample_noise")
|
||||
.expect("load rl_sample_noise function");
|
||||
|
||||
// Allocate device buffers. prng_state starts at zero — the kernel
|
||||
// self-seeds on first call per the kernel contract.
|
||||
let mut prng_state_d = stream.alloc_zeros::<u32>(1).expect("alloc prng_state");
|
||||
let mut eps_in_d = stream.alloc_zeros::<f32>(IN_DIM).expect("alloc eps_in");
|
||||
let mut eps_out_d = stream.alloc_zeros::<f32>(OUT_DIM).expect("alloc eps_out");
|
||||
|
||||
let in_dim_i: i32 = IN_DIM as i32;
|
||||
let out_dim_i: i32 = OUT_DIM as i32;
|
||||
let max_dim = IN_DIM.max(OUT_DIM) as u32;
|
||||
|
||||
// Launch: Grid=(1,1,1), Block=(max(in_dim, out_dim), 1, 1).
|
||||
unsafe {
|
||||
stream
|
||||
.launch_builder(&kernel_fn)
|
||||
.arg(&mut prng_state_d)
|
||||
.arg(&mut eps_in_d)
|
||||
.arg(&mut eps_out_d)
|
||||
.arg(&in_dim_i)
|
||||
.arg(&out_dim_i)
|
||||
.launch(LaunchConfig {
|
||||
grid_dim: (1, 1, 1),
|
||||
block_dim: (max_dim, 1, 1),
|
||||
shared_mem_bytes: 0,
|
||||
})
|
||||
.expect("launch rl_sample_noise");
|
||||
}
|
||||
stream.synchronize().expect("sync after rl_sample_noise");
|
||||
|
||||
// Read back results.
|
||||
let eps_in = download_f32(&stream, &eps_in_d);
|
||||
let eps_out = download_f32(&stream, &eps_out_d);
|
||||
|
||||
// Invariant 1: ALL eps_in values are finite and nonzero.
|
||||
for (i, &v) in eps_in.iter().enumerate() {
|
||||
assert!(v.is_finite(), "eps_in[{i}] is not finite: {v}");
|
||||
assert!(v != 0.0, "eps_in[{i}] is zero — noise was not generated");
|
||||
}
|
||||
|
||||
// Invariant 1: ALL eps_out values are finite and nonzero.
|
||||
for (i, &v) in eps_out.iter().enumerate() {
|
||||
assert!(v.is_finite(), "eps_out[{i}] is not finite: {v}");
|
||||
assert!(v != 0.0, "eps_out[{i}] is zero — noise was not generated");
|
||||
}
|
||||
|
||||
// Invariant 2: factored noise transform property.
|
||||
// f(x) = sign(x) * sqrt(|x|) where x in [-1, +1).
|
||||
// So for any output v = f(x): v^2 = |x| < 1.0, meaning
|
||||
// |v.signum() * v.powi(2)| = |x| < 1.0.
|
||||
// The kernel applies f() to raw xorshift32 values in [-1, +1).
|
||||
for (i, &v) in eps_in.iter().enumerate() {
|
||||
let recovered = v.signum() * v.powi(2);
|
||||
assert!(
|
||||
recovered.abs() <= 1.0,
|
||||
"eps_in[{i}]: |sign(v)*v^2| = {} > 1.0 — factored noise invariant violated (v={v})",
|
||||
recovered.abs()
|
||||
);
|
||||
}
|
||||
for (i, &v) in eps_out.iter().enumerate() {
|
||||
let recovered = v.signum() * v.powi(2);
|
||||
assert!(
|
||||
recovered.abs() <= 1.0,
|
||||
"eps_out[{i}]: |sign(v)*v^2| = {} > 1.0 — factored noise invariant violated (v={v})",
|
||||
recovered.abs()
|
||||
);
|
||||
}
|
||||
|
||||
eprintln!(
|
||||
"noisy_oracle.1 OK — rl_sample_noise: {} eps_in + {} eps_out all finite, nonzero, factored-transform invariant holds",
|
||||
eps_in.len(),
|
||||
eps_out.len()
|
||||
);
|
||||
}
|
||||
|
||||
// ── Test 2: noisy_linear zero noise equals mu ──────────────────────────
|
||||
|
||||
/// With eps_in and eps_out at zero (no resample_noise call), the forward
|
||||
/// pass computes `y = mu_w * x + mu_b` (pure mean parameters, no noise
|
||||
/// contribution).
|
||||
///
|
||||
/// Invariants checked:
|
||||
/// 1. All output values are finite.
|
||||
/// 2. At least one output value is nonzero (mu_w and mu_b are Xavier-
|
||||
/// initialized, so the dot product with a non-zero input is nonzero).
|
||||
#[test]
|
||||
#[ignore = "requires CUDA"]
|
||||
fn noisy_linear_zero_noise_equals_mu() {
|
||||
let Some(dev) = make_device() else { return };
|
||||
let stream = dev.cuda_stream().expect("cuda_stream").clone();
|
||||
|
||||
let cfg = NoisyLinearConfig {
|
||||
in_dim: IN_DIM,
|
||||
out_dim: OUT_DIM,
|
||||
sigma_init: 0.5,
|
||||
seed: 0xDEAD,
|
||||
};
|
||||
let layer = NoisyLinear::new(&dev, cfg).expect("NoisyLinear::new");
|
||||
|
||||
// Input: B=2 batches of constant 0.1 values.
|
||||
let x_host = vec![0.1_f32; B * IN_DIM];
|
||||
let x_d = upload_f32(&stream, &x_host);
|
||||
let mut y_d = stream.alloc_zeros::<f32>(B * OUT_DIM).expect("alloc y");
|
||||
|
||||
// Forward WITHOUT resample_noise — eps buffers are still zero from
|
||||
// alloc_zeros at construction time.
|
||||
layer
|
||||
.forward(&x_d, B, &mut y_d)
|
||||
.expect("noisy_linear forward");
|
||||
stream.synchronize().expect("sync after forward");
|
||||
|
||||
let y = download_f32(&stream, &y_d);
|
||||
|
||||
// Invariant 1: all values are finite.
|
||||
for (i, &v) in y.iter().enumerate() {
|
||||
assert!(v.is_finite(), "y[{i}] is not finite: {v}");
|
||||
}
|
||||
|
||||
// Invariant 2: not all zeros — Xavier mu weights + constant input
|
||||
// produce nonzero output.
|
||||
let any_nonzero = y.iter().any(|&v| v != 0.0);
|
||||
assert!(
|
||||
any_nonzero,
|
||||
"all y values are zero — mu_w forward should produce nonzero output with x=0.1"
|
||||
);
|
||||
|
||||
eprintln!(
|
||||
"noisy_oracle.2 OK — zero-noise forward: all {} outputs finite, nonzero output confirmed",
|
||||
y.len()
|
||||
);
|
||||
}
|
||||
|
||||
// ── Test 3: noisy_linear resample changes output ───────────────────────
|
||||
|
||||
/// The forward pass produces different outputs before and after
|
||||
/// `resample_noise()`, confirming that noise injection works.
|
||||
///
|
||||
/// Invariant: y1 (before resample) differs from y2 (after resample) in
|
||||
/// at least one element. Since resample_noise fills eps_in/eps_out with
|
||||
/// non-zero factored noise (verified by test 1), the noisy weight
|
||||
/// contribution `sigma_w * eps_w * x + sigma_b * eps_b` is non-zero,
|
||||
/// changing the output relative to the mu-only path.
|
||||
#[test]
|
||||
#[ignore = "requires CUDA"]
|
||||
fn noisy_linear_resample_changes_output() {
|
||||
let Some(dev) = make_device() else { return };
|
||||
let stream = dev.cuda_stream().expect("cuda_stream").clone();
|
||||
|
||||
let cfg = NoisyLinearConfig {
|
||||
in_dim: IN_DIM,
|
||||
out_dim: OUT_DIM,
|
||||
sigma_init: 0.5,
|
||||
seed: 0xBEEF,
|
||||
};
|
||||
let mut layer = NoisyLinear::new(&dev, cfg).expect("NoisyLinear::new");
|
||||
|
||||
let x_host = vec![0.1_f32; B * IN_DIM];
|
||||
let x_d = upload_f32(&stream, &x_host);
|
||||
|
||||
// Forward 1: eps=0 (construction default).
|
||||
let mut y1_d = stream.alloc_zeros::<f32>(B * OUT_DIM).expect("alloc y1");
|
||||
layer
|
||||
.forward(&x_d, B, &mut y1_d)
|
||||
.expect("forward before resample");
|
||||
stream.synchronize().expect("sync y1");
|
||||
let y1 = download_f32(&stream, &y1_d);
|
||||
|
||||
// Resample noise, then forward 2.
|
||||
layer.resample_noise().expect("resample_noise");
|
||||
let mut y2_d = stream.alloc_zeros::<f32>(B * OUT_DIM).expect("alloc y2");
|
||||
layer
|
||||
.forward(&x_d, B, &mut y2_d)
|
||||
.expect("forward after resample");
|
||||
stream.synchronize().expect("sync y2");
|
||||
let y2 = download_f32(&stream, &y2_d);
|
||||
|
||||
// Both outputs must be finite.
|
||||
for (i, (&a, &b)) in y1.iter().zip(y2.iter()).enumerate() {
|
||||
assert!(a.is_finite(), "y1[{i}] not finite: {a}");
|
||||
assert!(b.is_finite(), "y2[{i}] not finite: {b}");
|
||||
}
|
||||
|
||||
// At least one element must differ.
|
||||
let any_differ = y1.iter().zip(y2.iter()).any(|(&a, &b)| (a - b).abs() > 1e-9);
|
||||
assert!(
|
||||
any_differ,
|
||||
"y1 == y2 after resample_noise — noise injection had no effect"
|
||||
);
|
||||
|
||||
eprintln!(
|
||||
"noisy_oracle.3 OK — resample_noise changed output (y1 != y2 confirmed across {} elements)",
|
||||
y1.len()
|
||||
);
|
||||
}
|
||||
Reference in New Issue
Block a user