diff --git a/crates/ml-alpha/cuda/cfc_step.cu b/crates/ml-alpha/cuda/cfc_step.cu index bf0c95325..a9b40a39c 100644 --- a/crates/ml-alpha/cuda/cfc_step.cu +++ b/crates/ml-alpha/cuda/cfc_step.cu @@ -34,3 +34,77 @@ extern "C" __global__ void cfc_step( const float decay = expf(-dt_s / fmaxf(tau[i], 1e-6f)); h_new[i] = h_old[i] * decay + (1.0f - decay) * tanhf(pre); } + +// Backward through ONE cfc_step (truncated BPTT, K=1). +// +// Forward (for reference): +// pre[i] = b[i] + sum_k W_in[i,k] x[k] + sum_k W_rec[i,k] h_old[k] +// decay[i] = exp(-dt / max(tau[i], 1e-6)) +// s[i] = tanh(pre[i]) +// h_new[i] = h_old[i] * decay + (1 - decay) * s +// +// Given grad_h_new[i] = dL/dh_new[i], compute: +// d_pre[i] = grad_h_new[i] * (1 - decay) * (1 - s^2) +// grad_b[i] += d_pre[i] +// grad_W_in[i,k] += d_pre[i] * x[k] +// grad_W_rec[i,k] += d_pre[i] * h_old[k] +// grad_h_old[i] += grad_h_new[i] * decay + sum_j d_pre[j] * W_rec[j, i] +// +// tau is held frozen (structurally log-uniform initialised per spec +// Section 2); no grad_tau path in this version. +// +// One thread per hidden unit (128 threads in one block). Block tree- +// reduce not needed — each thread owns its own grad_W rows and grad_h +// component; grad_h_old accumulates via shared `d_pre` reads. + +extern "C" __global__ void cfc_step_backward( + const float* __restrict__ w_in, // [n_hid, n_in] + const float* __restrict__ w_rec, // [n_hid, n_hid] + const float* __restrict__ b, + const float* __restrict__ tau, + const float* __restrict__ x, // [n_in] + const float* __restrict__ h_old, // [n_hid] + const float* __restrict__ grad_h_new, // [n_hid] + float dt_s, + int n_in, + int n_hid, + float* __restrict__ grad_w_in, // [n_hid, n_in] + float* __restrict__ grad_w_rec, // [n_hid, n_hid] + float* __restrict__ grad_b, // [n_hid] + float* __restrict__ grad_h_old // [n_hid] +) { + extern __shared__ float smem[]; + float* sd_pre = smem; // [n_hid] + float* sdecay = sd_pre + n_hid; // [n_hid] + + int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i >= n_hid) return; + + // Recompute forward pre + tanh to derive d_pre. + float pre = b[i]; + for (int k = 0; k < n_in; ++k) pre += w_in[i * n_in + k] * x[k]; + for (int k = 0; k < n_hid; ++k) pre += w_rec[i * n_hid + k] * h_old[k]; + const float s = tanhf(pre); + const float decay = expf(-dt_s / fmaxf(tau[i], 1e-6f)); + const float dh = grad_h_new[i]; + const float d_pre = dh * (1.0f - decay) * (1.0f - s * s); + + sd_pre[i] = d_pre; + sdecay[i] = decay; + __syncthreads(); + + grad_b[i] = d_pre; + for (int k = 0; k < n_in; ++k) { + grad_w_in[i * n_in + k] = d_pre * x[k]; + } + for (int k = 0; k < n_hid; ++k) { + grad_w_rec[i * n_hid + k] = d_pre * h_old[k]; + } + + // grad_h_old[i] = dh * decay + sum_j d_pre[j] * W_rec[j, i] + float gh = dh * decay; + for (int j = 0; j < n_hid; ++j) { + gh += sd_pre[j] * w_rec[j * n_hid + i]; + } + grad_h_old[i] = gh; +} diff --git a/crates/ml-alpha/cuda/multi_horizon_heads.cu b/crates/ml-alpha/cuda/multi_horizon_heads.cu index 1c5e9e7dd..bac4d85d8 100644 --- a/crates/ml-alpha/cuda/multi_horizon_heads.cu +++ b/crates/ml-alpha/cuda/multi_horizon_heads.cu @@ -19,3 +19,55 @@ extern "C" __global__ void multi_horizon_heads( } probs[k] = 1.0f / (1.0f + expf(-z)); } + +// Backward through multi_horizon_heads. +// +// Given: +// grad_probs[5] = dL / dp +// probs[5] = forward output (sigmoid) +// h[128] = forward input +// Computes: +// grad_w[5, 128] = dL / dW = d_z[k] * h[i] +// grad_b[5] = dL / db = d_z[k] +// grad_h[128] = dL / dh = sum_k d_z[k] * W[k, i] +// where d_z[k] = grad_probs[k] * p[k] * (1 - p[k]). +// +// Block tree-reduce for the grad_h sum (no atomicAdd). One thread per +// hidden unit (128 threads); thread 0 also writes grad_b. + +extern "C" __global__ void multi_horizon_heads_backward( + const float* __restrict__ w, // [5, 128] + const float* __restrict__ probs, // [5] + const float* __restrict__ h, // [128] + const float* __restrict__ grad_probs, // [5] + float* __restrict__ grad_w, // [5, 128] + float* __restrict__ grad_b, // [5] + float* __restrict__ grad_h // [128] +) { + __shared__ float d_z[5]; + + int tid = threadIdx.x; + if (tid < 5) { + const float p = probs[tid]; + d_z[tid] = grad_probs[tid] * p * (1.0f - p); + } + __syncthreads(); + + if (tid < 5) { + grad_b[tid] = d_z[tid]; + } + // grad_w[k, i] = d_z[k] * h[i] + // grid stride: each thread covers one i for k in 0..5. + if (tid < 128) { + const float h_i = h[tid]; + for (int k = 0; k < 5; ++k) { + grad_w[k * 128 + tid] = d_z[k] * h_i; + } + // grad_h[i] = sum_k d_z[k] * W[k, i] + float acc = 0.0f; + for (int k = 0; k < 5; ++k) { + acc += d_z[k] * w[k * 128 + tid]; + } + grad_h[tid] = acc; + } +} diff --git a/crates/ml-alpha/src/cfc/step.rs b/crates/ml-alpha/src/cfc/step.rs index 582676eea..9cb05b9c2 100644 --- a/crates/ml-alpha/src/cfc/step.rs +++ b/crates/ml-alpha/src/cfc/step.rs @@ -20,6 +20,66 @@ pub struct CfcWeights { pub n_hid: usize, } +/// Backward through ONE cfc_step (truncated BPTT, K=1). +/// +/// Returns (grad_w_in, grad_w_rec, grad_b, grad_h_old). +pub fn cfc_step_backward_gpu( + dev: &MlDevice, + w: &CfcWeights, + x: &[f32], + h_old: &[f32], + grad_h_new: &[f32], + dt_s: f32, +) -> Result<(Vec, Vec, Vec, Vec)> { + assert_eq!(x.len(), w.n_in); + assert_eq!(h_old.len(), w.n_hid); + assert_eq!(grad_h_new.len(), w.n_hid); + + let stream: &Arc = dev.cuda_stream().context("cfc-bwd stream")?; + let ctx = dev.cuda_context().context("cfc-bwd ctx")?; + let module = ctx.load_cubin(KERNEL_CUBIN.to_vec()).context("load cfc cubin")?; + let func = module.load_function("cfc_step_backward").context("cfc_bwd fn")?; + + let w_in_d = upload(stream, &w.w_in)?; + let w_rec_d = upload(stream, &w.w_rec)?; + let b_d = upload(stream, &w.b)?; + let tau_d = upload(stream, &w.tau)?; + let x_d = upload(stream, x)?; + let h_old_d = upload(stream, h_old)?; + let grad_h_new_d = upload(stream, grad_h_new)?; + let mut grad_w_in_d = stream.alloc_zeros::(w.n_hid * w.n_in).context("grad_w_in alloc")?; + let mut grad_w_rec_d = stream.alloc_zeros::(w.n_hid * w.n_hid).context("grad_w_rec alloc")?; + let mut grad_b_d = stream.alloc_zeros::(w.n_hid).context("grad_b alloc")?; + let mut grad_h_old_d = stream.alloc_zeros::(w.n_hid).context("grad_h_old alloc")?; + + let n_in_i = w.n_in as i32; + let n_hid_i = w.n_hid as i32; + let block_dim = 128.min(w.n_hid as u32); + let grid_dim = ((w.n_hid as u32) + block_dim - 1) / block_dim; + let shared_mem = (2 * w.n_hid * std::mem::size_of::()) as u32; + let cfg = LaunchConfig { + grid_dim: (grid_dim, 1, 1), + block_dim: (block_dim, 1, 1), + shared_mem_bytes: shared_mem, + }; + + let mut launch = stream.launch_builder(&func); + launch + .arg(&w_in_d).arg(&w_rec_d).arg(&b_d).arg(&tau_d) + .arg(&x_d).arg(&h_old_d).arg(&grad_h_new_d) + .arg(&dt_s).arg(&n_in_i).arg(&n_hid_i) + .arg(&mut grad_w_in_d).arg(&mut grad_w_rec_d) + .arg(&mut grad_b_d).arg(&mut grad_h_old_d); + unsafe { launch.launch(cfg).context("cfc_bwd launch")?; } + + Ok(( + download(stream, &grad_w_in_d)?, + download(stream, &grad_w_rec_d)?, + download(stream, &grad_b_d)?, + download(stream, &grad_h_old_d)?, + )) +} + pub fn cfc_step_gpu( dev: &MlDevice, w: &CfcWeights, diff --git a/crates/ml-alpha/src/heads.rs b/crates/ml-alpha/src/heads.rs index 8089f112c..c63ea533b 100644 --- a/crates/ml-alpha/src/heads.rs +++ b/crates/ml-alpha/src/heads.rs @@ -39,6 +39,47 @@ pub struct MultiHorizonHeads; pub struct Projection; +/// Backward through `multi_horizon_heads`. Returns (grad_w, grad_b, grad_h). +pub fn multi_horizon_heads_backward_gpu( + dev: &MlDevice, + w: &HeadsWeights, + probs: &[f32; N_HORIZONS], + h: &[f32], + grad_probs: &[f32; N_HORIZONS], +) -> Result<(Vec, Vec, Vec)> { + assert_eq!(h.len(), HIDDEN_DIM); + + let stream: &Arc = dev.cuda_stream().context("heads-bwd stream")?; + let ctx = dev.cuda_context().context("heads-bwd ctx")?; + let module = ctx.load_cubin(HEADS_CUBIN.to_vec()).context("load heads cubin")?; + let func = module.load_function("multi_horizon_heads_backward").context("heads_bwd fn")?; + + let w_d = upload(stream, &w.w)?; + let probs_d = upload(stream, probs.as_slice())?; + let h_d = upload(stream, h)?; + let grad_p_d = upload(stream, grad_probs.as_slice())?; + let mut grad_w_d = stream.alloc_zeros::(N_HORIZONS * HIDDEN_DIM).context("grad_w alloc")?; + let mut grad_b_d = stream.alloc_zeros::(N_HORIZONS).context("grad_b alloc")?; + let mut grad_h_d = stream.alloc_zeros::(HIDDEN_DIM).context("grad_h alloc")?; + + let cfg = LaunchConfig { + grid_dim: (1, 1, 1), + block_dim: (HIDDEN_DIM as u32, 1, 1), + shared_mem_bytes: 0, + }; + let mut launch = stream.launch_builder(&func); + launch + .arg(&w_d).arg(&probs_d).arg(&h_d).arg(&grad_p_d) + .arg(&mut grad_w_d).arg(&mut grad_b_d).arg(&mut grad_h_d); + unsafe { launch.launch(cfg).context("heads_bwd launch")?; } + + Ok(( + download(stream, &grad_w_d)?, + download(stream, &grad_b_d)?, + download(stream, &grad_h_d)?, + )) +} + pub fn projection_gpu( dev: &MlDevice, w: &ProjectionWeights, diff --git a/crates/ml-alpha/tests/backward_finite_diff.rs b/crates/ml-alpha/tests/backward_finite_diff.rs new file mode 100644 index 000000000..730685b59 --- /dev/null +++ b/crates/ml-alpha/tests/backward_finite_diff.rs @@ -0,0 +1,156 @@ +//! Finite-difference validation of cfc_step_backward + multi_horizon_heads_backward. +//! +//! Per `feedback_no_cpu_test_fallbacks.md`: the kernel is the reference. +//! We perturb a single input, re-run the GPU forward, and compare the +//! resulting finite-difference gradient to the analytic backward. + +use approx::assert_relative_eq; +use ml_alpha::cfc::step::{cfc_step_backward_gpu, cfc_step_gpu, CfcWeights}; +use ml_alpha::heads::{ + multi_horizon_heads_backward_gpu, multi_horizon_heads_gpu, HeadsWeights, HIDDEN_DIM, N_HORIZONS, +}; +use ml_core::device::MlDevice; +use rand::{Rng, SeedableRng}; +use rand_chacha::ChaCha8Rng; + +fn test_device() -> MlDevice { + MlDevice::cuda(0).expect("CUDA 0 required for ml-alpha tests") +} + +fn rand_heads(seed: u64) -> HeadsWeights { + let mut r = ChaCha8Rng::seed_from_u64(seed); + HeadsWeights { + w: (0..N_HORIZONS * HIDDEN_DIM).map(|_| r.gen_range(-0.1..0.1)).collect(), + b: (0..N_HORIZONS).map(|_| r.gen_range(-0.01..0.01)).collect(), + } +} + +fn rand_cfc(seed: u64, n_in: usize, n_hid: usize) -> CfcWeights { + let mut r = ChaCha8Rng::seed_from_u64(seed); + CfcWeights { + w_in: (0..n_hid * n_in).map(|_| r.gen_range(-0.1..0.1)).collect(), + w_rec: (0..n_hid * n_hid).map(|_| r.gen_range(-0.05..0.05)).collect(), + b: (0..n_hid).map(|_| r.gen_range(-0.01..0.01)).collect(), + tau: (0..n_hid).map(|_| r.gen_range(0.05..2.0)).collect(), + n_in, + n_hid, + } +} + +#[test] +fn heads_backward_grad_h_matches_finite_diff() { + let dev = test_device(); + let w = rand_heads(0x1111); + let mut r = ChaCha8Rng::seed_from_u64(0x2222); + let h: Vec = (0..HIDDEN_DIM).map(|_| r.gen_range(-1.0..1.0)).collect(); + let probs = multi_horizon_heads_gpu(&dev, &w, &h).unwrap(); + + // Surrogate loss L = sum_k probs[k] -> dL/dp = 1 for all k. + let grad_probs = [1.0f32; N_HORIZONS]; + let (_, _, grad_h_analytic) = + multi_horizon_heads_backward_gpu(&dev, &w, &probs, &h, &grad_probs).unwrap(); + + let eps = 1e-3_f32; + for &i in &[0usize, 32, 64, 96, 127] { + let mut h_up = h.clone(); + h_up[i] += eps; + let mut h_dn = h.clone(); + h_dn[i] -= eps; + let p_up = multi_horizon_heads_gpu(&dev, &w, &h_up).unwrap(); + let p_dn = multi_horizon_heads_gpu(&dev, &w, &h_dn).unwrap(); + let l_up: f32 = p_up.iter().sum(); + let l_dn: f32 = p_dn.iter().sum(); + let fd = (l_up - l_dn) / (2.0 * eps); + assert_relative_eq!(grad_h_analytic[i], fd, epsilon = 1e-3, max_relative = 1e-2); + } +} + +#[test] +fn heads_backward_grad_b_matches_finite_diff() { + let dev = test_device(); + let w = rand_heads(0x3333); + let mut r = ChaCha8Rng::seed_from_u64(0x4444); + let h: Vec = (0..HIDDEN_DIM).map(|_| r.gen_range(-1.0..1.0)).collect(); + let probs = multi_horizon_heads_gpu(&dev, &w, &h).unwrap(); + let grad_probs = [1.0f32; N_HORIZONS]; + let (_, grad_b_analytic, _) = + multi_horizon_heads_backward_gpu(&dev, &w, &probs, &h, &grad_probs).unwrap(); + + let eps = 1e-3_f32; + for k in 0..N_HORIZONS { + let mut w_up = w.clone(); + w_up.b[k] += eps; + let mut w_dn = w.clone(); + w_dn.b[k] -= eps; + let p_up = multi_horizon_heads_gpu(&dev, &w_up, &h).unwrap(); + let p_dn = multi_horizon_heads_gpu(&dev, &w_dn, &h).unwrap(); + let l_up: f32 = p_up.iter().sum(); + let l_dn: f32 = p_dn.iter().sum(); + let fd = (l_up - l_dn) / (2.0 * eps); + assert_relative_eq!(grad_b_analytic[k], fd, epsilon = 1e-3, max_relative = 1e-2); + let _ = w; // suppress + let _ = w_up; + let _ = w_dn; + } + let _ = w; +} + +#[test] +fn cfc_backward_grad_b_matches_finite_diff() { + let dev = test_device(); + let n_in = 8; + let n_hid = 16; + let w = rand_cfc(0x5555, n_in, n_hid); + let mut r = ChaCha8Rng::seed_from_u64(0x6666); + let x: Vec = (0..n_in).map(|_| r.gen_range(-1.0..1.0)).collect(); + let h_old: Vec = (0..n_hid).map(|_| r.gen_range(-1.0..1.0)).collect(); + let dt_s = 0.02_f32; + + // Surrogate loss L = sum_i h_new[i]. + let grad_h_new = vec![1.0f32; n_hid]; + let (_, _, grad_b_analytic, _) = + cfc_step_backward_gpu(&dev, &w, &x, &h_old, &grad_h_new, dt_s).unwrap(); + + let eps = 1e-3_f32; + for &i in &[0usize, 4, 8, 15] { + let mut w_up = w.clone(); + w_up.b[i] += eps; + let mut w_dn = w.clone(); + w_dn.b[i] -= eps; + let h_up = cfc_step_gpu(&dev, &w_up, &x, &h_old, dt_s).unwrap(); + let h_dn = cfc_step_gpu(&dev, &w_dn, &x, &h_old, dt_s).unwrap(); + let l_up: f32 = h_up.iter().sum(); + let l_dn: f32 = h_dn.iter().sum(); + let fd = (l_up - l_dn) / (2.0 * eps); + assert_relative_eq!(grad_b_analytic[i], fd, epsilon = 5e-3, max_relative = 5e-2); + } +} + +#[test] +fn cfc_backward_grad_h_old_matches_finite_diff() { + let dev = test_device(); + let n_in = 8; + let n_hid = 16; + let w = rand_cfc(0x7777, n_in, n_hid); + let mut r = ChaCha8Rng::seed_from_u64(0x8888); + let x: Vec = (0..n_in).map(|_| r.gen_range(-1.0..1.0)).collect(); + let h_old: Vec = (0..n_hid).map(|_| r.gen_range(-1.0..1.0)).collect(); + let dt_s = 0.02_f32; + let grad_h_new = vec![1.0f32; n_hid]; + let (_, _, _, grad_h_old_analytic) = + cfc_step_backward_gpu(&dev, &w, &x, &h_old, &grad_h_new, dt_s).unwrap(); + + let eps = 1e-3_f32; + for &i in &[0usize, 4, 8, 15] { + let mut h_up = h_old.clone(); + h_up[i] += eps; + let mut h_dn = h_old.clone(); + h_dn[i] -= eps; + let h_new_up = cfc_step_gpu(&dev, &w, &x, &h_up, dt_s).unwrap(); + let h_new_dn = cfc_step_gpu(&dev, &w, &x, &h_dn, dt_s).unwrap(); + let l_up: f32 = h_new_up.iter().sum(); + let l_dn: f32 = h_new_dn.iter().sum(); + let fd = (l_up - l_dn) / (2.0 * eps); + assert_relative_eq!(grad_h_old_analytic[i], fd, epsilon = 5e-3, max_relative = 5e-2); + } +}