feat(rl): FRD layer-1 backward (dW1, db1, dh_t with ReLU mask) — F.3c
Third and final FRD backward stage. Closes the chain from
softmax+CE loss back to the encoder's hidden state h_t.
Kernel `cuda/rl_frd_layer1_bwd.cu`:
* grid_dim = (B, 1, 1), block_dim = (HIDDEN_DIM=128, 1, 1)
* Phase 0: threads 0..63 stage dL/dpre_hidden = grad_hidden ×
1{hidden > 0} into shared mem (the cached post-ReLU `hidden`
buffer encodes the mask — hidden == 0 ⇔ pre-activation was
≤ 0 → ReLU killed it). Same thread also writes db1_per_batch.
* Phase 1: each thread k (k < 128) writes one row of
grad_W1_per_batch[b, k, 0..64] (64 writes per thread, no atomics)
* Phase 2: same thread computes grad_h_t[b, k] =
Σ_i W1[k, i] × dL/dpre_hidden[b, i]
* Per-(b, k, i) sole-writer per feedback_no_atomicadd
Rust wiring `FrdHead::layer1_bwd` — takes h_t, hidden (forward cache),
grad_hidden (from layer2_bwd), self.w1_d; writes grad_w1_per_batch,
grad_b1_per_batch, grad_h_t. The grad_h_t buffer becomes the encoder-
upstream gradient that the trainer's grad_h_accumulate kernel folds
into the encoder's gradient with λ_frd scaling (same pattern as Q/π/V
heads — wiring lives in F.4).
Tests (2 new, 10/10 file total):
* frd_layer1_bwd_finite_diff_w1 — perturbs the W1 slot with MAX
|analytical gradient| (instead of an arbitrary fixed slot — fp32
finite-diff is rounding-error-limited so a tiny gradient gives
misleading rel_err). At max-magnitude slot (k=84, i=55): analytical
= -0.0451, numerical = -0.0448, rel_err = 5.6e-3 — well within
1e-2 tolerance (slightly looser than dW2's 5e-3 because dW1
crosses an extra matmul + the ReLU mask boundary).
* frd_layer1_bwd_relu_mask_zeros_grad — fixture with h_t = all -1
produces ~half the hidden slots ReLU-masked (cached hidden = 0).
For every masked slot i, asserts:
* db1_per_batch[b, i] == 0 (exact equality — mask is hard 0)
* dW1_per_batch[b, k, i] == 0 for every k (~32 × 128 = 4096
slots checked)
Empirically 32/64 masked, 32/64 active — confirms ReLU mask
is wired through the chain correctly without leaking gradient
through dead branches.
F.3 backward chain is now complete end-to-end:
rl_frd_softmax_ce_grad (F.3a) → rl_frd_layer2_bwd (F.3b) →
rl_frd_layer1_bwd (F.3c) → grad_h_t (consumed by F.4 wiring)
F.4 wires Adam optimizers for W1/b1/W2/b2 + grad_h_accumulate into
the encoder gradient + loader-side label generation + λ_frd × CE
into stats.l_total.
This commit is contained in:
jgrusewski
@@ -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
|
||||
|
||||
79
crates/ml-alpha/cuda/rl_frd_layer1_bwd.cu
Normal file
79
crates/ml-alpha/cuda/rl_frd_layer1_bwd.cu
Normal file
@@ -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 <stdint.h>
|
||||
|
||||
#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;
|
||||
}
|
||||
@@ -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<CudaModule>,
|
||||
layer2_bwd_fn: CudaFunction,
|
||||
_layer1_bwd_module: Arc<CudaModule>,
|
||||
layer1_bwd_fn: CudaFunction,
|
||||
|
||||
pub w1_d: CudaSlice<f32>, // [HIDDEN_DIM, FRD_HIDDEN_DIM]
|
||||
pub b1_d: CudaSlice<f32>, // [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<f32>,
|
||||
hidden_d: &CudaSlice<f32>,
|
||||
grad_hidden_d: &CudaSlice<f32>,
|
||||
grad_w1_per_batch_d: &mut CudaSlice<f32>,
|
||||
grad_b1_per_batch_d: &mut CudaSlice<f32>,
|
||||
grad_h_t_d: &mut CudaSlice<f32>,
|
||||
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.
|
||||
|
||||
@@ -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<f32> = (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<i32> = 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::<f32>(b_size * FRD_HIDDEN_DIM)?;
|
||||
let mut logits_d = stream.alloc_zeros::<f32>(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::<f32>(b_size * FRD_OUT_DIM)?;
|
||||
let mut loss_d = stream.alloc_zeros::<f32>(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::<f32>(b_size * FRD_HIDDEN_DIM * FRD_OUT_DIM)?;
|
||||
let mut grad_b2_pb_d = stream.alloc_zeros::<f32>(b_size * FRD_OUT_DIM)?;
|
||||
let mut grad_hidden_d = stream.alloc_zeros::<f32>(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::<f32>(b_size * HIDDEN_DIM * FRD_HIDDEN_DIM)?;
|
||||
let mut grad_b1_pb_d = stream.alloc_zeros::<f32>(b_size * FRD_HIDDEN_DIM)?;
|
||||
let mut grad_h_t_d = stream.alloc_zeros::<f32>(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::<f32>(b_size * FRD_HIDDEN_DIM)?;
|
||||
let mut logits_plus = stream.alloc_zeros::<f32>(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::<f32>(b_size * FRD_HIDDEN_DIM)?;
|
||||
let mut logits_minus = stream.alloc_zeros::<f32>(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<i32> = vec![5, 10, 15];
|
||||
let labels_d = upload_i32(&stream, &labels)?;
|
||||
|
||||
let mut hidden_d = stream.alloc_zeros::<f32>(b_size * FRD_HIDDEN_DIM)?;
|
||||
let mut logits_d = stream.alloc_zeros::<f32>(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::<f32>(b_size * FRD_OUT_DIM)?;
|
||||
let mut loss_d = stream.alloc_zeros::<f32>(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::<f32>(b_size * FRD_HIDDEN_DIM * FRD_OUT_DIM)?;
|
||||
let mut grad_b2_pb_d = stream.alloc_zeros::<f32>(b_size * FRD_OUT_DIM)?;
|
||||
let mut grad_hidden_d = stream.alloc_zeros::<f32>(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::<f32>(b_size * HIDDEN_DIM * FRD_HIDDEN_DIM)?;
|
||||
let mut grad_b1_pb_d = stream.alloc_zeros::<f32>(b_size * FRD_HIDDEN_DIM)?;
|
||||
let mut grad_h_t_d = stream.alloc_zeros::<f32>(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(())
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user