feat(rl): FRD layer-2 backward (dW2, db2, dhidden) — F.3b
Second of three FRD backward stages. Given dL/dlogits from F.3a's
softmax_ce_grad and the cached hidden activation from F.2's forward,
computes the layer-2 weight gradients via the standard chain rule
and emits the upstream gradient for layer-1 backward (F.3c).
Kernel `cuda/rl_frd_layer2_bwd.cu`:
* grid_dim = (B, 1, 1), block_dim = (FRD_HIDDEN_DIM=64, 1, 1)
* Phase 0: stage 63-slot grad_logits into shared (thread 63 idle)
* Phase 1: each thread i (i < 64) computes one row of per-batch
dW2 scratch: grad_w2_per_batch[b, i, 0..63] = h_bi × grad_logits[0..63]
(63 writes per thread, no atomics)
* Phase 2: each thread i computes dL/dhidden[b, i] = Σ_j W2[i, j] × grad_logits[j]
* Phase 3: thread i (i < 63) writes grad_b2_per_batch[b, i] = grad_logits[b, i]
* Per-batch scratch shape [B, FRD_HIDDEN_DIM, FRD_OUT_DIM] reduces
across batch via existing reduce_axis0 infra (caller's job, same
pattern as v_head_bwd / aux_heads_bwd)
Rust wiring `FrdHead::layer2_bwd`:
* Takes hidden (forward cache), grad_logits (from softmax_ce_grad),
self.w2_d
* Writes grad_w2_per_batch, grad_b2_per_batch, grad_hidden — all
sized to caller-allocated buffers
* Sole &self method (Adam step is the caller's responsibility)
Tests (2 new, 8/8 file total):
* frd_layer2_bwd_finite_diff_w2 — perturb W2[10, 5] by ±ε=1e-3,
compare (L(+) - L(-))/(2ε) to per-batch grad scratch. rel_err
= 6.27e-5 (better than F.3a's softmax-CE finite-diff because
the gradient magnitude here is larger so rounding error is
relatively smaller). Helper `ce_total_loss` re-uses
`softmax_ce_grad` to compute total CE for the perturbed forward
pass — pure GPU-oracle, no CPU softmax/CE reference impl.
* frd_layer2_bwd_db2_equals_grad_logits — analytical invariant:
db2_per_batch[b, j] must equal grad_logits[b, j] exactly (the
bias gradient is the identity passthrough at this layer). Cheap
structural check that catches dimension-shuffle bugs in the
kernel before they corrupt the reduce_axis0 step.
The kernel restores W2 to its original values after the perturbation
to keep test isolation clean — `&mut head` access pattern (proper
Rust borrowing, no UB const→mut casts).
F.3c (layer-1 backward: dW1, db1, dh_t with ReLU mask via the
cached hidden activation) is next.
This commit is contained in:
jgrusewski
@@ -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
|
||||
|
||||
80
crates/ml-alpha/cuda/rl_frd_layer2_bwd.cu
Normal file
80
crates/ml-alpha/cuda/rl_frd_layer2_bwd.cu
Normal file
@@ -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 <stdint.h>
|
||||
|
||||
#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];
|
||||
}
|
||||
}
|
||||
@@ -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<CudaModule>,
|
||||
softmax_ce_grad_fn: CudaFunction,
|
||||
_layer2_bwd_module: Arc<CudaModule>,
|
||||
layer2_bwd_fn: CudaFunction,
|
||||
|
||||
pub w1_d: CudaSlice<f32>, // [HIDDEN_DIM, FRD_HIDDEN_DIM]
|
||||
pub b1_d: CudaSlice<f32>, // [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<f32>,
|
||||
grad_logits_d: &CudaSlice<f32>,
|
||||
grad_w2_per_batch_d: &mut CudaSlice<f32>,
|
||||
grad_b2_per_batch_d: &mut CudaSlice<f32>,
|
||||
grad_hidden_d: &mut CudaSlice<f32>,
|
||||
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.
|
||||
|
||||
@@ -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<CudaStream>,
|
||||
logits: &[f32],
|
||||
labels_d: &CudaSlice<i32>,
|
||||
b_size: usize,
|
||||
) -> Result<f32> {
|
||||
let logits_d = upload_f32(stream, logits)?;
|
||||
let mut grad_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_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<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 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 labels: Vec<i32> = vec![5, 10, 15];
|
||||
let labels_d = upload_i32(&stream, &labels)?;
|
||||
|
||||
// Softmax+CE grad of logits.
|
||||
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)?;
|
||||
|
||||
// Layer-2 backward: produce per-batch grad_W2 scratch.
|
||||
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,
|
||||
)?;
|
||||
|
||||
// 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::<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(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::<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 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<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 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 labels: Vec<i32> = (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::<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 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(())
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user