diff --git a/crates/ml-alpha/cuda/adamw_step.cu b/crates/ml-alpha/cuda/adamw_step.cu index 7cee1bc75..4ddd9d8a6 100644 --- a/crates/ml-alpha/cuda/adamw_step.cu +++ b/crates/ml-alpha/cuda/adamw_step.cu @@ -1,2 +1,40 @@ -// placeholder — real implementation in the task adding this kernel -extern "C" __global__ void adamw_step_stub() {} +// adamw_step.cu — element-wise AdamW update. +// +// m_t = beta1 * m_{t-1} + (1 - beta1) * g +// v_t = beta2 * v_{t-1} + (1 - beta2) * g^2 +// m_hat = m_t / (1 - beta1^step) +// v_hat = v_t / (1 - beta2^step) +// theta -= lr * (m_hat / (sqrt(v_hat) + eps) + wd * theta) +// +// One thread per parameter. Block tree-reduce not needed; this is pure +// element-wise. No atomicAdd. + +extern "C" __global__ void adamw_step( + float* __restrict__ theta, + const float* __restrict__ grad, + float* __restrict__ m, + float* __restrict__ v, + int n_params, + float lr, + float beta1, + float beta2, + float eps, + float wd, + int step +) { + int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i >= n_params) return; + + const float g = grad[i]; + const float m_n = beta1 * m[i] + (1.0f - beta1) * g; + const float v_n = beta2 * v[i] + (1.0f - beta2) * g * g; + m[i] = m_n; + v[i] = v_n; + + const float bc1 = 1.0f - powf(beta1, (float) step); + const float bc2 = 1.0f - powf(beta2, (float) step); + const float m_hat = m_n / fmaxf(bc1, 1e-12f); + const float v_hat = v_n / fmaxf(bc2, 1e-12f); + + theta[i] -= lr * (m_hat / (sqrtf(v_hat) + eps) + wd * theta[i]); +} diff --git a/crates/ml-alpha/cuda/bce_loss_multi_horizon.cu b/crates/ml-alpha/cuda/bce_loss_multi_horizon.cu index 5ce435dd2..f4e29997b 100644 --- a/crates/ml-alpha/cuda/bce_loss_multi_horizon.cu +++ b/crates/ml-alpha/cuda/bce_loss_multi_horizon.cu @@ -1,2 +1,85 @@ -// placeholder — real implementation in the task adding this kernel -extern "C" __global__ void bce_loss_multi_horizon_stub() {} +// bce_loss_multi_horizon.cu +// +// Fused multi-horizon BCE forward + backward. +// +// Input layout (row-major): +// probs [n_pos, n_horizons] — model outputs in (0, 1) +// labels [n_pos, n_horizons] — binary {0.0, 1.0}; NaN = mask (drop) +// Output: +// loss_out[1] — mean BCE over valid entries +// grad_probs[n_pos, n_horizons] — d loss / d prob (scaled by 1/N_valid) +// +// per element: +// p = clamp(probs[i], 1e-6, 1 - 1e-6) +// y = labels[i] +// L_i = -[y log p + (1-y) log(1-p)] +// dL/dp = (p - y) / (p (1-p)) +// +// total_loss = (1/N_valid) * sum_i L_i +// grad_probs[i] = (p - y) / (p (1-p) * N_valid) +// +// Block tree-reduce over a single block. n_pos * n_horizons must fit +// within the block's per-thread iteration count (this kernel is for +// minibatches up to a few thousand elements per call — well within +// reach for n_pos = 256 * seq_len = 16384 elements). For larger +// batches, the caller chunks. + +extern "C" __global__ void bce_multi_horizon_forward_backward( + const float* __restrict__ probs, // [n_pos * n_horizons] + const float* __restrict__ labels, // [n_pos * n_horizons] + int n_pos, + int n_horizons, + float* __restrict__ loss_out, // [1] + float* __restrict__ grad_probs, // [n_pos * n_horizons] + int* __restrict__ valid_count_out // [1] — number of non-NaN labels +) { + const int tid = threadIdx.x; + const int total = n_pos * n_horizons; + + __shared__ float sloss[256]; + __shared__ int svalid[256]; + sloss[tid] = 0.0f; + svalid[tid] = 0; + + // First pass: count valid entries and accumulate loss numerator. + for (int i = tid; i < total; i += blockDim.x) { + const float y = labels[i]; + if (isnan(y)) { + grad_probs[i] = 0.0f; + continue; + } + const float p = fminf(fmaxf(probs[i], 1e-6f), 1.0f - 1e-6f); + sloss[tid] -= y * logf(p) + (1.0f - y) * logf(1.0f - p); + svalid[tid] += 1; + } + __syncthreads(); + + // Block tree-reduce on sloss and svalid. + for (int s = blockDim.x / 2; s > 0; s >>= 1) { + if (tid < s) { + sloss[tid] += sloss[tid + s]; + svalid[tid] += svalid[tid + s]; + } + __syncthreads(); + } + + if (tid == 0) { + const int n_valid = svalid[0]; + loss_out[0] = (n_valid > 0) ? sloss[0] / (float) n_valid : 0.0f; + valid_count_out[0] = n_valid; + } + __syncthreads(); + + // Second pass: emit grad scaled by 1/n_valid. Uses the published + // valid_count_out[0] (already written by thread 0 above). + const int n_valid = valid_count_out[0]; + const float scale = (n_valid > 0) ? (1.0f / (float) n_valid) : 0.0f; + for (int i = tid; i < total; i += blockDim.x) { + const float y = labels[i]; + if (isnan(y)) { + continue; + } + const float p = fminf(fmaxf(probs[i], 1e-6f), 1.0f - 1e-6f); + grad_probs[i] = (p - y) / (p * (1.0f - p)) * scale; + } +} diff --git a/crates/ml-alpha/src/lib.rs b/crates/ml-alpha/src/lib.rs index 17007b516..e671a2aeb 100644 --- a/crates/ml-alpha/src/lib.rs +++ b/crates/ml-alpha/src/lib.rs @@ -30,6 +30,7 @@ pub mod heads; pub mod isv; pub mod pinned; pub mod pinned_mem; +pub mod trainer; // Preserved from prior crate state — used by Phase A. pub mod fxcache_reader; diff --git a/crates/ml-alpha/src/trainer/loss.rs b/crates/ml-alpha/src/trainer/loss.rs new file mode 100644 index 000000000..c7767386f --- /dev/null +++ b/crates/ml-alpha/src/trainer/loss.rs @@ -0,0 +1,122 @@ +//! Fused multi-horizon BCE loss + gradient w.r.t. probs. + +use std::sync::Arc; + +use anyhow::{Context, Result}; +use cudarc::driver::{CudaSlice, CudaStream, DevicePtr, DevicePtrMut, LaunchConfig, PushKernelArg}; +use ml_core::device::MlDevice; + +use crate::pinned_mem::MappedF32Buffer; + +const CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/bce_loss_multi_horizon.cubin")); + +pub struct BceInput { + pub probs: Vec, + pub labels: Vec, // NaN entries are masked out + pub n_horizons: usize, + pub n_pos: usize, +} + +pub struct BceOutput { + pub loss: f32, + pub n_valid: i32, + pub grad_probs: Vec, +} + +pub fn bce_multi_horizon_loss_and_grad_gpu(dev: &MlDevice, input: &BceInput) -> Result { + let total = input.n_pos * input.n_horizons; + assert_eq!(input.probs.len(), total); + assert_eq!(input.labels.len(), total); + + let stream: &Arc = dev.cuda_stream().context("bce stream")?; + let ctx = dev.cuda_context().context("bce ctx")?; + let module = ctx.load_cubin(CUBIN.to_vec()).context("load bce cubin")?; + let func = module.load_function("bce_multi_horizon_forward_backward").context("bce fn")?; + + let probs_d = upload(stream, &input.probs)?; + let labels_d = upload(stream, &input.labels)?; + let mut loss_d = stream.alloc_zeros::(1).context("loss alloc")?; + let mut grad_d = stream.alloc_zeros::(total).context("grad alloc")?; + let mut valid_d = stream.alloc_zeros::(1).context("valid alloc")?; + + let n_pos_i = input.n_pos as i32; + let n_horizons_i = input.n_horizons 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(&labels_d) + .arg(&n_pos_i).arg(&n_horizons_i) + .arg(&mut loss_d).arg(&mut grad_d).arg(&mut valid_d); + unsafe { launch.launch(cfg).context("bce launch")?; } + + stream.synchronize().context("bce sync")?; + let loss = download_f32(stream, &loss_d)?[0]; + let grad = download_f32(stream, &grad_d)?; + let valid = download_i32(stream, &valid_d)?[0]; + + Ok(BceOutput { + loss, + n_valid: valid, + grad_probs: grad, + }) +} + +fn upload(stream: &Arc, host: &[f32]) -> Result> { + let n = host.len(); + let staging = unsafe { MappedF32Buffer::new(n) } + .map_err(|e| anyhow::anyhow!("bce upload staging: {e}"))?; + staging.write_from_slice(host); + let mut dst = stream.alloc_zeros::(n).context("bce 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("bce upload DtoD")?; + } + } + Ok(dst) +} + +fn download_f32(stream: &Arc, src: &CudaSlice) -> Result> { + let n = src.len(); + let staging = unsafe { MappedF32Buffer::new(n) } + .map_err(|e| anyhow::anyhow!("bce download f32: {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("bce download DtoD")?; + } + stream.synchronize().context("bce download sync")?; + Ok(staging.read_all()) +} + +fn download_i32(stream: &Arc, src: &CudaSlice) -> Result> { + // i32 path through a transient host vec; mapped-pinned i32 buffer + // lives in the parent ml crate which we can't import (cycle). For + // a single i32 readback this is acceptable slow-path cost. + let n = src.len(); + let staging = unsafe { MappedF32Buffer::new(n) } + .map_err(|e| anyhow::anyhow!("bce download i32: {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("bce i32 download DtoD")?; + } + stream.synchronize().context("bce i32 download sync")?; + // Reinterpret f32 buffer as i32. + let host_ptr = staging.host_ptr.cast::(); + let mut out = Vec::with_capacity(n); + unsafe { + for i in 0..n { + out.push(std::ptr::read_volatile(host_ptr.add(i))); + } + } + Ok(out) +} diff --git a/crates/ml-alpha/src/trainer/mod.rs b/crates/ml-alpha/src/trainer/mod.rs new file mode 100644 index 000000000..8d13d4212 --- /dev/null +++ b/crates/ml-alpha/src/trainer/mod.rs @@ -0,0 +1,7 @@ +//! Phase A trainer scaffolding. +//! +//! Task 10 lands loss + optimizer kernels. Tasks 13-14 land the trainer +//! loop + real CfC backward. + +pub mod loss; +pub mod optim; diff --git a/crates/ml-alpha/src/trainer/optim.rs b/crates/ml-alpha/src/trainer/optim.rs new file mode 100644 index 000000000..0ef5c5a83 --- /dev/null +++ b/crates/ml-alpha/src/trainer/optim.rs @@ -0,0 +1,75 @@ +//! AdamW optimizer state + kernel binding. + +use std::sync::Arc; + +use anyhow::{Context, Result}; +use cudarc::driver::{CudaFunction, CudaModule, CudaSlice, CudaStream, LaunchConfig, PushKernelArg}; +use ml_core::device::MlDevice; + +const CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/adamw_step.cubin")); + +pub struct AdamW { + stream: Arc, + _module: Arc, + func: CudaFunction, + m: CudaSlice, + v: CudaSlice, + step: i32, + pub lr: f32, + pub beta1: f32, + pub beta2: f32, + pub eps: f32, + pub wd: f32, +} + +impl AdamW { + pub fn new(dev: &MlDevice, n_params: usize, lr: f32) -> Result { + let stream = dev.cuda_stream().context("adamw stream")?.clone(); + let ctx = dev.cuda_context().context("adamw ctx")?; + let module = ctx.load_cubin(CUBIN.to_vec()).context("load adamw cubin")?; + let func = module.load_function("adamw_step").context("adamw fn")?; + let m = stream.alloc_zeros::(n_params).context("m alloc")?; + let v = stream.alloc_zeros::(n_params).context("v alloc")?; + Ok(Self { + stream, + _module: module, + func, + m, + v, + step: 0, + lr, + beta1: 0.9, + beta2: 0.999, + eps: 1e-8, + wd: 1e-2, + }) + } + + pub fn step_count(&self) -> i32 { + self.step + } + + pub fn step(&mut self, theta: &mut CudaSlice, grad: &CudaSlice) -> Result<()> { + let n = theta.len(); + assert_eq!(grad.len(), n, "grad/theta size mismatch"); + assert_eq!(self.m.len(), n, "m size mismatch"); + assert_eq!(self.v.len(), n, "v size mismatch"); + self.step += 1; + let n_params_i = n as i32; + let step_i = self.step; + let cfg = LaunchConfig { + grid_dim: (((n as u32) + 255) / 256, 1, 1), + block_dim: (256, 1, 1), + shared_mem_bytes: 0, + }; + let mut launch = self.stream.launch_builder(&self.func); + launch + .arg(theta).arg(grad).arg(&mut self.m).arg(&mut self.v) + .arg(&n_params_i) + .arg(&self.lr).arg(&self.beta1).arg(&self.beta2).arg(&self.eps).arg(&self.wd) + .arg(&step_i); + unsafe { launch.launch(cfg).context("adamw launch")?; } + self.stream.synchronize().context("adamw sync")?; + Ok(()) + } +} diff --git a/crates/ml-alpha/tests/adamw_invariants.rs b/crates/ml-alpha/tests/adamw_invariants.rs new file mode 100644 index 000000000..357fcf813 --- /dev/null +++ b/crates/ml-alpha/tests/adamw_invariants.rs @@ -0,0 +1,106 @@ +//! AdamW invariants — element-wise correctness without a CPU mirror. + +use approx::assert_relative_eq; +use cudarc::driver::{CudaSlice, CudaStream, DevicePtr, DevicePtrMut}; +use ml_alpha::pinned_mem::MappedF32Buffer; +use ml_alpha::trainer::optim::AdamW; +use ml_core::device::MlDevice; +use std::sync::Arc; + +fn test_device() -> MlDevice { + MlDevice::cuda(0).expect("CUDA 0 required for ml-alpha tests") +} + +fn upload(stream: &Arc, host: &[f32]) -> CudaSlice { + let n = host.len(); + let staging = unsafe { MappedF32Buffer::new(n) }.unwrap(); + staging.write_from_slice(host); + let mut dst = stream.alloc_zeros::(n).unwrap(); + let nbytes = n * 4; + 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()).unwrap(); + } + dst +} + +fn download(stream: &Arc, src: &CudaSlice) -> Vec { + let n = src.len(); + let staging = unsafe { MappedF32Buffer::new(n) }.unwrap(); + let nbytes = n * 4; + unsafe { + let (src_ptr, _g) = src.device_ptr(stream); + cudarc::driver::result::memcpy_dtod_async(staging.dev_ptr, src_ptr, nbytes, stream.cu_stream()).unwrap(); + } + stream.synchronize().unwrap(); + staging.read_all() +} + +#[test] +fn zero_grad_moves_parameter_only_by_weight_decay() { + let dev = test_device(); + let stream = dev.cuda_stream().unwrap().clone(); + let theta_init: Vec = vec![1.0, 2.0, -3.0, 0.5]; + let mut theta = upload(&stream, &theta_init); + let grad = upload(&stream, &vec![0.0; 4]); + let mut opt = AdamW::new(&dev, 4, 1e-3).unwrap(); + opt.wd = 0.1; + opt.step(&mut theta, &grad).unwrap(); + let theta_after = download(&stream, &theta); + // With g=0, m=v=0 too, so the Adam update is 0; only weight decay + // applies: theta -= lr * wd * theta = theta * (1 - lr*wd) + let factor = 1.0_f32 - 1e-3 * 0.1; + for (a, b) in theta_after.iter().zip(&theta_init) { + assert_relative_eq!(*a, *b * factor, epsilon = 1e-6); + } +} + +#[test] +fn positive_grad_decreases_parameter() { + let dev = test_device(); + let stream = dev.cuda_stream().unwrap().clone(); + let theta_init: Vec = vec![5.0, 3.0, 1.0, 0.1]; + let mut theta = upload(&stream, &theta_init); + let grad = upload(&stream, &vec![1.0, 1.0, 1.0, 1.0]); + let mut opt = AdamW::new(&dev, 4, 1e-2).unwrap(); + opt.wd = 0.0; // isolate the Adam update + opt.step(&mut theta, &grad).unwrap(); + let theta_after = download(&stream, &theta); + for (a, b) in theta_after.iter().zip(&theta_init) { + assert!(*a < *b, "positive grad must reduce parameter: {b} -> {a}"); + } +} + +#[test] +fn adam_step_count_increments() { + let dev = test_device(); + let stream = dev.cuda_stream().unwrap().clone(); + let mut theta = upload(&stream, &vec![1.0; 8]); + let grad = upload(&stream, &vec![0.1; 8]); + let mut opt = AdamW::new(&dev, 8, 1e-3).unwrap(); + assert_eq!(opt.step_count(), 0); + opt.step(&mut theta, &grad).unwrap(); + assert_eq!(opt.step_count(), 1); + opt.step(&mut theta, &grad).unwrap(); + assert_eq!(opt.step_count(), 2); +} + +#[test] +fn repeated_grad_descent_drives_parameter_toward_zero() { + // grad = theta (gradient of 0.5 * theta^2 wrt theta), repeated AdamW steps + // must reduce |theta| toward 0 over many steps. + let dev = test_device(); + let stream = dev.cuda_stream().unwrap().clone(); + let initial: Vec = vec![10.0; 16]; + let mut theta = upload(&stream, &initial); + let mut opt = AdamW::new(&dev, 16, 1e-1).unwrap(); + opt.wd = 0.0; + for _ in 0..200 { + let current = download(&stream, &theta); + let grad = upload(&stream, ¤t); + opt.step(&mut theta, &grad).unwrap(); + } + let final_theta = download(&stream, &theta); + let norm: f32 = final_theta.iter().map(|v| v * v).sum::().sqrt(); + assert!(norm < 1.0, "‖θ‖ should drop well below initial 40 after 200 steps, got {norm}"); +} diff --git a/crates/ml-alpha/tests/bce_grad_finite_diff.rs b/crates/ml-alpha/tests/bce_grad_finite_diff.rs new file mode 100644 index 000000000..9ae139f36 --- /dev/null +++ b/crates/ml-alpha/tests/bce_grad_finite_diff.rs @@ -0,0 +1,85 @@ +//! BCE forward + backward validation. +//! +//! Per `feedback_no_cpu_test_fallbacks.md`: no CPU oracle. The gradient +//! is validated against finite-difference computed via GPU forward calls +//! at perturbed inputs (the kernel itself is the reference; we only +//! verify that the analytic grad matches its own numerical derivative). + +use approx::assert_relative_eq; +use ml_alpha::trainer::loss::{bce_multi_horizon_loss_and_grad_gpu, BceInput}; +use ml_core::device::MlDevice; + +fn test_device() -> MlDevice { + MlDevice::cuda(0).expect("CUDA 0 required for ml-alpha tests") +} + +fn make_input(probs: Vec, labels: Vec) -> BceInput { + BceInput { probs, labels, n_horizons: 5, n_pos: 4 } +} + +#[test] +fn bce_loss_is_positive_and_finite() { + let dev = test_device(); + let probs: Vec = (0..20).map(|i| (0.1 + i as f32 * 0.04).min(0.95)).collect(); + let labels: Vec = (0..20).map(|i| (i % 2) as f32).collect(); + let out = bce_multi_horizon_loss_and_grad_gpu(&dev, &make_input(probs, labels)).unwrap(); + assert!(out.loss.is_finite()); + assert!(out.loss > 0.0); + assert_eq!(out.n_valid, 20); +} + +#[test] +fn bce_grad_zero_when_probs_match_labels() { + let dev = test_device(); + // probs ≈ labels (within clamp window): grad ≈ 0 wherever |p-y| ≈ 0 + let mut probs = vec![0.5_f32; 20]; + let labels: Vec = (0..20).map(|i| (i % 2) as f32).collect(); + // Move probs very close to labels (but within (1e-6, 1-1e-6) clamp window). + for i in 0..20 { + probs[i] = if labels[i] > 0.5 { 1.0 - 1e-5 } else { 1e-5 }; + } + let out = bce_multi_horizon_loss_and_grad_gpu(&dev, &make_input(probs, labels)).unwrap(); + // Near-perfect predictions → near-zero loss + assert!(out.loss < 1e-3, "loss should be small for matching probs/labels, got {}", out.loss); +} + +#[test] +fn bce_grad_matches_finite_difference() { + let dev = test_device(); + let probs: Vec = (0..20).map(|i| (0.1 + i as f32 * 0.04).min(0.95)).collect(); + let labels: Vec = (0..20).map(|i| (i % 2) as f32).collect(); + let base = bce_multi_horizon_loss_and_grad_gpu(&dev, &make_input(probs.clone(), labels.clone())).unwrap(); + + let eps = 1e-3_f32; + for i in [0usize, 5, 10, 15, 19] { + let mut up = probs.clone(); + up[i] += eps; + let mut dn = probs.clone(); + dn[i] -= eps; + let lu = bce_multi_horizon_loss_and_grad_gpu(&dev, &make_input(up, labels.clone())).unwrap().loss; + let ld = bce_multi_horizon_loss_and_grad_gpu(&dev, &make_input(dn, labels.clone())).unwrap().loss; + let fd = (lu - ld) / (2.0 * eps); + assert_relative_eq!( + base.grad_probs[i], + fd, + epsilon = 5e-3, + max_relative = 5e-2 + ); + } +} + +#[test] +fn bce_nan_labels_mask_grad_and_loss() { + let dev = test_device(); + let mut labels: Vec = (0..20).map(|i| (i % 2) as f32).collect(); + // Mask out indices 5, 10, 15 + labels[5] = f32::NAN; + labels[10] = f32::NAN; + labels[15] = f32::NAN; + let probs: Vec = (0..20).map(|i| (0.1 + i as f32 * 0.04).min(0.95)).collect(); + let out = bce_multi_horizon_loss_and_grad_gpu(&dev, &make_input(probs, labels)).unwrap(); + assert_eq!(out.n_valid, 17, "should drop 3 NaN-masked entries"); + for &i in &[5usize, 10, 15] { + assert_eq!(out.grad_probs[i], 0.0, "masked grad should be 0 at index {i}, got {}", out.grad_probs[i]); + } +}