diff --git a/crates/ml-alpha/build.rs b/crates/ml-alpha/build.rs index d91972b2d..81122e5f6 100644 --- a/crates/ml-alpha/build.rs +++ b/crates/ml-alpha/build.rs @@ -77,6 +77,7 @@ const KERNELS: &[&str] = &[ "rl_trail_stop_check", // SP20 P1+P5 audit fix: per-unit trail breach check; OVERRIDE action to FlatFromLong/Short on breach (close routes through existing flat plumbing) "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 ]; // Cache bust v31 — five new reduce / derive kernels populate the input diff --git a/crates/ml-alpha/cuda/rl_frd_layer2_bwd.cu b/crates/ml-alpha/cuda/rl_frd_layer2_bwd.cu new file mode 100644 index 000000000..1b2c47788 --- /dev/null +++ b/crates/ml-alpha/cuda/rl_frd_layer2_bwd.cu @@ -0,0 +1,80 @@ +// rl_frd_layer2_bwd.cu — FRD head backward stage 2 (SP20 P3 F.3b). +// +// Given dL/dlogits [B, FRD_OUT_DIM] (produced by rl_frd_softmax_ce_grad), +// compute the layer-2 gradients via the chain rule: +// +// logits[b, j] = Σ_i hidden[b, i] × W2[i, j] + b2[j] +// +// dL/dW2[i, j] = Σ_b hidden[b, i] × dL/dlogits[b, j] ← reduce over B +// dL/db2[j] = Σ_b dL/dlogits[b, j] ← reduce over B +// dL/dhidden[b, i] = Σ_j W2[i, j] × dL/dlogits[b, j] ← per-batch overwrite +// +// Per-batch scratch buffers (caller reduces across batch via the +// existing reduce_axis0 infra, same pattern as v_head_bwd): +// grad_W2_per_batch[b, i, j] = hidden[b, i] × dL/dlogits[b, j] +// grad_b2_per_batch[b, j] = dL/dlogits[b, j] +// +// Block layout: 1 block per batch, FRD_HIDDEN_DIM=64 threads. +// * Thread i (i < 64) computes dL/dhidden[b, i] (full inner product +// across all 63 output slots — sequential at 63 reads per thread). +// * Thread i also writes grad_W2_per_batch[b, i, 0..63] (63 writes +// per thread, no cross-thread coordination). +// * Threads 0..62 also write grad_b2_per_batch[b, tid] in a single +// pass (one slot per thread, thread 63 idle for that step). +// +// Per `feedback_no_atomicadd`: per-(b, i, j) sole-writer pattern. +// Per `feedback_cpu_is_read_only`: pure device-side. + +#include + +#define FRD_HIDDEN_DIM 64 +#define FRD_N_HORIZONS 3 +#define FRD_N_ATOMS 21 +#define FRD_OUT_DIM (FRD_N_HORIZONS * FRD_N_ATOMS) + +extern "C" __global__ void rl_frd_layer2_bwd( + const float* __restrict__ hidden, // [B, FRD_HIDDEN_DIM] + const float* __restrict__ grad_logits, // [B, FRD_OUT_DIM] + const float* __restrict__ w2, // [FRD_HIDDEN_DIM, FRD_OUT_DIM] + int b_size, + float* __restrict__ grad_w2_per_batch, // [B, FRD_HIDDEN_DIM, FRD_OUT_DIM] + float* __restrict__ grad_b2_per_batch, // [B, FRD_OUT_DIM] + float* __restrict__ grad_hidden // [B, FRD_HIDDEN_DIM] +) { + const int b = blockIdx.x; + const int i = threadIdx.x; + if (b >= b_size || i >= FRD_HIDDEN_DIM) return; + + // Stage grad_logits into shared so the inner loop reads from + // shared, not global. + __shared__ float s_grad_logits[FRD_OUT_DIM]; + // FRD_OUT_DIM = 63, FRD_HIDDEN_DIM = 64 threads → each thread + // covers one slot, thread 63 idle for staging. + if (i < FRD_OUT_DIM) { + s_grad_logits[i] = grad_logits[b * FRD_OUT_DIM + i]; + } + __syncthreads(); + + const float h_bi = hidden[b * FRD_HIDDEN_DIM + i]; + + // grad_W2_per_batch[b, i, j] = h_bi × s_grad_logits[j] + const int row_off = (b * FRD_HIDDEN_DIM + i) * FRD_OUT_DIM; + #pragma unroll + for (int j = 0; j < FRD_OUT_DIM; ++j) { + grad_w2_per_batch[row_off + j] = h_bi * s_grad_logits[j]; + } + + // grad_hidden[b, i] = Σ_j W2[i, j] × s_grad_logits[j] + float acc = 0.0f; + #pragma unroll + for (int j = 0; j < FRD_OUT_DIM; ++j) { + acc += w2[i * FRD_OUT_DIM + j] * s_grad_logits[j]; + } + grad_hidden[b * FRD_HIDDEN_DIM + i] = acc; + + // grad_b2_per_batch[b, j] = s_grad_logits[j] — sole writer per + // (b, j); thread i writes slot j == i (covers j < 63 only). + if (i < FRD_OUT_DIM) { + grad_b2_per_batch[b * FRD_OUT_DIM + i] = s_grad_logits[i]; + } +} diff --git a/crates/ml-alpha/src/rl/frd.rs b/crates/ml-alpha/src/rl/frd.rs index 89e8cb0db..157f50e40 100644 --- a/crates/ml-alpha/src/rl/frd.rs +++ b/crates/ml-alpha/src/rl/frd.rs @@ -60,6 +60,8 @@ use crate::trainer::integrated::write_slice_f32_d_pub; const FRD_FWD_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/rl_frd_fwd.cubin")); 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")); /// Per-batch output width: `FRD_N_HORIZONS × FRD_N_ATOMS`. Each batch /// row contains 3 contiguous horizon blocks of 21 atom logits. @@ -89,6 +91,8 @@ pub struct FrdHead { fwd_fn: CudaFunction, _softmax_ce_grad_module: Arc, softmax_ce_grad_fn: CudaFunction, + _layer2_bwd_module: Arc, + layer2_bwd_fn: CudaFunction, pub w1_d: CudaSlice, // [HIDDEN_DIM, FRD_HIDDEN_DIM] pub b1_d: CudaSlice, // [FRD_HIDDEN_DIM] @@ -112,6 +116,12 @@ impl FrdHead { let softmax_ce_grad_fn = softmax_ce_grad_module .load_function("rl_frd_softmax_ce_grad") .context("load rl_frd_softmax_ce_grad fn")?; + let layer2_bwd_module = ctx + .load_cubin(FRD_LAYER2_BWD_CUBIN.to_vec()) + .context("load rl_frd_layer2_bwd cubin")?; + let layer2_bwd_fn = layer2_bwd_module + .load_function("rl_frd_layer2_bwd") + .context("load rl_frd_layer2_bwd fn")?; // Per pearl_scoped_init_seed_for_reproducibility — guard around // Xavier draws so GPU init helpers downstream see the same RNG. @@ -147,6 +157,8 @@ impl FrdHead { fwd_fn, _softmax_ce_grad_module: softmax_ce_grad_module, softmax_ce_grad_fn, + _layer2_bwd_module: layer2_bwd_module, + layer2_bwd_fn, w1_d, b1_d, w2_d, @@ -198,6 +210,56 @@ impl FrdHead { Ok(()) } + /// Backward stage 2 (F.3b): layer-2 weight gradients. + /// + /// Given `hidden` (cached from forward) and `grad_logits` (from + /// `softmax_ce_grad`), computes: + /// * `grad_w2_per_batch[b, i, j] = hidden[b, i] × grad_logits[b, j]` + /// * `grad_b2_per_batch[b, j] = grad_logits[b, j]` + /// * `grad_hidden[b, i] = Σ_j W2[i, j] × grad_logits[b, j]` + /// + /// Per-batch scratch buffers feed `reduce_axis0` to produce the + /// final accumulated weight grads. `grad_hidden` is the upstream + /// gradient for layer-1 backward (F.3c). + #[allow(clippy::too_many_arguments)] + pub fn layer2_bwd( + &self, + hidden_d: &CudaSlice, + grad_logits_d: &CudaSlice, + grad_w2_per_batch_d: &mut CudaSlice, + grad_b2_per_batch_d: &mut CudaSlice, + grad_hidden_d: &mut CudaSlice, + b_size: usize, + ) -> Result<()> { + debug_assert_eq!(hidden_d.len(), b_size * FRD_HIDDEN_DIM); + debug_assert_eq!(grad_logits_d.len(), b_size * FRD_OUT_DIM); + debug_assert_eq!( + grad_w2_per_batch_d.len(), + b_size * FRD_HIDDEN_DIM * FRD_OUT_DIM + ); + debug_assert_eq!(grad_b2_per_batch_d.len(), b_size * FRD_OUT_DIM); + debug_assert_eq!(grad_hidden_d.len(), b_size * FRD_HIDDEN_DIM); + let cfg = LaunchConfig { + grid_dim: (b_size as u32, 1, 1), + block_dim: (FRD_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.layer2_bwd_fn); + launch + .arg(hidden_d) + .arg(grad_logits_d) + .arg(&self.w2_d) + .arg(&b_size_i) + .arg(grad_w2_per_batch_d) + .arg(grad_b2_per_batch_d) + .arg(grad_hidden_d); + unsafe { + launch.launch(cfg).context("rl_frd_layer2_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 bfb119fa3..1ad9744ec 100644 --- a/crates/ml-alpha/tests/frd_head.rs +++ b/crates/ml-alpha/tests/frd_head.rs @@ -352,3 +352,169 @@ fn frd_softmax_ce_grad_finite_diff_matches_analytical() -> Result<()> { ); Ok(()) } + +/// Helper: compute total loss for given logits by calling the +/// softmax_ce_grad kernel and summing per-(b, h) CE. +fn ce_total_loss( + head: &FrdHead, + stream: &Arc, + logits: &[f32], + labels_d: &CudaSlice, + b_size: usize, +) -> Result { + let logits_d = upload_f32(stream, logits)?; + let mut grad_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_d, &mut loss_d, b_size)?; + let loss = read_slice_d_pub(stream, &loss_d, b_size * FRD_N_HORIZONS)?; + Ok(loss.iter().sum()) +} + +#[test] +#[ignore = "requires CUDA (MlDevice::cuda(0))"] +fn frd_layer2_bwd_finite_diff_w2() -> Result<()> { + let Some((dev, mut head)) = build_head() else { return Ok(()) }; + let stream = dev.cuda_stream()?.clone(); + let b_size = 1; + + // Forward pass with random h_t to produce hidden + logits. + use rand::{Rng, SeedableRng}; + use rand_chacha::ChaCha8Rng; + let mut r = ChaCha8Rng::seed_from_u64(0xB2); + 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 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 labels: Vec = vec![5, 10, 15]; + let labels_d = upload_i32(&stream, &labels)?; + + // Softmax+CE grad of logits. + 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)?; + + // Layer-2 backward: produce per-batch grad_W2 scratch. + 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, + )?; + + // For b_size=1 the per-batch scratch IS the reduced gradient. + let grad_w2 = read_slice_d_pub(&stream, &grad_w2_pb_d, FRD_HIDDEN_DIM * FRD_OUT_DIM)?; + + // Finite-diff probe: perturb W2[i=10, j=5] by ±ε, measure ΔL. + let probe_i = 10_usize; + let probe_j = 5_usize; + let probe_off = probe_i * FRD_OUT_DIM + probe_j; + let eps = 1e-3_f32; + + let mut w2_host = read_slice_d_pub(&stream, &head.w2_d, FRD_HIDDEN_DIM * FRD_OUT_DIM)?; + let original = w2_host[probe_off]; + + // L(W2 + ε · e_(i,j)) — perturb weight on device, re-run forward, + // compute total CE loss across all (b, h) rows. + w2_host[probe_off] = original + eps; + write_slice_f32_d_pub(&stream, &w2_host, &mut head.w2_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(W2 - ε · e_(i,j)) + w2_host[probe_off] = original - eps; + write_slice_f32_d_pub(&stream, &w2_host, &mut head.w2_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 W2 to keep test isolation clean (next test gets default init). + w2_host[probe_off] = original; + write_slice_f32_d_pub(&stream, &w2_host, &mut head.w2_d)?; + + let numerical = (l_plus - l_minus) / (2.0 * eps); + let analytical = grad_w2[probe_off]; + let rel_err = (numerical - analytical).abs() / (analytical.abs().max(1e-6)); + assert!( + rel_err < 5e-3, + "dW2 finite-diff mismatch at (i={probe_i}, j={probe_j}): \ + analytical={analytical:.6}, numerical={numerical:.6}, rel_err={rel_err:.6}" + ); + + eprintln!( + "PASS — dW2 finite-diff: analytical={:.6} numerical={:.6} rel_err={:.2e}", + analytical, numerical, rel_err + ); + Ok(()) +} + +#[test] +#[ignore = "requires CUDA (MlDevice::cuda(0))"] +fn frd_layer2_bwd_db2_equals_grad_logits() -> Result<()> { + // db2 invariant: per-batch grad_b2[b, j] = grad_logits[b, j]. + // After reduce_axis0 across batch this becomes Σ_b grad_logits[b, j] + // (the standard bias gradient). Verify the per-batch scratch + // matches the input grad_logits exactly. + let Some((dev, head)) = build_head() else { return Ok(()) }; + let stream = dev.cuda_stream()?.clone(); + let b_size = 4; + + use rand::{Rng, SeedableRng}; + use rand_chacha::ChaCha8Rng; + let mut r = ChaCha8Rng::seed_from_u64(0xB22); + 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 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 labels: Vec = (0..b_size * FRD_N_HORIZONS) + .map(|i| ((i * 3) % FRD_N_ATOMS) as i32) + .collect(); + let labels_d = upload_i32(&stream, &labels)?; + + 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 grad_logits = read_slice_d_pub(&stream, &grad_logits_d, b_size * FRD_OUT_DIM)?; + let grad_b2_pb = read_slice_d_pub(&stream, &grad_b2_pb_d, b_size * FRD_OUT_DIM)?; + for (i, (gl, gb)) in grad_logits.iter().zip(grad_b2_pb.iter()).enumerate() { + assert_eq!(*gl, *gb, "db2_per_batch[{i}] must equal grad_logits[{i}]"); + } + + eprintln!( + "PASS — db2_per_batch ≡ grad_logits across {} slots", + grad_logits.len() + ); + Ok(()) +}