From 6cfd7e6691fe4f838bb9fd0f467e2a7d6ea2e9cb Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Sun, 24 May 2026 18:31:03 +0200 Subject: [PATCH] feat(rl): FRD softmax + CE + dL/dlogits backward stage 1 (F.3a) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Per-(batch, horizon) softmax + cross-entropy loss + gradient w.r.t. the 21 atom logits. First of three backward stages — F.3b adds layer-2 weight grads (dW2, db2, dhidden), F.3c adds layer-1 weight grads (dW1, db1, dh_t with ReLU mask). Kernel `cuda/rl_frd_softmax_ce_grad.cu`: * grid_dim = (B, FRD_N_HORIZONS, 1), block_dim = (FRD_N_ATOMS=21, 1, 1) — one block per (batch, horizon) pair, threads cooperate over the 21 atoms via shared mem * Standard numerically-stable softmax: shift by row_max, exponentiate, normalize by row_sum (thread 0 does the serial reductions — 21 atoms is small enough warp-shuffle overhead isn't worth it) * Gradient: (p[a] - 1{a==label}) / B at the source per v_head_bwd convention (mean-reduce over batch) * Loss: -log(p[label]) with 1e-30 floor against log(0) * Sentinel label (-1) zeros both gradient row and loss — for the missing-horizon case at the rightmost edge of the snapshot stream (forward returns at h=300 ticks aren't realized for the last 300 snapshots; loader marks those labels with -1) * Per feedback_no_atomicadd: per-(b, h, a) sole-writer pattern Rust wiring `src/rl/frd.rs::FrdHead::softmax_ce_grad`: * Second cubin loaded alongside fwd (separate module per the aux_heads pattern; small handle, no impact on init time) * Caller provides labels_d [B, FRD_N_HORIZONS] of i32 and gets back grad_logits + per-(b, h) raw CE; sum + λ_frd scaling left to the caller (F.4 will hook this into stats.l_total + Adam step) Tests `tests/frd_head.rs` — 3 new GPU-oracle tests (6/6 file total), all PASS on RTX 3050 Ti: 1. frd_softmax_ce_grad_uniform_logits_match_log_n_atoms — for any label, uniform logits → CE = ln(FRD_N_ATOMS) = ln(21) ≈ 3.0445. Also asserts per-row Σ grad_logits = 0 (softmax-CE invariant). 2. frd_softmax_ce_grad_sentinel_label_zeros_row — label=-1 with non-trivial random logits produces exactly zero loss + grad for every row (no leak through the sentinel path). 3. frd_softmax_ce_grad_finite_diff_matches_analytical — perturbs one logit slot by ±ε=1e-3, compares (L(+ε) - L(-ε))/(2ε) to the kernel's analytical gradient. rel_err ≈ 1.3e-3 (fp32 finite-diff is rounding-error-limited at this ε; tolerance set to 5e-3 with explanatory comment). The first two tests provide strong analytical oracles (no CPU reference impl per feedback_no_cpu_test_fallbacks). The finite-diff test cross-validates the full softmax+CE chain via a numerical gradient — the standard ground-truth for autodiff kernels. --- crates/ml-alpha/build.rs | 1 + .../ml-alpha/cuda/rl_frd_softmax_ce_grad.cu | 115 ++++++++++++ crates/ml-alpha/src/rl/frd.rs | 62 ++++++- crates/ml-alpha/tests/frd_head.rs | 167 +++++++++++++++++- 4 files changed, 342 insertions(+), 3 deletions(-) create mode 100644 crates/ml-alpha/cuda/rl_frd_softmax_ce_grad.cu diff --git a/crates/ml-alpha/build.rs b/crates/ml-alpha/build.rs index 8798c0a83..d91972b2d 100644 --- a/crates/ml-alpha/build.rs +++ b/crates/ml-alpha/build.rs @@ -76,6 +76,7 @@ const KERNELS: &[&str] = &[ "rl_trail_mutate", // SP20 P1+P5 audit fix (a7/a8 dead): TrailTighten/TrailLoosen mutate unit_trail_distance bounded MIN/MAX, symmetric reciprocal adjust "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 ]; // Cache bust v31 — five new reduce / derive kernels populate the input diff --git a/crates/ml-alpha/cuda/rl_frd_softmax_ce_grad.cu b/crates/ml-alpha/cuda/rl_frd_softmax_ce_grad.cu new file mode 100644 index 000000000..ae65b22cd --- /dev/null +++ b/crates/ml-alpha/cuda/rl_frd_softmax_ce_grad.cu @@ -0,0 +1,115 @@ +// rl_frd_softmax_ce_grad.cu — FRD head per-horizon softmax + CE loss + +// gradient w.r.t. logits (SP20 P3 backward stage 1). +// +// Per batch (b) and horizon (h), the 21 atom logits feed a softmax: +// p[b, h, a] = exp(logits[b, h, a] - max) / Σ_a' exp(logits[b, h, a'] - max) +// +// Cross-entropy against one-hot label: +// CE[b, h] = -log p[b, h, label[b, h]] +// +// Gradient (standard softmax-CE derivative, mean-reduced over batch +// per the v_head_bwd convention so downstream reduce_axis0 + Adam +// step land mean-grad without an explicit /B): +// grad_logits[b, h, a] = (1/B) × (p[b, h, a] - 1{a == label[b, h]}) +// +// Label sentinel: a negative label index (e.g. -1) means "missing +// horizon" (future return not yet realized at the rightmost edge of +// the snapshot stream); the kernel zeros grad_logits for that row +// and writes loss = 0 so the caller's sum is masked-out. +// +// Block layout: 1 block per (batch, horizon), 21 threads per block. +// * Cooperative max-reduction across 21 atoms (thread 0 serial in +// shared mem — 21 atoms is small enough that the warp-shuffle +// pattern's overhead isn't worth it) +// * Cooperative sum-exp across 21 atoms (same pattern) +// * Each thread writes one grad_logits slot +// * Thread 0 writes loss_per_batch_h[b, h] +// +// Per `feedback_no_atomicadd`: no atomics — per-(batch, horizon, atom) +// sole-writer pattern. +// Per `feedback_no_nvrtc`: pre-compiled cubin via build.rs. + +#include + +#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_softmax_ce_grad( + const float* __restrict__ logits, // [B, FRD_OUT_DIM] + const int* __restrict__ labels, // [B, FRD_N_HORIZONS] + int b_size, + float* __restrict__ grad_logits, // [B, FRD_OUT_DIM] OUT + float* __restrict__ loss_per_batch_h // [B, FRD_N_HORIZONS] OUT +) { + const int b = blockIdx.x; + const int h = blockIdx.y; + const int a = threadIdx.x; + if (b >= b_size || h >= FRD_N_HORIZONS || a >= FRD_N_ATOMS) return; + + const int row_off = b * FRD_OUT_DIM + h * FRD_N_ATOMS; + + __shared__ float exp_buf[FRD_N_ATOMS]; + __shared__ float row_max; + __shared__ float row_sum; + __shared__ int s_label; + + // Phase 0: stage label (thread 0) — same value broadcast to all threads. + if (a == 0) { + s_label = labels[b * FRD_N_HORIZONS + h]; + } + __syncthreads(); + + const int label = s_label; + + // Sentinel label → emit zero gradient + zero loss for this row. + // (Caller masks out missing-horizon contributions from the total + // before backward propagation continues.) + if (label < 0 || label >= FRD_N_ATOMS) { + grad_logits[row_off + a] = 0.0f; + if (a == 0) { + loss_per_batch_h[b * FRD_N_HORIZONS + h] = 0.0f; + } + return; + } + + // Phase 1: stage logits + compute max (thread 0 serial). + const float my_logit = logits[row_off + a]; + exp_buf[a] = my_logit; + __syncthreads(); + if (a == 0) { + float m = exp_buf[0]; + #pragma unroll + for (int i = 1; i < FRD_N_ATOMS; ++i) { + m = fmaxf(m, exp_buf[i]); + } + row_max = m; + } + __syncthreads(); + + // Phase 2: each thread computes its exp(logit - max), stages. + const float my_exp = expf(my_logit - row_max); + exp_buf[a] = my_exp; + __syncthreads(); + if (a == 0) { + float s = 0.0f; + #pragma unroll + for (int i = 0; i < FRD_N_ATOMS; ++i) { + s += exp_buf[i]; + } + row_sum = s; + } + __syncthreads(); + + // Phase 3: per-atom softmax probability + gradient. + const float p_a = my_exp / row_sum; + const float target = (a == label) ? 1.0f : 0.0f; + const float inv_B = 1.0f / (float) b_size; + grad_logits[row_off + a] = (p_a - target) * inv_B; + + // Phase 4: loss (thread 0 only). p_label = exp_buf[label] / row_sum. + if (a == 0) { + const float p_label = exp_buf[label] / row_sum; + loss_per_batch_h[b * FRD_N_HORIZONS + h] = -logf(fmaxf(p_label, 1e-30f)); + } +} diff --git a/crates/ml-alpha/src/rl/frd.rs b/crates/ml-alpha/src/rl/frd.rs index 2b96545ef..89e8cb0db 100644 --- a/crates/ml-alpha/src/rl/frd.rs +++ b/crates/ml-alpha/src/rl/frd.rs @@ -58,6 +58,8 @@ use crate::rl::common::{FRD_HIDDEN_DIM, FRD_N_ATOMS, FRD_N_HORIZONS}; 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")); /// Per-batch output width: `FRD_N_HORIZONS × FRD_N_ATOMS`. Each batch /// row contains 3 contiguous horizon blocks of 21 atom logits. @@ -77,12 +79,16 @@ impl Default for FrdHeadConfig { } } -/// FRD head — owns device weights for the 2-layer MLP + cached forward -/// cubin function. Backward kernel is added in F.3. +/// FRD head — owns device weights for the 2-layer MLP + cached cubin +/// function handles. Forward (F.2) ships separately from backward +/// (F.3a/b/c) so each stage has its own GPU-oracle test before +/// composition. pub struct FrdHead { stream: Arc, _module: Arc, fwd_fn: CudaFunction, + _softmax_ce_grad_module: Arc, + softmax_ce_grad_fn: CudaFunction, pub w1_d: CudaSlice, // [HIDDEN_DIM, FRD_HIDDEN_DIM] pub b1_d: CudaSlice, // [FRD_HIDDEN_DIM] @@ -100,6 +106,12 @@ impl FrdHead { let fwd_fn = module .load_function("rl_frd_fwd") .context("load rl_frd_fwd fn")?; + let softmax_ce_grad_module = ctx + .load_cubin(FRD_SOFTMAX_CE_GRAD_CUBIN.to_vec()) + .context("load rl_frd_softmax_ce_grad cubin")?; + let softmax_ce_grad_fn = softmax_ce_grad_module + .load_function("rl_frd_softmax_ce_grad") + .context("load rl_frd_softmax_ce_grad fn")?; // Per pearl_scoped_init_seed_for_reproducibility — guard around // Xavier draws so GPU init helpers downstream see the same RNG. @@ -133,6 +145,8 @@ impl FrdHead { stream, _module: module, fwd_fn, + _softmax_ce_grad_module: softmax_ce_grad_module, + softmax_ce_grad_fn, w1_d, b1_d, w2_d, @@ -140,6 +154,50 @@ impl FrdHead { }) } + /// Backward stage 1 (F.3a): per-(batch, horizon) softmax + CE loss + /// + gradient w.r.t. logits. Pairs with the forward kernel's + /// `logits_out_d`. The `labels_d` buffer is `[B, FRD_N_HORIZONS]` + /// of i32 — pass `-1` as the sentinel for missing horizons (the + /// kernel zeros that row's gradient + loss). + /// + /// Outputs: + /// * `grad_logits_d [B, FRD_OUT_DIM]` — mean-reduced over batch + /// (multiplied by `1/B` at the source per v_head_bwd convention) + /// * `loss_per_batch_h_d [B, FRD_N_HORIZONS]` — raw CE per row; + /// caller sums + scales by `λ_frd` before backward continues + pub fn softmax_ce_grad( + &self, + logits_d: &CudaSlice, + labels_d: &CudaSlice, + grad_logits_d: &mut CudaSlice, + loss_per_batch_h_d: &mut CudaSlice, + b_size: usize, + ) -> Result<()> { + debug_assert_eq!(logits_d.len(), b_size * FRD_OUT_DIM); + debug_assert_eq!(labels_d.len(), b_size * FRD_N_HORIZONS); + debug_assert_eq!(grad_logits_d.len(), b_size * FRD_OUT_DIM); + debug_assert_eq!(loss_per_batch_h_d.len(), b_size * FRD_N_HORIZONS); + let cfg = LaunchConfig { + grid_dim: (b_size as u32, FRD_N_HORIZONS as u32, 1), + block_dim: (FRD_N_ATOMS as u32, 1, 1), + shared_mem_bytes: 0, + }; + let b_size_i = b_size as i32; + let mut launch = self.stream.launch_builder(&self.softmax_ce_grad_fn); + launch + .arg(logits_d) + .arg(labels_d) + .arg(&b_size_i) + .arg(grad_logits_d) + .arg(loss_per_batch_h_d); + unsafe { + launch + .launch(cfg) + .context("rl_frd_softmax_ce_grad 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 a14dcb341..bfb119fa3 100644 --- a/crates/ml-alpha/tests/frd_head.rs +++ b/crates/ml-alpha/tests/frd_head.rs @@ -30,7 +30,7 @@ use ml_alpha::heads::HIDDEN_DIM; use ml_alpha::rl::common::{FRD_HIDDEN_DIM, FRD_N_ATOMS, FRD_N_HORIZONS}; use ml_alpha::rl::frd::{FrdHead, FrdHeadConfig, FRD_OUT_DIM}; use ml_alpha::trainer::integrated::{ - read_slice_d_pub, write_slice_f32_d_pub, + read_slice_d_pub, write_slice_f32_d_pub, write_slice_i32_d_pub, }; use ml_core::device::MlDevice; use std::sync::Arc; @@ -187,3 +187,168 @@ fn frd_forward_relu_mask_consistent_with_cached_hidden() -> Result<()> { ); Ok(()) } + +fn upload_i32(stream: &Arc, host: &[i32]) -> Result> { + let mut d = stream.alloc_zeros::(host.len())?; + write_slice_i32_d_pub(stream, host, &mut d)?; + Ok(d) +} + +#[test] +#[ignore = "requires CUDA (MlDevice::cuda(0))"] +fn frd_softmax_ce_grad_uniform_logits_match_log_n_atoms() -> Result<()> { + let Some((dev, head)) = build_head() else { return Ok(()) }; + let stream = dev.cuda_stream()?.clone(); + let b_size = 4; + + // Uniform logits (all zeros) → softmax = 1/FRD_N_ATOMS everywhere, + // so CE = -log(1/21) = ln(21) ≈ 3.0445 for ANY label index. + let logits = vec![0.0_f32; b_size * FRD_OUT_DIM]; + let logits_d = upload_f32(&stream, &logits)?; + // Set label = 10 (mid-bucket) for every (batch, horizon) row. + let labels: Vec = vec![10; b_size * FRD_N_HORIZONS]; + let labels_d = upload_i32(&stream, &labels)?; + 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)?; + let expected = (FRD_N_ATOMS as f32).ln(); + for (i, v) in loss.iter().enumerate() { + assert!( + (v - expected).abs() < 1e-5, + "uniform-logit CE[{i}] should be ln({}) = {expected}; got {v}", + FRD_N_ATOMS + ); + } + + // Softmax-CE gradient invariant: Σ_a (p[a] - 1{a==label}) = 0 + // (mean-reduced over B doesn't change the per-row sum being zero). + let grad = read_slice_d_pub(&stream, &grad_d, b_size * FRD_OUT_DIM)?; + for b in 0..b_size { + for h in 0..FRD_N_HORIZONS { + let off = b * FRD_OUT_DIM + h * FRD_N_ATOMS; + let row_sum: f32 = grad[off..off + FRD_N_ATOMS].iter().sum(); + assert!( + row_sum.abs() < 1e-6, + "Σ grad_logits per (b={b}, h={h}) must be 0 (softmax-CE invariant); got {row_sum}" + ); + } + } + + eprintln!( + "PASS — uniform logits → CE = ln({}) = {:.4}; per-row grad sum = 0", + FRD_N_ATOMS, expected + ); + Ok(()) +} + +#[test] +#[ignore = "requires CUDA (MlDevice::cuda(0))"] +fn frd_softmax_ce_grad_sentinel_label_zeros_row() -> Result<()> { + let Some((dev, head)) = build_head() else { return Ok(()) }; + let stream = dev.cuda_stream()?.clone(); + let b_size = 2; + + // Non-trivial logits so a non-sentinel label would produce non-zero + // loss and gradient — proving the sentinel really is masking. + use rand::{Rng, SeedableRng}; + use rand_chacha::ChaCha8Rng; + let mut r = ChaCha8Rng::seed_from_u64(0x5EE); + let logits: Vec = (0..b_size * FRD_OUT_DIM) + .map(|_| r.gen_range(-1.0..1.0)) + .collect(); + let logits_d = upload_f32(&stream, &logits)?; + // Labels: ALL sentinel (-1) → every row should be zeroed. + let labels: Vec = vec![-1; b_size * FRD_N_HORIZONS]; + let labels_d = upload_i32(&stream, &labels)?; + 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)?; + let grad = read_slice_d_pub(&stream, &grad_d, b_size * FRD_OUT_DIM)?; + for (i, v) in loss.iter().enumerate() { + assert_eq!(*v, 0.0, "sentinel label loss[{i}] must be 0; got {v}"); + } + for (i, v) in grad.iter().enumerate() { + assert_eq!( + *v, 0.0, + "sentinel label grad_logits[{i}] must be 0; got {v}" + ); + } + + eprintln!("PASS — sentinel label (-1) zeros loss + grad rows"); + Ok(()) +} + +#[test] +#[ignore = "requires CUDA (MlDevice::cuda(0))"] +fn frd_softmax_ce_grad_finite_diff_matches_analytical() -> Result<()> { + let Some((dev, head)) = build_head() else { return Ok(()) }; + let stream = dev.cuda_stream()?.clone(); + let b_size = 1; + + // Single batch, fixed (h=0, target_a=5). Pick a smooth random logit + // pattern so probabilities aren't degenerate at 0/1. + use rand::{Rng, SeedableRng}; + use rand_chacha::ChaCha8Rng; + let mut r = ChaCha8Rng::seed_from_u64(0xFD1FF); + let mut logits: Vec = (0..b_size * FRD_OUT_DIM) + .map(|_| r.gen_range(-0.5..0.5)) + .collect(); + let labels: Vec = vec![5, 10, 15]; // distinct per horizon + let labels_d = upload_i32(&stream, &labels)?; + + // Analytical gradient via the kernel. + 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 grad_analytical = read_slice_d_pub(&stream, &grad_d, b_size * FRD_OUT_DIM)?; + + // Finite-difference for slot (b=0, h=0, a=3). Note: gradient was + // mean-reduced by 1/B at the kernel source; for b_size=1 the 1/B + // factor is 1.0 so finite-diff matches directly. + let probe_h = 0_usize; + let probe_a = 3_usize; + let probe_off = probe_h * FRD_N_ATOMS + probe_a; + let eps = 1e-3_f32; + + // L(logits + ε · e_j) — perturb only the target slot upward. + logits[probe_off] += eps; + let logits_plus_d = upload_f32(&stream, &logits)?; + head.softmax_ce_grad(&logits_plus_d, &labels_d, &mut grad_d, &mut loss_d, b_size)?; + let loss_plus = read_slice_d_pub(&stream, &loss_d, b_size * FRD_N_HORIZONS)?; + let l_plus = loss_plus[probe_h]; // only h=0 affected — h=1,2 share the perturbation only if probe was in their horizon block + + // L(logits - ε · e_j) + logits[probe_off] -= 2.0 * eps; + let logits_minus_d = upload_f32(&stream, &logits)?; + head.softmax_ce_grad(&logits_minus_d, &labels_d, &mut grad_d, &mut loss_d, b_size)?; + let loss_minus = read_slice_d_pub(&stream, &loss_d, b_size * FRD_N_HORIZONS)?; + let l_minus = loss_minus[probe_h]; + + let numerical = (l_plus - l_minus) / (2.0 * eps); + let analytical = grad_analytical[probe_off]; + + let rel_err = (numerical - analytical).abs() / (analytical.abs().max(1e-6)); + // 5e-3 tolerance: fp32 single-precision finite-difference at ε=1e-3 + // saturates around the per-evaluation rounding error (~1e-7) divided + // by ε, so 1e-4-ish absolute error is intrinsic to the test method, + // not the kernel. Tighter threshold would force fp64 finite-diff + // (not worth the complexity here). + assert!( + rel_err < 5e-3, + "finite-diff gradient mismatch at (h={probe_h}, a={probe_a}): \ + analytical={analytical:.6}, numerical={numerical:.6}, rel_err={rel_err:.6}" + ); + + eprintln!( + "PASS — finite-diff matches analytical: analytical={:.6} numerical={:.6} rel_err={:.2e}", + analytical, numerical, rel_err + ); + Ok(()) +}