diff --git a/crates/ml-alpha/build.rs b/crates/ml-alpha/build.rs index 81122e5f6..708b40194 100644 --- a/crates/ml-alpha/build.rs +++ b/crates/ml-alpha/build.rs @@ -78,6 +78,7 @@ const KERNELS: &[&str] = &[ "rl_frd_fwd", // SP20 P3: Forward-Return-Distribution head fwd — 2-layer MLP [HIDDEN_DIM → FRD_HIDDEN_DIM → FRD_N_HORIZONS × FRD_N_ATOMS]; ReLU hidden cached for bwd; softmax + CE happen in bwd "rl_frd_softmax_ce_grad", // SP20 P3 F.3a: per-(batch, horizon) softmax + CE loss + dL/dlogits; 1 block per (b, h), 21 threads; label = -1 sentinel masks the row "rl_frd_layer2_bwd", // SP20 P3 F.3b: FRD head layer-2 backward — dW2 (per-batch scratch), db2 (per-batch scratch), dhidden (per-batch overwrite); 1 block per batch, 64 threads + "rl_frd_layer1_bwd", // SP20 P3 F.3c: FRD head layer-1 backward — dW1 (per-batch scratch), db1 (per-batch scratch), dh_t (per-batch overwrite); applies ReLU mask via cached post-ReLU hidden; 1 block per batch, 128 threads ]; // Cache bust v31 — five new reduce / derive kernels populate the input diff --git a/crates/ml-alpha/cuda/rl_frd_layer1_bwd.cu b/crates/ml-alpha/cuda/rl_frd_layer1_bwd.cu new file mode 100644 index 000000000..8bc9e4f76 --- /dev/null +++ b/crates/ml-alpha/cuda/rl_frd_layer1_bwd.cu @@ -0,0 +1,79 @@ +// rl_frd_layer1_bwd.cu — FRD head backward stage 3 (SP20 P3 F.3c). +// +// Given dL/dhidden [B, FRD_HIDDEN_DIM] (from rl_frd_layer2_bwd) and +// the cached post-ReLU hidden [B, FRD_HIDDEN_DIM] (from rl_frd_fwd), +// compute the layer-1 gradients via the chain rule with ReLU mask: +// +// pre_hidden[b, i] = Σ_k h_t[b, k] × W1[k, i] + b1[i] +// hidden[b, i] = ReLU(pre_hidden[b, i]) +// +// dL/dpre_hidden[b, i] = dL/dhidden[b, i] × 1{hidden[b, i] > 0} +// dL/dW1[k, i] = Σ_b h_t[b, k] × dL/dpre_hidden[b, i] ← reduce over B +// dL/db1[i] = Σ_b dL/dpre_hidden[b, i] ← reduce over B +// dL/dh_t[b, k] = Σ_i W1[k, i] × dL/dpre_hidden[b, i] ← per-batch overwrite +// +// Per-batch scratch: +// grad_W1_per_batch[b, k, i] = h_t[b, k] × dL/dpre_hidden[b, i] +// grad_b1_per_batch[b, i] = dL/dpre_hidden[b, i] +// +// Block layout: 1 block per batch, HIDDEN_DIM=128 threads. +// * Phase 0: stage post-ReLU `hidden` (64 slots) + dL/dhidden into +// shared, apply ReLU mask in-place to produce dL/dpre_hidden +// (only first 64 threads work; remaining 64 idle). +// * Phase 1: each thread k (k < 128) is the sole writer of: +// - grad_W1_per_batch[b, k, 0..64] (64 writes per thread) +// - grad_h_t[b, k] (one Σ_i across 64 outputs) +// * Phase 2: thread i (i < 64) writes grad_b1_per_batch[b, i]. +// +// Per `feedback_no_atomicadd`: per-(b, k, i) sole-writer pattern. +// Per `feedback_cpu_is_read_only`: pure device-side. + +#include + +#define HIDDEN_DIM 128 +#define FRD_HIDDEN_DIM 64 + +extern "C" __global__ void rl_frd_layer1_bwd( + const float* __restrict__ h_t, // [B, HIDDEN_DIM] + const float* __restrict__ hidden, // [B, FRD_HIDDEN_DIM] (post-ReLU cache) + const float* __restrict__ grad_hidden, // [B, FRD_HIDDEN_DIM] (from layer2_bwd) + const float* __restrict__ w1, // [HIDDEN_DIM, FRD_HIDDEN_DIM] + int b_size, + float* __restrict__ grad_w1_per_batch, // [B, HIDDEN_DIM, FRD_HIDDEN_DIM] + float* __restrict__ grad_b1_per_batch, // [B, FRD_HIDDEN_DIM] + float* __restrict__ grad_h_t // [B, HIDDEN_DIM] +) { + const int b = blockIdx.x; + const int k = threadIdx.x; + if (b >= b_size || k >= HIDDEN_DIM) return; + + // Stage dL/dpre_hidden in shared — apply ReLU mask using the + // cached post-ReLU hidden (hidden > 0 ⇒ mask=1, else mask=0). + __shared__ float s_grad_pre[FRD_HIDDEN_DIM]; + if (k < FRD_HIDDEN_DIM) { + const int idx = b * FRD_HIDDEN_DIM + k; + const float h_val = hidden[idx]; + const float g_val = grad_hidden[idx]; + s_grad_pre[k] = (h_val > 0.0f) ? g_val : 0.0f; + // grad_b1_per_batch[b, i] = dL/dpre_hidden[b, i] — same value. + grad_b1_per_batch[idx] = s_grad_pre[k]; + } + __syncthreads(); + + const float h_t_bk = h_t[b * HIDDEN_DIM + k]; + + // grad_W1_per_batch[b, k, i] = h_t_bk × s_grad_pre[i] + const int row_off = (b * HIDDEN_DIM + k) * FRD_HIDDEN_DIM; + #pragma unroll + for (int i = 0; i < FRD_HIDDEN_DIM; ++i) { + grad_w1_per_batch[row_off + i] = h_t_bk * s_grad_pre[i]; + } + + // grad_h_t[b, k] = Σ_i W1[k, i] × s_grad_pre[i] + float acc = 0.0f; + #pragma unroll + for (int i = 0; i < FRD_HIDDEN_DIM; ++i) { + acc += w1[k * FRD_HIDDEN_DIM + i] * s_grad_pre[i]; + } + grad_h_t[b * HIDDEN_DIM + k] = acc; +} diff --git a/crates/ml-alpha/src/rl/frd.rs b/crates/ml-alpha/src/rl/frd.rs index 157f50e40..a5bef1aa0 100644 --- a/crates/ml-alpha/src/rl/frd.rs +++ b/crates/ml-alpha/src/rl/frd.rs @@ -62,6 +62,8 @@ const FRD_SOFTMAX_CE_GRAD_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/rl_frd_softmax_ce_grad.cubin")); const FRD_LAYER2_BWD_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/rl_frd_layer2_bwd.cubin")); +const FRD_LAYER1_BWD_CUBIN: &[u8] = + include_bytes!(concat!(env!("OUT_DIR"), "/rl_frd_layer1_bwd.cubin")); /// Per-batch output width: `FRD_N_HORIZONS × FRD_N_ATOMS`. Each batch /// row contains 3 contiguous horizon blocks of 21 atom logits. @@ -93,6 +95,8 @@ pub struct FrdHead { softmax_ce_grad_fn: CudaFunction, _layer2_bwd_module: Arc, layer2_bwd_fn: CudaFunction, + _layer1_bwd_module: Arc, + layer1_bwd_fn: CudaFunction, pub w1_d: CudaSlice, // [HIDDEN_DIM, FRD_HIDDEN_DIM] pub b1_d: CudaSlice, // [FRD_HIDDEN_DIM] @@ -122,6 +126,12 @@ impl FrdHead { let layer2_bwd_fn = layer2_bwd_module .load_function("rl_frd_layer2_bwd") .context("load rl_frd_layer2_bwd fn")?; + let layer1_bwd_module = ctx + .load_cubin(FRD_LAYER1_BWD_CUBIN.to_vec()) + .context("load rl_frd_layer1_bwd cubin")?; + let layer1_bwd_fn = layer1_bwd_module + .load_function("rl_frd_layer1_bwd") + .context("load rl_frd_layer1_bwd fn")?; // Per pearl_scoped_init_seed_for_reproducibility — guard around // Xavier draws so GPU init helpers downstream see the same RNG. @@ -159,6 +169,8 @@ impl FrdHead { softmax_ce_grad_fn, _layer2_bwd_module: layer2_bwd_module, layer2_bwd_fn, + _layer1_bwd_module: layer1_bwd_module, + layer1_bwd_fn, w1_d, b1_d, w2_d, @@ -260,6 +272,63 @@ impl FrdHead { Ok(()) } + /// Backward stage 3 (F.3c): layer-1 weight gradients with ReLU mask. + /// + /// Given `h_t` (encoder hidden, the layer-1 input), `hidden` (post- + /// ReLU cache from forward — used for the ReLU mask), `grad_hidden` + /// (from `layer2_bwd`), and `self.w1_d`, computes: + /// * `grad_w1_per_batch[b, k, i] = h_t[b, k] × dpre_hidden[b, i]` + /// * `grad_b1_per_batch[b, i] = dpre_hidden[b, i]` + /// * `grad_h_t[b, k] = Σ_i W1[k, i] × dpre_hidden[b, i]` + /// + /// where `dpre_hidden[b, i] = grad_hidden[b, i] × 1{hidden[b, i] > 0}` + /// is the ReLU-masked upstream gradient. + /// + /// `grad_h_t` is the encoder-upstream gradient — the caller's + /// `grad_h_accumulate` kernel folds it into the encoder's grad with + /// a `λ_frd`-scaled accumulator (same pattern as Q / π / V heads). + #[allow(clippy::too_many_arguments)] + pub fn layer1_bwd( + &self, + h_t_d: &CudaSlice, + hidden_d: &CudaSlice, + grad_hidden_d: &CudaSlice, + grad_w1_per_batch_d: &mut CudaSlice, + grad_b1_per_batch_d: &mut CudaSlice, + grad_h_t_d: &mut CudaSlice, + b_size: usize, + ) -> Result<()> { + debug_assert_eq!(h_t_d.len(), b_size * HIDDEN_DIM); + debug_assert_eq!(hidden_d.len(), b_size * FRD_HIDDEN_DIM); + debug_assert_eq!(grad_hidden_d.len(), b_size * FRD_HIDDEN_DIM); + debug_assert_eq!( + grad_w1_per_batch_d.len(), + b_size * HIDDEN_DIM * FRD_HIDDEN_DIM + ); + debug_assert_eq!(grad_b1_per_batch_d.len(), b_size * FRD_HIDDEN_DIM); + debug_assert_eq!(grad_h_t_d.len(), b_size * HIDDEN_DIM); + let cfg = LaunchConfig { + grid_dim: (b_size as u32, 1, 1), + block_dim: (HIDDEN_DIM as u32, 1, 1), + shared_mem_bytes: 0, + }; + let b_size_i = b_size as i32; + let mut launch = self.stream.launch_builder(&self.layer1_bwd_fn); + launch + .arg(h_t_d) + .arg(hidden_d) + .arg(grad_hidden_d) + .arg(&self.w1_d) + .arg(&b_size_i) + .arg(grad_w1_per_batch_d) + .arg(grad_b1_per_batch_d) + .arg(grad_h_t_d); + unsafe { + launch.launch(cfg).context("rl_frd_layer1_bwd launch")?; + } + Ok(()) + } + /// Forward pass — fills `logits_out_d [B, FRD_OUT_DIM]` and caches /// the post-ReLU hidden activations in `hidden_out_d [B, /// FRD_HIDDEN_DIM]` for the backward kernel. diff --git a/crates/ml-alpha/tests/frd_head.rs b/crates/ml-alpha/tests/frd_head.rs index 1ad9744ec..e7a4bc1d9 100644 --- a/crates/ml-alpha/tests/frd_head.rs +++ b/crates/ml-alpha/tests/frd_head.rs @@ -518,3 +518,220 @@ fn frd_layer2_bwd_db2_equals_grad_logits() -> Result<()> { ); Ok(()) } + +#[test] +#[ignore = "requires CUDA (MlDevice::cuda(0))"] +fn frd_layer1_bwd_finite_diff_w1() -> Result<()> { + let Some((dev, mut head)) = build_head() else { return Ok(()) }; + let stream = dev.cuda_stream()?.clone(); + let b_size = 1; + + use rand::{Rng, SeedableRng}; + use rand_chacha::ChaCha8Rng; + let mut r = ChaCha8Rng::seed_from_u64(0xB1); + let h_t: Vec = (0..b_size * HIDDEN_DIM) + .map(|_| r.gen_range(-1.0..1.0)) + .collect(); + let h_t_d = upload_f32(&stream, &h_t)?; + let labels: Vec = vec![5, 10, 15]; + let labels_d = upload_i32(&stream, &labels)?; + + // Full backward chain at the unperturbed weights. + let mut hidden_d = stream.alloc_zeros::(b_size * FRD_HIDDEN_DIM)?; + let mut logits_d = stream.alloc_zeros::(b_size * FRD_OUT_DIM)?; + head.forward(&h_t_d, &mut hidden_d, &mut logits_d, b_size)?; + + let mut grad_logits_d = stream.alloc_zeros::(b_size * FRD_OUT_DIM)?; + let mut loss_d = stream.alloc_zeros::(b_size * FRD_N_HORIZONS)?; + head.softmax_ce_grad(&logits_d, &labels_d, &mut grad_logits_d, &mut loss_d, b_size)?; + + let mut grad_w2_pb_d = stream.alloc_zeros::(b_size * FRD_HIDDEN_DIM * FRD_OUT_DIM)?; + let mut grad_b2_pb_d = stream.alloc_zeros::(b_size * FRD_OUT_DIM)?; + let mut grad_hidden_d = stream.alloc_zeros::(b_size * FRD_HIDDEN_DIM)?; + head.layer2_bwd( + &hidden_d, + &grad_logits_d, + &mut grad_w2_pb_d, + &mut grad_b2_pb_d, + &mut grad_hidden_d, + b_size, + )?; + + let mut grad_w1_pb_d = stream.alloc_zeros::(b_size * HIDDEN_DIM * FRD_HIDDEN_DIM)?; + let mut grad_b1_pb_d = stream.alloc_zeros::(b_size * FRD_HIDDEN_DIM)?; + let mut grad_h_t_d = stream.alloc_zeros::(b_size * HIDDEN_DIM)?; + head.layer1_bwd( + &h_t_d, + &hidden_d, + &grad_hidden_d, + &mut grad_w1_pb_d, + &mut grad_b1_pb_d, + &mut grad_h_t_d, + b_size, + )?; + + // Finite-diff probe: perturb W1[k=30, i=20] by ±ε. Pick an + // (k, i) likely to land on an active ReLU branch (use a moderate + // i within the 64 hidden slots — random init usually has ~50% + // active so any pick has ~50% chance of being a non-mask slot). + let probe_k = 30_usize; + let probe_i = 20_usize; + let probe_off = probe_k * FRD_HIDDEN_DIM + probe_i; + let eps = 1e-3_f32; + + let grad_w1 = read_slice_d_pub(&stream, &grad_w1_pb_d, HIDDEN_DIM * FRD_HIDDEN_DIM)?; + // Pick the slot with MAX |analytical| — fp32 finite-diff rel_err + // is dominated by absolute rounding noise (~1e-5 at ε=1e-3 across + // the deep matmul chain), so a tiny-magnitude gradient produces + // misleading relative error. Probing the largest gradient gives + // the cleanest sanity check. + let _ = (probe_k, probe_i, probe_off); + let (chosen_off, &chosen_analytical) = grad_w1 + .iter() + .enumerate() + .max_by(|(_, a), (_, b)| a.abs().partial_cmp(&b.abs()).unwrap()) + .expect("non-empty grad_w1"); + let chosen_k = chosen_off / FRD_HIDDEN_DIM; + let chosen_i = chosen_off % FRD_HIDDEN_DIM; + assert!( + chosen_analytical.abs() > 1e-3, + "max |dW1| should be at least 1e-3 for a meaningful finite-diff; got {chosen_analytical}" + ); + + let mut w1_host = read_slice_d_pub(&stream, &head.w1_d, HIDDEN_DIM * FRD_HIDDEN_DIM)?; + let original = w1_host[chosen_off]; + + // L(W1 + ε · e_(k,i)) + w1_host[chosen_off] = original + eps; + write_slice_f32_d_pub(&stream, &w1_host, &mut head.w1_d)?; + let mut hidden_plus = stream.alloc_zeros::(b_size * FRD_HIDDEN_DIM)?; + let mut logits_plus = stream.alloc_zeros::(b_size * FRD_OUT_DIM)?; + head.forward(&h_t_d, &mut hidden_plus, &mut logits_plus, b_size)?; + let logits_plus_h = read_slice_d_pub(&stream, &logits_plus, b_size * FRD_OUT_DIM)?; + let l_plus = ce_total_loss(&head, &stream, &logits_plus_h, &labels_d, b_size)?; + + // L(W1 - ε · e_(k,i)) + w1_host[chosen_off] = original - eps; + write_slice_f32_d_pub(&stream, &w1_host, &mut head.w1_d)?; + let mut hidden_minus = stream.alloc_zeros::(b_size * FRD_HIDDEN_DIM)?; + let mut logits_minus = stream.alloc_zeros::(b_size * FRD_OUT_DIM)?; + head.forward(&h_t_d, &mut hidden_minus, &mut logits_minus, b_size)?; + let logits_minus_h = read_slice_d_pub(&stream, &logits_minus, b_size * FRD_OUT_DIM)?; + let l_minus = ce_total_loss(&head, &stream, &logits_minus_h, &labels_d, b_size)?; + + // Restore W1. + w1_host[chosen_off] = original; + write_slice_f32_d_pub(&stream, &w1_host, &mut head.w1_d)?; + + let numerical = (l_plus - l_minus) / (2.0 * eps); + let rel_err = (numerical - chosen_analytical).abs() / (chosen_analytical.abs().max(1e-6)); + // 1e-2 tolerance is slightly looser than dW2's 5e-3 because the + // dW1 chain crosses an extra matmul + the ReLU mask boundary + // (which is a non-differentiable point — at hidden ≈ 0 the + // finite-diff straddles the mask discontinuity). + assert!( + rel_err < 1e-2, + "dW1 finite-diff mismatch at (k={chosen_k}, i={chosen_i}): \ + analytical={chosen_analytical:.6}, numerical={numerical:.6}, rel_err={rel_err:.6}" + ); + + eprintln!( + "PASS — dW1 finite-diff at (k={chosen_k}, i={chosen_i}): \ + analytical={:.6} numerical={:.6} rel_err={:.2e}", + chosen_analytical, numerical, rel_err + ); + Ok(()) +} + +#[test] +#[ignore = "requires CUDA (MlDevice::cuda(0))"] +fn frd_layer1_bwd_relu_mask_zeros_grad() -> Result<()> { + // Invariant: if cached `hidden[b, i] == 0` (i.e. the pre-activation + // landed on the negative half of ReLU), then: + // * grad_b1_per_batch[b, i] = 0 (db1 = dpre_hidden = grad_hidden × mask) + // * grad_W1_per_batch[b, k, i] = 0 for every k + // + // Construct a fixture where some hidden slots are guaranteed + // negative: feed h_t = all -1.0 (so half of Xavier-init W1's + // products average negative). + let Some((dev, head)) = build_head() else { return Ok(()) }; + let stream = dev.cuda_stream()?.clone(); + let b_size = 1; + + let h_t = vec![-1.0_f32; b_size * HIDDEN_DIM]; + let h_t_d = upload_f32(&stream, &h_t)?; + let labels: Vec = vec![5, 10, 15]; + let labels_d = upload_i32(&stream, &labels)?; + + let mut hidden_d = stream.alloc_zeros::(b_size * FRD_HIDDEN_DIM)?; + let mut logits_d = stream.alloc_zeros::(b_size * FRD_OUT_DIM)?; + head.forward(&h_t_d, &mut hidden_d, &mut logits_d, b_size)?; + + let mut grad_logits_d = stream.alloc_zeros::(b_size * FRD_OUT_DIM)?; + let mut loss_d = stream.alloc_zeros::(b_size * FRD_N_HORIZONS)?; + head.softmax_ce_grad(&logits_d, &labels_d, &mut grad_logits_d, &mut loss_d, b_size)?; + + let mut grad_w2_pb_d = stream.alloc_zeros::(b_size * FRD_HIDDEN_DIM * FRD_OUT_DIM)?; + let mut grad_b2_pb_d = stream.alloc_zeros::(b_size * FRD_OUT_DIM)?; + let mut grad_hidden_d = stream.alloc_zeros::(b_size * FRD_HIDDEN_DIM)?; + head.layer2_bwd( + &hidden_d, + &grad_logits_d, + &mut grad_w2_pb_d, + &mut grad_b2_pb_d, + &mut grad_hidden_d, + b_size, + )?; + + let mut grad_w1_pb_d = stream.alloc_zeros::(b_size * HIDDEN_DIM * FRD_HIDDEN_DIM)?; + let mut grad_b1_pb_d = stream.alloc_zeros::(b_size * FRD_HIDDEN_DIM)?; + let mut grad_h_t_d = stream.alloc_zeros::(b_size * HIDDEN_DIM)?; + head.layer1_bwd( + &h_t_d, + &hidden_d, + &grad_hidden_d, + &mut grad_w1_pb_d, + &mut grad_b1_pb_d, + &mut grad_h_t_d, + b_size, + )?; + + let hidden = read_slice_d_pub(&stream, &hidden_d, b_size * FRD_HIDDEN_DIM)?; + let grad_b1 = read_slice_d_pub(&stream, &grad_b1_pb_d, b_size * FRD_HIDDEN_DIM)?; + let grad_w1 = read_slice_d_pub(&stream, &grad_w1_pb_d, b_size * HIDDEN_DIM * FRD_HIDDEN_DIM)?; + + let mut masked_count = 0; + let mut unmasked_count = 0; + for i in 0..FRD_HIDDEN_DIM { + if hidden[i] == 0.0 { + masked_count += 1; + // ReLU mask zeros dpre_hidden → db1 + entire column of dW1 must be 0. + assert_eq!( + grad_b1[i], 0.0, + "db1[{i}] should be 0 under ReLU mask (hidden=0); got {}", + grad_b1[i] + ); + for k in 0..HIDDEN_DIM { + let off = k * FRD_HIDDEN_DIM + i; + assert_eq!( + grad_w1[off], 0.0, + "dW1[{k}, {i}] should be 0 under ReLU mask (hidden=0); got {}", + grad_w1[off] + ); + } + } else { + unmasked_count += 1; + } + } + assert!( + masked_count > 0, + "expected at least one ReLU-masked hidden slot under h_t = all -1; \ + got 0 (all-positive hidden — Xavier init may have flipped signs)" + ); + + eprintln!( + "PASS — ReLU mask zeros dW1 + db1 for {} of {} hidden slots ({} active)", + masked_count, FRD_HIDDEN_DIM, unmasked_count + ); + Ok(()) +}