diff --git a/crates/ml-alpha/build.rs b/crates/ml-alpha/build.rs index d069041ac..99c443dbf 100644 --- a/crates/ml-alpha/build.rs +++ b/crates/ml-alpha/build.rs @@ -157,6 +157,8 @@ const KERNELS: &[&str] = &[ "rl_deterministic_checksum", // Determinism foundation Phase 1 (2026-06-02): provably deterministic sum-of-squares (single-block / single-thread / f64 accumulation) for per-step component checksums in diag; localizes non-determinism source. Spec: docs/superpowers/specs/2026-06-02-determinism-foundation.md §1.1 "multi_head_policy_forward", // Phase 2A-A (2026-06-03): K-head policy logit + softmax + mixture combination; per-batch deterministic sequential mixture sum. Spec: docs/superpowers/specs/2026-06-02-multi-head-policy-with-r-multiple.md §R.6 "multi_head_policy_gate_forward",// Phase 2A-A (2026-06-03): gating head from raw regime features [B × 6] (parallel channel bypassing VSN/Mamba2). Spec ADDENDUM 2026-06-03 §R.3 Option B + "multi_head_policy_backward", // Phase 2A-B (2026-06-03): backward through mixture + per-head softmax + gating softmax. Two kernels: `_backward_pi` and `_backward_gate`. Per-batch grad scratch; caller reduces via reduce_axis0. Spec ADDENDUM 2026-06-03 §R.6 + "multi_head_policy_aux_prior", // Phase 2A-B (2026-06-03): per-head auxiliary KL prior gradient — additive into grad_pi_logits_k. β read from ISV slot 764. Spec ADDENDUM 2026-06-03 §R.4 ]; // Cache bust v31 — five new reduce / derive kernels populate the input diff --git a/crates/ml-alpha/cuda/multi_head_policy_aux_prior.cu b/crates/ml-alpha/cuda/multi_head_policy_aux_prior.cu new file mode 100644 index 000000000..5793d5570 --- /dev/null +++ b/crates/ml-alpha/cuda/multi_head_policy_aux_prior.cu @@ -0,0 +1,91 @@ +// multi_head_policy_aux_prior.cu — Per-head auxiliary KL prior gradient. +// +// Spec: docs/superpowers/specs/2026-06-02-multi-head-policy-with-r-multiple.md +// ADDENDUM 2026-06-03 §R.4 (per-head aux KL priors) +// Plan: docs/superpowers/plans/2026-06-03-multi-head-policy-implementation.md (Phase 2A-B) +// +// For each head k, we add an auxiliary KL regularizer that anchors the +// head distribution to a fixed prior: +// +// L_aux_k = β · KL(π_k || prior_k) +// = β · Σ_a π_k(a) · log( π_k(a) / prior_k(a) ) +// +// The grad of L_aux on the head's pre-softmax logits is the standard +// softmax-Jacobian form (parameterization-invariant): +// +// ∂L_aux_k / ∂pi_logits_k[a] = β · π_k(a) · ( log(π_k(a)/prior_k(a)) − KL_k ) +// +// This kernel ADDS that contribution to `grad_pi_logits_k` (which the +// backward pi kernel populated with the Q-distill chain rule). +// +// β is read from ISV slot `RL_POLICY_AUX_PRIOR_BETA_INDEX = 764`. +// +// ── Block layout ──────────────────────────────────────────────────────── +// grid = (B, K, 1) +// block = (N_ACTIONS = 11, 1, 1) +// +// One block per (batch, head). N_ACTIONS threads cooperate via shared +// memory to compute KL_k[b] then write the additive grad per action. +// +// Per feedback_no_atomicadd: sole-writer per (b, k, a) cell within this +// kernel; `grad_pi_logits_k` is read-modify-write but additive into a +// slot that was just exclusively written by the preceding pi backward +// kernel on the same stream — sequential stream ordering keeps the RMW +// safe (no races). +// +// Per feedback_no_nvrtc: pre-compiled cubin via build.rs. + +#include + +#define N_ACTIONS 11 +#define RL_POLICY_AUX_PRIOR_BETA_INDEX 764 +#define KL_PROB_EPS 1e-12f // log domain stability + +extern "C" __global__ void multi_head_policy_aux_prior( + const float* __restrict__ pi_probs_k, // [B × K × N_ACTIONS] + const float* __restrict__ priors, // [K × N_ACTIONS] + const float* __restrict__ isv, // ISV bus — slot 764 holds β + const int B, + const int K, + float* __restrict__ grad_pi_logits_k // [B × K × N_ACTIONS] RW (+=) +) { + const int b = blockIdx.x; + const int k = blockIdx.y; + if (b >= B || k >= K) return; + + const int a = threadIdx.x; + if (a >= N_ACTIONS) return; + + const float beta = isv[RL_POLICY_AUX_PRIOR_BETA_INDEX]; + // Zero-β short circuit — leaves grad_pi_logits_k untouched. + if (beta == 0.0f) return; + + // ── Stage probs + per-action log(π/prior) ───────────────────────── + __shared__ float s_pi[N_ACTIONS]; + __shared__ float s_log_ratio[N_ACTIONS]; + __shared__ float s_kl; + + const int pk_idx = (b * K + k) * N_ACTIONS + a; + const float pi_a = pi_probs_k[pk_idx]; + const float prior_a = priors[k * N_ACTIONS + a]; + + s_pi[a] = pi_a; + // log(π_a / prior_a) with both args clamped against fp32 underflow. + s_log_ratio[a] = logf(fmaxf(pi_a, KL_PROB_EPS) / fmaxf(prior_a, KL_PROB_EPS)); + __syncthreads(); + + // ── KL_k[b] = Σ_a π_a · log(π_a / prior_a) ──────────────────────── + // Single-thread reduction over N_ACTIONS=11 for deterministic order. + if (a == 0) { + float kl = 0.0f; + for (int aa = 0; aa < N_ACTIONS; ++aa) { + kl += s_pi[aa] * s_log_ratio[aa]; + } + s_kl = kl; + } + __syncthreads(); + + // ── grad += β · π_a · (log_ratio − KL) ──────────────────────────── + const float g_aux = beta * s_pi[a] * (s_log_ratio[a] - s_kl); + grad_pi_logits_k[pk_idx] += g_aux; +} diff --git a/crates/ml-alpha/cuda/multi_head_policy_backward.cu b/crates/ml-alpha/cuda/multi_head_policy_backward.cu new file mode 100644 index 000000000..27035352f --- /dev/null +++ b/crates/ml-alpha/cuda/multi_head_policy_backward.cu @@ -0,0 +1,300 @@ +// multi_head_policy_backward.cu — Backward pass for the K-head policy mixture. +// +// Spec: docs/superpowers/specs/2026-06-02-multi-head-policy-with-r-multiple.md +// ADDENDUM 2026-06-03 §R.4 (per-head aux KL priors), §R.6 (gradient flow) +// Plan: docs/superpowers/plans/2026-06-03-multi-head-policy-implementation.md (Phase 2A-B) +// +// ── Forward recap (multi_head_policy_forward.cu) ──────────────────────── +// pi_logits_k[b,k,a] = W_heads[k,a,:] · h_t[b,:] + b_heads[k,a] +// pi_probs_k[b,k,a] = softmax_a(pi_logits_k[b,k,:]) +// pi_probs[b,a] = Σ_k gate_probs[b,k] · pi_probs_k[b,k,a] +// pi_logits[b,a] = log(pi_probs[b,a] + ε) +// +// ── Backward chain implemented by this file (two kernels) ─────────────── +// +// (1) `multi_head_policy_backward_pi` — distributes incoming `grad_pi_logits` +// through the mixture and per-head softmax. Outputs: +// a. grad_pi_logits_k [B × K × N_ACTIONS] (for aux prior + diag) +// b. grad_gate_probs [B × K] (consumed by kernel 2) +// c. grad_W_heads_pb [B × K × N_ACTIONS × HIDDEN_DIM] +// d. grad_b_heads_pb [B × K × N_ACTIONS] +// e. grad_h_t_scratch [B × HIDDEN_DIM] (OVERWRITE; caller folds +// via grad_h_accumulate_scaled) +// +// (2) `multi_head_policy_backward_gate` — distributes `grad_gate_probs` +// through the gating softmax. Outputs: +// a. grad_gate_logits [B × K] (diag) +// b. grad_W_gate_pb [B × K × REGIME_DIM] +// c. grad_b_gate_pb [B × K] +// d. grad_regime_h [B × REGIME_DIM] (computed; NOT used +// upstream in Phase 2A-B) +// +// ── Math (kernel 1) ───────────────────────────────────────────────────── +// grad_pi_probs[b,a] = grad_pi_logits[b,a] / (pi_probs[b,a] + ε) +// grad_pi_probs_k[b,k,a] = gate_probs[b,k] · grad_pi_probs[b,a] +// grad_gate_probs[b,k] = Σ_a pi_probs_k[b,k,a] · grad_pi_probs[b,a] +// +// (per-head softmax-Jacobian; standard form): +// grad_pi_logits_k[b,k,a] = pi_probs_k[b,k,a] +// · ( grad_pi_probs_k[b,k,a] +// − Σ_{a'} pi_probs_k[b,k,a'] · grad_pi_probs_k[b,k,a'] ) +// +// ── Math (kernel 2) ───────────────────────────────────────────────────── +// grad_gate_logits[b,k] = gate_probs[b,k] +// · ( grad_gate_probs[b,k] +// − Σ_{k'} gate_probs[b,k'] · grad_gate_probs[b,k'] ) +// +// grad_W_gate_pb[b,k,r] = grad_gate_logits[b,k] · regime_h[b,r] +// grad_b_gate_pb[b,k] = grad_gate_logits[b,k] +// grad_regime_h[b,r] = Σ_k W_gate[k,r] · grad_gate_logits[b,k] +// +// ── Per-batch grad scratch convention ─────────────────────────────────── +// Caller reduces `_pb` buffers via `reduce_axis0`: +// grad_W_heads_pb [B × K·N·H] → grad_W_heads [K·N·H] +// grad_b_heads_pb [B × K·N] → grad_b_heads [K·N] +// grad_W_gate_pb [B × K·R] → grad_W_gate [K·R] +// grad_b_gate_pb [B × K] → grad_b_gate [K] +// +// Per feedback_no_atomicadd: NO atomicAdd. Each output cell has exactly one +// writer per launch. Per-batch scratch + reduce_axis0 is the standard +// foxhunt pattern (see ppo_clipped_surrogate.cu::ppo_grad_w_b_h_t). +// +// Per feedback_no_nvrtc: pre-compiled cubin via build.rs. + +#include + +#define HIDDEN_DIM 128 +#define N_ACTIONS 11 +#define REGIME_DIM 6 +#define MAX_K_HEADS 8 // Compile-time max; runtime K from ISV +#define LOG_PROB_EPS 1e-12f // mirrors forward kernel's `+1e-12` guard + +// ───────────────────────────────────────────────────────────────────── +// Kernel 1: backward through the mixture + per-head softmax. +// +// Block layout: +// grid = (B, 1, 1) +// block = (HIDDEN_DIM = 128, 1, 1) +// shared = MAX_K_HEADS*N_ACTIONS*3 + MAX_K_HEADS + N_ACTIONS ≈ 304 floats +// +// Thread c handles one hidden-dim index. Threads 0..N_ACTIONS-1 first +// reconstruct pi_probs[b,:] and the softmax-Jacobian — only threads with +// `c < N_ACTIONS` participate in the small reductions. All HIDDEN_DIM +// threads participate in the per-batch weight-grad emit. +// ───────────────────────────────────────────────────────────────────── +extern "C" __global__ void multi_head_policy_backward_pi( + const float* __restrict__ grad_pi_logits, // [B × N_ACTIONS] incoming grad on log-mixture + const float* __restrict__ h_t, // [B × HIDDEN_DIM] + const float* __restrict__ W_heads, // [K × N_ACTIONS × HIDDEN_DIM] + const float* __restrict__ gate_probs, // [B × K] forward output + const float* __restrict__ pi_probs_k, // [B × K × N_ACTIONS] forward output + const int B, + const int K, + float* __restrict__ grad_pi_logits_k, // [B × K × N_ACTIONS] OUT (for aux + diag) + float* __restrict__ grad_gate_probs, // [B × K] OUT (→ kernel 2) + float* __restrict__ grad_W_heads_pb, // [B × K × N_ACTIONS × HIDDEN_DIM] + float* __restrict__ grad_b_heads_pb, // [B × K × N_ACTIONS] + float* __restrict__ grad_h_t_scratch // [B × HIDDEN_DIM] (OVERWRITE) +) { + const int b = blockIdx.x; + if (b >= B) return; + const int c = threadIdx.x; // hidden-dim index, 0..HIDDEN_DIM-1 + + // ── Shared state (per-batch) ────────────────────────────────────── + __shared__ float s_pi_probs[N_ACTIONS]; // mixture probs + __shared__ float s_grad_pi_probs[N_ACTIONS]; // dL/dp[a] + __shared__ float s_gate_probs[MAX_K_HEADS]; // forward gate + __shared__ float s_pi_probs_k[MAX_K_HEADS * N_ACTIONS]; // forward per-head probs + __shared__ float s_grad_pi_probs_k[MAX_K_HEADS * N_ACTIONS]; // step 2 intermediate + __shared__ float s_grad_pi_logits_k[MAX_K_HEADS * N_ACTIONS]; // final per-head logit grad + __shared__ float s_dot_k[MAX_K_HEADS]; // softmax-Jacobian scalar per head + + // ── Stage forward outputs (small, per-batch) into shared mem ─────── + if (c < K) { + s_gate_probs[c] = gate_probs[b * K + c]; + } + if (c < N_ACTIONS) { + const int gpi_idx = b * N_ACTIONS + c; + s_grad_pi_probs[c] = 0.0f; // accumulator, populated below + // Stage pi_probs_k[b, :, c] across all heads. + for (int k = 0; k < K; ++k) { + s_pi_probs_k[k * N_ACTIONS + c] = + pi_probs_k[(b * K + k) * N_ACTIONS + c]; + } + // Reconstruct pi_probs[b, c] = Σ_k gate · pi_probs_k. + // Sequential summation order matches the forward kernel for + // determinism. + float p = 0.0f; + #pragma unroll + for (int k = 0; k < K; ++k) { + p += gate_probs[b * K + k] * pi_probs_k[(b * K + k) * N_ACTIONS + c]; + } + s_pi_probs[c] = p; + // dL/dp[a] = dL/dlog(p+ε) × 1/(p+ε) + const float gl = grad_pi_logits[gpi_idx]; + s_grad_pi_probs[c] = gl / (s_pi_probs[c] + LOG_PROB_EPS); + } + __syncthreads(); + + // ── Step 2: grad_pi_probs_k + grad_gate_probs ───────────────────── + // grad_pi_probs_k[b,k,a] = gate_probs[b,k] · grad_pi_probs[b,a] + // grad_gate_probs[b,k] = Σ_a pi_probs_k[b,k,a] · grad_pi_probs[b,a] + if (c < N_ACTIONS) { + for (int k = 0; k < K; ++k) { + s_grad_pi_probs_k[k * N_ACTIONS + c] = + s_gate_probs[k] * s_grad_pi_probs[c]; + } + } + __syncthreads(); + + // Thread k computes Σ_a pi_probs_k[k,a] · grad_pi_probs[a] for its head. + if (c < K) { + float gp = 0.0f; + #pragma unroll + for (int a = 0; a < N_ACTIONS; ++a) { + gp += s_pi_probs_k[c * N_ACTIONS + a] * s_grad_pi_probs[a]; + } + grad_gate_probs[b * K + c] = gp; + } + __syncthreads(); + + // ── Step 3: per-head softmax-Jacobian → grad_pi_logits_k ────────── + // dot_k[b,k] = Σ_{a'} pi_probs_k[k,a'] · grad_pi_probs_k[k,a'] + if (c < K) { + float dot = 0.0f; + #pragma unroll + for (int a = 0; a < N_ACTIONS; ++a) { + dot += s_pi_probs_k[c * N_ACTIONS + a] + * s_grad_pi_probs_k[c * N_ACTIONS + a]; + } + s_dot_k[c] = dot; + } + __syncthreads(); + + if (c < N_ACTIONS) { + for (int k = 0; k < K; ++k) { + const int kak = k * N_ACTIONS + c; + const float gl = s_pi_probs_k[kak] + * ( s_grad_pi_probs_k[kak] - s_dot_k[k] ); + s_grad_pi_logits_k[kak] = gl; + // Persist to global memory for aux prior kernel + diagnostics. + grad_pi_logits_k[(b * K + k) * N_ACTIONS + c] = gl; + } + } + __syncthreads(); + + // ── Step 4: per-batch weight grads + grad_h_t ───────────────────── + // Each thread c is the sole writer of: + // grad_W_heads_pb[b, k, a, c] for all (k, a) + // grad_h_t_scratch[b, c] + // Threads c= B) return; + const int tid = threadIdx.x; + + __shared__ float s_gate_probs[MAX_K_HEADS]; + __shared__ float s_grad_gate_probs[MAX_K_HEADS]; + __shared__ float s_grad_gate_logits[MAX_K_HEADS]; + __shared__ float s_dot; // Σ_{k'} gate · grad_gate_probs + + // ── Stage forward + incoming grad ───────────────────────────────── + if (tid < K) { + s_gate_probs[tid] = gate_probs[b * K + tid]; + s_grad_gate_probs[tid] = grad_gate_probs[b * K + tid]; + } + __syncthreads(); + + // ── Softmax-Jacobian reduction (single thread for determinism) ──── + if (tid == 0) { + float dot = 0.0f; + for (int k = 0; k < K; ++k) { + dot += s_gate_probs[k] * s_grad_gate_probs[k]; + } + s_dot = dot; + } + __syncthreads(); + + // ── grad_gate_logits[b, k] = p · (dp − dot) ─────────────────────── + if (tid < K) { + const float gl = s_gate_probs[tid] * (s_grad_gate_probs[tid] - s_dot); + s_grad_gate_logits[tid] = gl; + grad_gate_logits[b * K + tid] = gl; + grad_b_gate_pb[b * K + tid] = gl; + } + __syncthreads(); + + // ── grad_W_gate_pb[b, k, r] = grad_gate_logits[k] · regime_h[r] ─── + // Cover all (k, r) slots — K*REGIME_DIM ≤ 48 — using a single loop + // per active thread. Thread `tid` handles its column across all k. + if (tid < REGIME_DIM) { + const float r_b = regime_h[b * REGIME_DIM + tid]; + for (int k = 0; k < K; ++k) { + grad_W_gate_pb[((b * K + k) * REGIME_DIM) + tid] = + s_grad_gate_logits[k] * r_b; + } + } + + // ── grad_regime_h[b, r] = Σ_k W_gate[k, r] · grad_gate_logits[k] ── + // Computed but NOT propagated upstream in Phase 2A-B (regime_h is + // loader-precomputed in foxhunt). Allocate + write the buffer so a + // future phase can route it if needed. + if (tid < REGIME_DIM) { + float gr = 0.0f; + #pragma unroll + for (int k = 0; k < K; ++k) { + gr += W_gate[k * REGIME_DIM + tid] * s_grad_gate_logits[k]; + } + grad_regime_h[b * REGIME_DIM + tid] = gr; + } +} diff --git a/crates/ml-alpha/src/rl/multi_head_policy.rs b/crates/ml-alpha/src/rl/multi_head_policy.rs index 88f5c0dbc..e9838bce7 100644 --- a/crates/ml-alpha/src/rl/multi_head_policy.rs +++ b/crates/ml-alpha/src/rl/multi_head_policy.rs @@ -68,6 +68,16 @@ const GATE_FORWARD_CUBIN: &[u8] = include_bytes!(concat!( env!("OUT_DIR"), "/multi_head_policy_gate_forward.cubin" )); +const BACKWARD_CUBIN: &[u8] = include_bytes!(concat!( + env!("OUT_DIR"), + "/multi_head_policy_backward.cubin" +)); +const AUX_PRIOR_CUBIN: &[u8] = include_bytes!(concat!( + env!("OUT_DIR"), + "/multi_head_policy_aux_prior.cubin" +)); +const REDUCE_AXIS0_CUBIN: &[u8] = + include_bytes!(concat!(env!("OUT_DIR"), "/reduce_axis0.cubin")); /// Compile-time upper bound on K — mirrors the `#define MAX_K_HEADS 8` in /// both kernel sources. The runtime K (from `RL_POLICY_NUM_HEADS_INDEX`) @@ -115,6 +125,13 @@ pub struct MultiHeadPolicy { forward_fn: CudaFunction, _gate_forward_module: Arc, gate_forward_fn: CudaFunction, + _backward_module: Arc, + backward_pi_fn: CudaFunction, + backward_gate_fn: CudaFunction, + _aux_prior_module: Arc, + aux_prior_fn: CudaFunction, + _reduce_axis0_module: Arc, + reduce_axis0_fn: CudaFunction, // ── Policy-head weights ([K × N_ACTIONS × HIDDEN_DIM] row-major) ── /// `W_heads[k, a, j]` — row-major flat over `(k, a, j)`. @@ -137,6 +154,50 @@ pub struct MultiHeadPolicy { /// Per-head post-softmax probabilities `[B × K × N_ACTIONS]`. Sums to /// 1 per (batch, head) row. pub pi_probs_k_d: CudaSlice, + + // ── Per-head action priors ([K × N_ACTIONS]) ────────────────────── + /// Distribution priors for the aux KL term. Per-head priors mirror + /// the init bias of [`MultiHeadPolicy::new`] so the aux loss does not + /// fight the asymmetry-break initialization (`pearl_bootstrap_must_respect_clamp_range` + /// applied to priors). Heads 3..K (if K > 3) carry uniform priors. + pub priors_d: CudaSlice, + + // ── Backward output buffers ─────────────────────────────────────── + /// `[B × K × N_ACTIONS]` — gradient on per-head logits. Populated by + /// the backward pi kernel (mixture + softmax-Jacobian); aux prior + /// kernel adds its β-weighted contribution. Available for diagnostics + /// or further chain-rule routing in Phase 2A-C. + pub grad_pi_logits_k_d: CudaSlice, + /// `[B × K]` — gradient on the gating distribution (before softmax-Jacobian), + /// computed by the pi backward kernel and consumed by the gate backward kernel. + pub grad_gate_probs_d: CudaSlice, + /// `[B × K]` — gradient on the pre-softmax gate logits (for diagnostics). + pub grad_gate_logits_d: CudaSlice, + /// `[B × HIDDEN_DIM]` — head's contribution to the encoder grad. The + /// `backward()` method writes (OVERWRITE) into this scratch; the + /// caller folds it into the encoder's `grad_h_t` via + /// `grad_h_accumulate_scaled`. + pub grad_h_t_scratch_d: CudaSlice, + /// `[B × REGIME_DIM]` — gradient on the regime input. Computed for + /// completeness but NOT routed upstream in Phase 2A-B (regime is + /// loader-precomputed). + pub grad_regime_h_d: CudaSlice, + + // ── Per-batch grad scratch (for reduce_axis0) ───────────────────── + grad_w_heads_pb_d: CudaSlice, + grad_b_heads_pb_d: CudaSlice, + grad_w_gate_pb_d: CudaSlice, + grad_b_gate_pb_d: CudaSlice, + + // ── Reduced grad buffers (post `reduce_axis0`) ──────────────────── + /// Reduced grad on `W_heads`, ready for an Adam step in Phase 2A-C. + pub grad_w_heads_d: CudaSlice, + /// Reduced grad on `b_heads`. + pub grad_b_heads_d: CudaSlice, + /// Reduced grad on `W_gate`. + pub grad_w_gate_d: CudaSlice, + /// Reduced grad on `b_gate`. + pub grad_b_gate_d: CudaSlice, } impl MultiHeadPolicy { @@ -186,6 +247,27 @@ impl MultiHeadPolicy { let gate_forward_fn = gate_forward_module .load_function("multi_head_policy_gate_forward") .context("load multi_head_policy_gate_forward fn")?; + let backward_module = ctx + .load_cubin(BACKWARD_CUBIN.to_vec()) + .context("load multi_head_policy_backward cubin")?; + let backward_pi_fn = backward_module + .load_function("multi_head_policy_backward_pi") + .context("load multi_head_policy_backward_pi fn")?; + let backward_gate_fn = backward_module + .load_function("multi_head_policy_backward_gate") + .context("load multi_head_policy_backward_gate fn")?; + let aux_prior_module = ctx + .load_cubin(AUX_PRIOR_CUBIN.to_vec()) + .context("load multi_head_policy_aux_prior cubin")?; + let aux_prior_fn = aux_prior_module + .load_function("multi_head_policy_aux_prior") + .context("load multi_head_policy_aux_prior fn")?; + let reduce_axis0_module = ctx + .load_cubin(REDUCE_AXIS0_CUBIN.to_vec()) + .context("load reduce_axis0 cubin (multi_head_policy)")?; + let reduce_axis0_fn = reduce_axis0_module + .load_function("reduce_axis0") + .context("load reduce_axis0 fn (multi_head_policy)")?; // Install determinism guard before drawing any random samples // (pearl_scoped_init_seed_for_reproducibility). @@ -250,6 +332,59 @@ impl MultiHeadPolicy { .alloc_zeros::(cfg.b_size * cfg.k * N_ACTIONS) .context("alloc pi_probs_k_d")?; + // ── Per-head action priors ──────────────────────────────────── + // Mirror the per-head init bias so the aux KL term does not fight + // the asymmetry-break initialization. Heads 3..K (if K > 3) get a + // uniform prior. + let priors_host = build_priors(cfg.k); + let priors_d = upload(&stream, &priors_host) + .context("upload aux KL priors")?; + + // ── Backward output buffers ─────────────────────────────────── + let grad_pi_logits_k_d = stream + .alloc_zeros::(cfg.b_size * cfg.k * N_ACTIONS) + .context("alloc grad_pi_logits_k_d")?; + let grad_gate_probs_d = stream + .alloc_zeros::(cfg.b_size * cfg.k) + .context("alloc grad_gate_probs_d")?; + let grad_gate_logits_d = stream + .alloc_zeros::(cfg.b_size * cfg.k) + .context("alloc grad_gate_logits_d")?; + let grad_h_t_scratch_d = stream + .alloc_zeros::(cfg.b_size * HIDDEN_DIM) + .context("alloc grad_h_t_scratch_d")?; + let grad_regime_h_d = stream + .alloc_zeros::(cfg.b_size * REGIME_DIM) + .context("alloc grad_regime_h_d")?; + + // ── Per-batch grad scratch (for reduce_axis0) ──────────────── + let grad_w_heads_pb_d = stream + .alloc_zeros::(cfg.b_size * cfg.k * N_ACTIONS * HIDDEN_DIM) + .context("alloc grad_w_heads_pb_d")?; + let grad_b_heads_pb_d = stream + .alloc_zeros::(cfg.b_size * cfg.k * N_ACTIONS) + .context("alloc grad_b_heads_pb_d")?; + let grad_w_gate_pb_d = stream + .alloc_zeros::(cfg.b_size * cfg.k * REGIME_DIM) + .context("alloc grad_w_gate_pb_d")?; + let grad_b_gate_pb_d = stream + .alloc_zeros::(cfg.b_size * cfg.k) + .context("alloc grad_b_gate_pb_d")?; + + // ── Reduced grad buffers ────────────────────────────────────── + let grad_w_heads_d = stream + .alloc_zeros::(cfg.k * N_ACTIONS * HIDDEN_DIM) + .context("alloc grad_w_heads_d")?; + let grad_b_heads_d = stream + .alloc_zeros::(cfg.k * N_ACTIONS) + .context("alloc grad_b_heads_d")?; + let grad_w_gate_d = stream + .alloc_zeros::(cfg.k * REGIME_DIM) + .context("alloc grad_w_gate_d")?; + let grad_b_gate_d = stream + .alloc_zeros::(cfg.k) + .context("alloc grad_b_gate_d")?; + let raw_stream = stream.cu_stream(); Ok(Self { @@ -260,6 +395,13 @@ impl MultiHeadPolicy { forward_fn, _gate_forward_module: gate_forward_module, gate_forward_fn, + _backward_module: backward_module, + backward_pi_fn, + backward_gate_fn, + _aux_prior_module: aux_prior_module, + aux_prior_fn, + _reduce_axis0_module: reduce_axis0_module, + reduce_axis0_fn, w_heads_d, b_heads_d, w_gate_d, @@ -268,6 +410,20 @@ impl MultiHeadPolicy { gate_probs_d, pi_logits_k_d, pi_probs_k_d, + priors_d, + grad_pi_logits_k_d, + grad_gate_probs_d, + grad_gate_logits_d, + grad_h_t_scratch_d, + grad_regime_h_d, + grad_w_heads_pb_d, + grad_b_heads_pb_d, + grad_w_gate_pb_d, + grad_b_gate_pb_d, + grad_w_heads_d, + grad_b_heads_d, + grad_w_gate_d, + grad_b_gate_d, }) } @@ -376,6 +532,192 @@ impl MultiHeadPolicy { Ok(()) } + + /// Backward pass through the mixture, per-head softmax, and gating softmax. + /// + /// Inputs: + /// * `grad_pi_logits_d [B × N_ACTIONS]` — incoming gradient on the + /// mixture log-probs (forward output). In Phase 2A-C this will be + /// produced by `rl_q_pi_distill_grad`. + /// * `h_t_d [B × HIDDEN_DIM]` — same encoder hidden state passed to + /// forward (re-passed because backward needs it for the weight grad). + /// * `regime_h_d [B × REGIME_DIM]` — same regime input passed to + /// forward. + /// * `isv_dev_ptr` — raw device pointer to the trainer's ISV bus. + /// The aux prior kernel reads slot `RL_POLICY_AUX_PRIOR_BETA_INDEX`. + /// + /// Pipeline (single stream, sequential ordering): + /// 1. `multi_head_policy_backward_pi` — distributes incoming grad + /// through mixture + per-head softmax. Writes + /// `grad_pi_logits_k_d`, `grad_gate_probs_d`, per-batch + /// `grad_W_heads_pb` + `grad_b_heads_pb`, and + /// `grad_h_t_scratch_d` (OVERWRITE). + /// 2. `multi_head_policy_aux_prior` — adds β-weighted KL-prior grad + /// to `grad_pi_logits_k_d`. (β=0 → no-op.) NOTE: Phase 2A-B does + /// not back-propagate the aux contribution into weight grads; the + /// aux term shapes the per-head LOGIT distribution only. This + /// matches §R.4 of the spec (the aux is a regularizer on the head + /// output distribution, not a chain through the full network). + /// Phase 2A-C will wire the aux contribution into `grad_W_heads` if + /// empirically needed. + /// 3. `multi_head_policy_backward_gate` — distributes + /// `grad_gate_probs_d` through the gating softmax. Writes + /// `grad_gate_logits_d`, per-batch `grad_W_gate_pb` + + /// `grad_b_gate_pb`, and `grad_regime_h_d` (computed; not routed). + /// 4. `reduce_axis0` over B for all four per-batch grad scratches + /// yields the reduced `grad_W_heads_d`, `grad_b_heads_d`, + /// `grad_W_gate_d`, `grad_b_gate_d` ready for an Adam step. + /// + /// `grad_h_t_scratch_d` is the head's contribution to the encoder grad + /// and is exposed for the caller to fold via `grad_h_accumulate_scaled` + /// (Phase 2A-C wires this into the trainer's encoder backward chain). + pub fn backward( + &mut self, + grad_pi_logits_d: &CudaSlice, + h_t_d: &CudaSlice, + regime_h_d: &CudaSlice, + isv_dev_ptr: u64, + ) -> Result<()> { + let b = self.cfg.b_size; + let k = self.cfg.k; + debug_assert_eq!( + grad_pi_logits_d.len(), + b * N_ACTIONS, + "grad_pi_logits shape mismatch" + ); + debug_assert_eq!(h_t_d.len(), b * HIDDEN_DIM, "h_t shape mismatch"); + debug_assert_eq!( + regime_h_d.len(), + b * REGIME_DIM, + "regime_h shape mismatch" + ); + + let b_i = b as i32; + let k_i = k as i32; + + // ── Step 1: backward through mixture + per-head softmax ─────── + // Grid = (B), Block = (HIDDEN_DIM). Each thread covers one + // hidden-dim index across all (k, a). + { + let mut args = RawArgs::new(); + args.push_ptr(grad_pi_logits_d.raw_ptr()); + args.push_ptr(h_t_d.raw_ptr()); + args.push_ptr(self.w_heads_d.raw_ptr()); + args.push_ptr(self.gate_probs_d.raw_ptr()); + args.push_ptr(self.pi_probs_k_d.raw_ptr()); + args.push_i32(b_i); + args.push_i32(k_i); + args.push_ptr(self.grad_pi_logits_k_d.raw_ptr()); + args.push_ptr(self.grad_gate_probs_d.raw_ptr()); + args.push_ptr(self.grad_w_heads_pb_d.raw_ptr()); + args.push_ptr(self.grad_b_heads_pb_d.raw_ptr()); + args.push_ptr(self.grad_h_t_scratch_d.raw_ptr()); + let mut ptrs = args.build_arg_ptrs(); + unsafe { + raw_launch( + self.backward_pi_fn.cu_function(), + (b as u32, 1, 1), + (HIDDEN_DIM as u32, 1, 1), + 0, + self.raw_stream, + &mut ptrs[..args.len()], + ) + .map_err(|e| anyhow::anyhow!("multi_head_policy_backward_pi: {e:?}"))?; + } + } + + // ── Step 2: aux KL prior (additive into grad_pi_logits_k_d) ── + // Grid = (B, K), Block = (N_ACTIONS). Runs AFTER step 1 so the + // base softmax-Jacobian grad is already in place when we add the + // β-weighted aux contribution. β=0 → kernel short-circuits. + { + let mut args = RawArgs::new(); + args.push_ptr(self.pi_probs_k_d.raw_ptr()); + args.push_ptr(self.priors_d.raw_ptr()); + args.push_ptr(isv_dev_ptr); + args.push_i32(b_i); + args.push_i32(k_i); + args.push_ptr(self.grad_pi_logits_k_d.raw_ptr()); + let mut ptrs = args.build_arg_ptrs(); + unsafe { + raw_launch( + self.aux_prior_fn.cu_function(), + (b as u32, k as u32, 1), + (N_ACTIONS as u32, 1, 1), + 0, + self.raw_stream, + &mut ptrs[..args.len()], + ) + .map_err(|e| anyhow::anyhow!("multi_head_policy_aux_prior: {e:?}"))?; + } + } + + // ── Step 3: backward through gating softmax ─────────────────── + // Grid = (B), Block = (MAX_K_HEADS=8) — covers both K and REGIME_DIM + // (both ≤ MAX_K_HEADS). Single-stream hand-off from pi backward. + { + let mut args = RawArgs::new(); + args.push_ptr(self.grad_gate_probs_d.raw_ptr()); + args.push_ptr(self.gate_probs_d.raw_ptr()); + args.push_ptr(regime_h_d.raw_ptr()); + args.push_ptr(self.w_gate_d.raw_ptr()); + args.push_i32(b_i); + args.push_i32(k_i); + args.push_ptr(self.grad_gate_logits_d.raw_ptr()); + args.push_ptr(self.grad_w_gate_pb_d.raw_ptr()); + args.push_ptr(self.grad_b_gate_pb_d.raw_ptr()); + args.push_ptr(self.grad_regime_h_d.raw_ptr()); + let mut ptrs = args.build_arg_ptrs(); + unsafe { + raw_launch( + self.backward_gate_fn.cu_function(), + (b as u32, 1, 1), + (MAX_K_HEADS as u32, 1, 1), + 0, + self.raw_stream, + &mut ptrs[..args.len()], + ) + .map_err(|e| anyhow::anyhow!("multi_head_policy_backward_gate: {e:?}"))?; + } + } + + // ── Step 4: reduce per-batch grads across B → final shapes ──── + // Standard foxhunt reduce_axis0 pattern (see trainer integrated.rs). + reduce_axis0_local( + &self.stream, + &self.reduce_axis0_fn, + &self.grad_w_heads_pb_d, + b, + k * N_ACTIONS * HIDDEN_DIM, + &mut self.grad_w_heads_d, + )?; + reduce_axis0_local( + &self.stream, + &self.reduce_axis0_fn, + &self.grad_b_heads_pb_d, + b, + k * N_ACTIONS, + &mut self.grad_b_heads_d, + )?; + reduce_axis0_local( + &self.stream, + &self.reduce_axis0_fn, + &self.grad_w_gate_pb_d, + b, + k * REGIME_DIM, + &mut self.grad_w_gate_d, + )?; + reduce_axis0_local( + &self.stream, + &self.reduce_axis0_fn, + &self.grad_b_gate_pb_d, + b, + k, + &mut self.grad_b_gate_d, + )?; + + Ok(()) + } } // ── pinned-staging upload helper ───────────────────────────────────── @@ -402,3 +744,81 @@ fn upload(stream: &Arc, host: &[f32]) -> Result> { } Ok(dst) } + +// ── Per-head action priors ────────────────────────────────────────── +// +// Mirror the per-head init biases from `MultiHeadPolicy::new` so the aux +// KL term does not fight the asymmetry-break initialization. +// +// prior[0] = softmax([+0.5, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]) (SHORT) +// prior[1] = softmax([0, 0, 0, 0, 0, +0.5, +0.5, 0, 0, 0, 0]) (LONG) +// prior[2] = softmax([0, 0, +0.5, 0, 0, 0, 0, 0, 0, 0, 0]) (HOLD) +// prior[k≥3] = uniform = 1 / N_ACTIONS +// +// Returns a `K × N_ACTIONS` flat row-major host buffer. +fn build_priors(k: usize) -> Vec { + let mut priors = vec![0.0_f32; k * N_ACTIONS]; + for head in 0..k { + // Logit ramp matching the init bias (zero for unbiased actions). + let mut logit = vec![0.0_f32; N_ACTIONS]; + match head { + 0 => logit[0] = 0.5, + 1 => { + logit[5] = 0.5; + logit[6] = 0.5; + } + 2 => logit[2] = 0.5, + _ => {} // uniform — all zeros → softmax = 1/N + } + // softmax — single-batch row, no overflow risk at this scale. + let mx = logit.iter().cloned().fold(f32::NEG_INFINITY, f32::max); + let mut s = 0.0_f32; + let mut exps = [0.0_f32; N_ACTIONS]; + for (i, &l) in logit.iter().enumerate() { + let e = (l - mx).exp(); + exps[i] = e; + s += e; + } + for (i, e) in exps.iter().enumerate() { + priors[head * N_ACTIONS + i] = e / s; + } + } + priors +} + +// ── reduce_axis0 launcher (local copy of the trainer's helper) ───── +// +// Sums `[B × N]` per-batch grad scratch to `[N]` final form. Single +// block per output-tile column, 32×8 threads (REDAX0_TILE_J × REDAX0_TILE_B +// matching `cuda/reduce_axis0.cu`). No host-side helper import to keep +// `MultiHeadPolicy` independent of `trainer::integrated` internals. +fn reduce_axis0_local( + stream: &Arc, + reduce_fn: &CudaFunction, + per_batch: &CudaSlice, + b_size: usize, + n_tail: usize, + out: &mut CudaSlice, +) -> Result<()> { + debug_assert_eq!(per_batch.len(), b_size * n_tail); + debug_assert_eq!(out.len(), n_tail); + let grid_x = (n_tail as u32).div_ceil(32); + let mut args = RawArgs::new(); + args.push_ptr(per_batch.raw_ptr()); + args.push_i32(b_size as i32); + args.push_i32(n_tail as i32); + args.push_ptr(out.raw_ptr()); + let mut ptrs = args.build_arg_ptrs(); + unsafe { + raw_launch( + reduce_fn.cu_function(), + (grid_x, 1, 1), + (32, 8, 1), + 0, + stream.cu_stream(), + &mut ptrs[..args.len()], + ) + .map_err(|e| anyhow::anyhow!("reduce_axis0 (multi_head_policy): {e:?}"))?; + } + Ok(()) +} diff --git a/crates/ml-alpha/tests/multi_head_policy_invariants.rs b/crates/ml-alpha/tests/multi_head_policy_invariants.rs index 9b2d31580..0b2236d73 100644 --- a/crates/ml-alpha/tests/multi_head_policy_invariants.rs +++ b/crates/ml-alpha/tests/multi_head_policy_invariants.rs @@ -1,7 +1,6 @@ -//! GPU-oracle invariants for [`MultiHeadPolicy`] (Phase 2A-A). +//! GPU-oracle invariants for [`MultiHeadPolicy`] (Phase 2A-A + 2A-B). //! -//! Forward-pass invariants verified (no trainer dependency — tests stand -//! up an `MlDevice`, allocate inputs, run forward, read back outputs): +//! Phase 2A-A — forward-pass invariants (5 tests): //! //! 1. `gate_probs_sum_to_one` — softmax invariant on the gating head. //! 2. `pi_probs_mixture_sums_to_one` — exp(pi_logits[b]).sum() ≈ 1 for @@ -20,10 +19,23 @@ //! bit-equal outputs on the same input. Verifies no determinism //! regression vs `pearl_determinism_achieved`. //! +//! Phase 2A-B — backward-pass invariants (4 tests): +//! +//! 6. `backward_gradcheck_w_heads` — numerical gradient check on +//! `grad_W_heads` via 5-point finite-difference at random entries. +//! β=0 (aux KL disabled) for a clean Q-distill-chain-only test. +//! 7. `backward_gradcheck_w_gate` — same pattern for `grad_W_gate`. +//! 8. `aux_prior_grad_sums_to_zero_per_head` — the softmax-Jacobian +//! identity Σ_a (π · (log_ratio − KL))[a] = 0 must hold per (b, k). +//! 9. `backward_is_deterministic_across_contexts` — bit-equal `grad_W_heads` +//! + `grad_W_gate` across two fresh same-seed instances on identical +//! inputs. Determinism regression guard. +//! //! Per `feedback_no_cpu_test_fallbacks`: oracles are analytical (softmax -//! sum, per-head reduction, deterministic seed). +//! sum, per-head reduction, deterministic seed, finite-difference identity). //! Per `feedback_no_htod_htoh_only_mapped_pinned`: tests use -//! `write_slice_f32_d_pub` / `read_slice_d_pub` for all transfers. +//! `write_slice_f32_d_pub` / `read_slice_d_pub` for all transfers; the +//! ISV bus uses `MappedF32Buffer`. //! //! Run with: //! `cargo test -p ml-alpha --test multi_head_policy_invariants --release -- --ignored --nocapture` @@ -32,7 +44,9 @@ use anyhow::Result; use cudarc::driver::{CudaSlice, CudaStream}; use ml_alpha::cfc::snap_features::REGIME_DIM; use ml_alpha::heads::HIDDEN_DIM; +use ml_alpha::pinned_mem::MappedF32Buffer; use ml_alpha::rl::common::N_ACTIONS; +use ml_alpha::rl::isv_slots::{RL_POLICY_AUX_PRIOR_BETA_INDEX, RL_SLOTS_END}; use ml_alpha::rl::multi_head_policy::{MultiHeadPolicy, MultiHeadPolicyConfig}; use ml_alpha::trainer::integrated::{read_slice_d_pub, write_slice_f32_d_pub}; use ml_core::device::MlDevice; @@ -387,3 +401,688 @@ fn forward_is_deterministic_across_contexts() -> Result<()> { ); Ok(()) } + +// ───────────────────────────────────────────────────────────────────── +// Phase 2A-B — backward-pass tests +// ───────────────────────────────────────────────────────────────────── + +/// Allocate a mapped-pinned ISV bus, set `β = beta`. Returns the buffer +/// (caller keeps it alive for the lifetime of the kernel launches). +fn make_isv_with_beta(beta: f32) -> Result { + let isv = unsafe { MappedF32Buffer::new(RL_SLOTS_END) } + .map_err(|e| anyhow::anyhow!("isv MappedF32Buffer alloc: {e}"))?; + // Zero the whole bus deterministically. + for slot in 0..RL_SLOTS_END { + isv.write_record(slot, 0.0); + } + isv.write_record(RL_POLICY_AUX_PRIOR_BETA_INDEX, beta); + Ok(isv) +} + +/// Tolerance for the finite-difference gradcheck. +/// +/// ── Why ε=1e-2 ──────────────────────────────────────────────────────── +/// In fp32 the gradcheck error has TWO components: +/// * truncation: ~ε²·f'''/6 (decreases with smaller ε) +/// * fp32 noise: ~ulp(L)/(2ε·|grad|) ≈ 2^-23·|L|/(2ε·|grad|) +/// (INCREASES with smaller ε — catastrophic cancellation) +/// +/// The forward chain `pi_logits = log(softmax_K(gate_logits) · softmax_A(W_h·h))` +/// gives |L| ~ |log(1/N_ACTIONS)| × |grad_pi_logits| ≈ 22 · 0.3 = 6.6, with +/// `|grad| ~ 1e-2`. At ε=1e-3 noise dominates: ulp(L)/(2ε·|grad|) ≈ 4%, +/// matching the observed 1.7-1.9% relative error (averaged over probes). +/// +/// At ε=1e-2 the noise term drops by 10× (to ~0.4%) and truncation rises +/// by 100× — but for this chain truncation at ε=1e-2 is still well under +/// noise, so total error MINIMIZES around ε=1e-2 at ~0.5% relative. +/// Verified by `backward_gradcheck_w_heads_eps_sweep` / +/// `backward_gradcheck_w_gate_eps_sweep`: +/// +/// | ε | max_rel_err W_heads | max_rel_err W_gate | +/// |-------|---------------------|--------------------| +/// | 1e-2 | 1.6e-3 | 5.0e-3 | ← sweet spot +/// | 5e-3 | 1.3e-3 | 1.8e-2 | +/// | 1e-3 | 1.9e-2 | 3.3e-2 | ← old setting +/// | 5e-4 | 1.0e-2 | 4.3e-2 | +/// | 1e-4 | 1.0e-1 | 4.7e-1 | ← cancellation +/// +/// Since the error GROWS as ε shrinks, it is NOT truncation (which would +/// shrink ∝ ε²). It is fp32 cancellation in the per-output `pi_logits` +/// values. The analytical gradient is therefore correct to ≤1% — anything +/// larger would not vanish at the cancellation-minimum ε. +/// +/// Tolerance 1e-2 matches `tests/backward_finite_diff.rs::max_relative`. +const GRADCHECK_REL_TOL: f32 = 1e-2; +const GRADCHECK_ABS_TOL: f32 = 5e-4; +const GRADCHECK_EPS: f32 = 1e-2; + +/// Build a small reproducible `(h_t, regime_h, grad_pi_logits)` triple. The +/// `grad_pi_logits` is the "outer loss" gradient — any nonzero vector +/// will do for gradcheck as long as it's deterministic. +fn build_gradcheck_inputs( + seed: u64, + b_size: usize, +) -> (Vec, Vec, Vec) { + use rand::{Rng, SeedableRng}; + use rand_chacha::ChaCha8Rng; + let mut r = ChaCha8Rng::seed_from_u64(seed); + let h_t: Vec = (0..b_size * HIDDEN_DIM) + .map(|_| r.gen_range(-0.3..0.3)) + .collect(); + let regime: Vec = (0..b_size * REGIME_DIM) + .map(|_| r.gen_range(-0.5..0.5)) + .collect(); + let grad_pi_logits: Vec = (0..b_size * N_ACTIONS) + .map(|_| r.gen_range(-0.5..0.5)) + .collect(); + (h_t, regime, grad_pi_logits) +} + +/// Compute the "loss" implied by an `(outputs, grad_outputs)` pair: +/// L(outputs) = Σ_i outputs[i] · grad_outputs[i] +/// This is the canonical pairing such that dL/d(outputs) = grad_outputs. +/// Finite-differencing this L through a forward pass validates the +/// analytical backward of `grad_outputs → grad_inputs`. +fn linear_loss(outputs: &[f32], grad_outputs: &[f32]) -> f64 { + debug_assert_eq!(outputs.len(), grad_outputs.len()); + let mut s = 0.0_f64; + for (o, g) in outputs.iter().zip(grad_outputs.iter()) { + s += (*o as f64) * (*g as f64); + } + s +} + +#[test] +#[ignore = "requires CUDA (MlDevice::cuda(0))"] +fn backward_gradcheck_w_heads() -> Result<()> { + let Some(dev) = build_device() else { + return Ok(()); + }; + let stream = dev.cuda_stream()?.clone(); + let b_size = 2usize; + let k = 2usize; + + // β=0 disables the aux KL contribution so this test isolates the + // softmax-mixture chain rule on its own. Aux is checked separately + // by `aux_prior_grad_sums_to_zero_per_head`. + let isv = make_isv_with_beta(0.0)?; + let isv_dev_ptr = isv.dev_ptr; + + let mut mhp = MultiHeadPolicy::new( + &dev, + MultiHeadPolicyConfig { + k, + b_size, + seed: 0x6E_AD_AC_AC, + }, + )?; + + let (h_t_host, regime_host, grad_pi_logits_host) = + build_gradcheck_inputs(0xF00D_BABE, b_size); + let h_t_d = upload_f32(&stream, &h_t_host)?; + let regime_d = upload_f32(&stream, ®ime_host)?; + let grad_pi_logits_d = upload_f32(&stream, &grad_pi_logits_host)?; + let mut pi_logits_out = stream.alloc_zeros::(b_size * N_ACTIONS)?; + + // 1. Forward + backward at the baseline W_heads to get analytical grad. + mhp.forward(&h_t_d, ®ime_d, &mut pi_logits_out)?; + mhp.backward(&grad_pi_logits_d, &h_t_d, ®ime_d, isv_dev_ptr)?; + + let analytical = read_slice_d_pub( + &stream, + &mhp.grad_w_heads_d, + k * N_ACTIONS * HIDDEN_DIM, + )?; + let w_heads_baseline = read_slice_d_pub( + &stream, + &mhp.w_heads_d, + k * N_ACTIONS * HIDDEN_DIM, + )?; + + // 2. Choose 5 deterministic perturbation indices. + use rand::{Rng, SeedableRng}; + use rand_chacha::ChaCha8Rng; + let mut r = ChaCha8Rng::seed_from_u64(0x1234_5678); + let n_params = k * N_ACTIONS * HIDDEN_DIM; + let probe_idxs: Vec = (0..5).map(|_| r.gen_range(0..n_params)).collect(); + + let mut max_abs_err = 0.0_f32; + let mut max_rel_err = 0.0_f32; + + for &idx in &probe_idxs { + // L(+ε) + let mut w_plus = w_heads_baseline.clone(); + w_plus[idx] += GRADCHECK_EPS; + write_slice_f32_d_pub(&stream, &w_plus, &mut mhp.w_heads_d)?; + mhp.forward(&h_t_d, ®ime_d, &mut pi_logits_out)?; + let pl_plus = read_slice_d_pub(&stream, &pi_logits_out, b_size * N_ACTIONS)?; + let l_plus = linear_loss(&pl_plus, &grad_pi_logits_host); + + // L(-ε) + let mut w_minus = w_heads_baseline.clone(); + w_minus[idx] -= GRADCHECK_EPS; + write_slice_f32_d_pub(&stream, &w_minus, &mut mhp.w_heads_d)?; + mhp.forward(&h_t_d, ®ime_d, &mut pi_logits_out)?; + let pl_minus = read_slice_d_pub(&stream, &pi_logits_out, b_size * N_ACTIONS)?; + let l_minus = linear_loss(&pl_minus, &grad_pi_logits_host); + + let num_grad = ((l_plus - l_minus) / (2.0 * GRADCHECK_EPS as f64)) as f32; + let ana_grad = analytical[idx]; + let abs_err = (num_grad - ana_grad).abs(); + let denom = num_grad.abs().max(ana_grad.abs()).max(1e-8); + let rel_err = abs_err / denom; + + if abs_err > max_abs_err { + max_abs_err = abs_err; + } + if rel_err > max_rel_err { + max_rel_err = rel_err; + } + + // Accept if EITHER relative OR absolute tolerance is met. + let pass = rel_err < GRADCHECK_REL_TOL || abs_err < GRADCHECK_ABS_TOL; + assert!( + pass, + "grad_W_heads[{idx}]: analytical={ana_grad}, numerical={num_grad}, \ + abs_err={abs_err}, rel_err={rel_err} (tol rel={GRADCHECK_REL_TOL} abs={GRADCHECK_ABS_TOL})" + ); + + // Restore W_heads to baseline. + write_slice_f32_d_pub(&stream, &w_heads_baseline, &mut mhp.w_heads_d)?; + } + + eprintln!( + "PASS — grad_W_heads gradcheck (5 probes, ε={GRADCHECK_EPS}): max_abs_err={max_abs_err:.3e}, max_rel_err={max_rel_err:.3e}" + ); + Ok(()) +} + +/// Diagnostic: sweep ε across {1e-2, 5e-3, 1e-3, 5e-4, 1e-4} and check +/// 10 probe indices to characterize whether the gradcheck error is +/// truncation (∝ ε²) or a real bug (~constant across ε). +/// +/// Also reports per-(k,a,h) localization of the maximum error to identify +/// whether errors cluster around saturated softmax / init-bias positions. +#[test] +#[ignore = "requires CUDA (MlDevice::cuda(0))"] +fn backward_gradcheck_w_heads_eps_sweep() -> Result<()> { + let Some(dev) = build_device() else { + return Ok(()); + }; + let stream = dev.cuda_stream()?.clone(); + let b_size = 2usize; + let k = 2usize; + + let isv = make_isv_with_beta(0.0)?; + let isv_dev_ptr = isv.dev_ptr; + + let mut mhp = MultiHeadPolicy::new( + &dev, + MultiHeadPolicyConfig { + k, + b_size, + seed: 0x6E_AD_AC_AC, + }, + )?; + + let (h_t_host, regime_host, grad_pi_logits_host) = + build_gradcheck_inputs(0xF00D_BABE, b_size); + let h_t_d = upload_f32(&stream, &h_t_host)?; + let regime_d = upload_f32(&stream, ®ime_host)?; + let grad_pi_logits_d = upload_f32(&stream, &grad_pi_logits_host)?; + let mut pi_logits_out = stream.alloc_zeros::(b_size * N_ACTIONS)?; + + // Analytical grad (computed once). + mhp.forward(&h_t_d, ®ime_d, &mut pi_logits_out)?; + mhp.backward(&grad_pi_logits_d, &h_t_d, ®ime_d, isv_dev_ptr)?; + let analytical = read_slice_d_pub( + &stream, + &mhp.grad_w_heads_d, + k * N_ACTIONS * HIDDEN_DIM, + )?; + let w_heads_baseline = read_slice_d_pub( + &stream, + &mhp.w_heads_d, + k * N_ACTIONS * HIDDEN_DIM, + )?; + + // 10 deterministic probe indices — same seed as production gradcheck + // so we land on a comparable set. + use rand::{Rng, SeedableRng}; + use rand_chacha::ChaCha8Rng; + let mut r = ChaCha8Rng::seed_from_u64(0x1234_5678); + let n_params = k * N_ACTIONS * HIDDEN_DIM; + let probe_idxs: Vec = (0..10).map(|_| r.gen_range(0..n_params)).collect(); + + let epsilons: [f32; 5] = [1e-2, 5e-3, 1e-3, 5e-4, 1e-4]; + eprintln!("\n── ε-sweep gradcheck on grad_W_heads ─────────────────"); + eprintln!("(k = head, a = action, h = hidden-dim)"); + eprintln!("| ε | max_abs_err | max_rel_err | argmax (k,a,h) | ana@argmax | num@argmax |"); + + for &eps in epsilons.iter() { + let mut max_abs_err = 0.0_f32; + let mut max_rel_err = 0.0_f32; + let mut argmax_idx = 0usize; + let mut ana_at = 0.0_f32; + let mut num_at = 0.0_f32; + + for &idx in &probe_idxs { + let mut w_plus = w_heads_baseline.clone(); + w_plus[idx] += eps; + write_slice_f32_d_pub(&stream, &w_plus, &mut mhp.w_heads_d)?; + mhp.forward(&h_t_d, ®ime_d, &mut pi_logits_out)?; + let pl_plus = read_slice_d_pub(&stream, &pi_logits_out, b_size * N_ACTIONS)?; + let l_plus = linear_loss(&pl_plus, &grad_pi_logits_host); + + let mut w_minus = w_heads_baseline.clone(); + w_minus[idx] -= eps; + write_slice_f32_d_pub(&stream, &w_minus, &mut mhp.w_heads_d)?; + mhp.forward(&h_t_d, ®ime_d, &mut pi_logits_out)?; + let pl_minus = read_slice_d_pub(&stream, &pi_logits_out, b_size * N_ACTIONS)?; + let l_minus = linear_loss(&pl_minus, &grad_pi_logits_host); + + let num_grad = ((l_plus - l_minus) / (2.0 * eps as f64)) as f32; + let ana_grad = analytical[idx]; + let abs_err = (num_grad - ana_grad).abs(); + let denom = num_grad.abs().max(ana_grad.abs()).max(1e-8); + let rel_err = abs_err / denom; + + if abs_err > max_abs_err { + max_abs_err = abs_err; + } + if rel_err > max_rel_err { + max_rel_err = rel_err; + argmax_idx = idx; + ana_at = ana_grad; + num_at = num_grad; + } + + write_slice_f32_d_pub(&stream, &w_heads_baseline, &mut mhp.w_heads_d)?; + } + + // Decode (k, a, h) from flat index. Layout: [K × N_ACTIONS × HIDDEN_DIM]. + let h_idx = argmax_idx % HIDDEN_DIM; + let a_idx = (argmax_idx / HIDDEN_DIM) % N_ACTIONS; + let k_idx = argmax_idx / (HIDDEN_DIM * N_ACTIONS); + eprintln!( + "| {eps:.0e} | {max_abs_err:.3e} | {max_rel_err:.3e} | ({k_idx},{a_idx:>2},{h_idx:>3}) | {ana_at:+.4e} | {num_at:+.4e} |" + ); + } + Ok(()) +} + +#[test] +#[ignore = "requires CUDA (MlDevice::cuda(0))"] +fn backward_gradcheck_w_gate_eps_sweep() -> Result<()> { + let Some(dev) = build_device() else { + return Ok(()); + }; + let stream = dev.cuda_stream()?.clone(); + let b_size = 2usize; + let k = 2usize; + + let isv = make_isv_with_beta(0.0)?; + let isv_dev_ptr = isv.dev_ptr; + + let mut mhp = MultiHeadPolicy::new( + &dev, + MultiHeadPolicyConfig { + k, + b_size, + seed: 0xAA_BB_CC_DD, + }, + )?; + + let w_gate_init: Vec = (0..k * REGIME_DIM) + .map(|i| ((i as f32) * 0.37).sin() * 0.5) + .collect(); + write_slice_f32_d_pub(&stream, &w_gate_init, &mut mhp.w_gate_d)?; + + let (h_t_host, regime_host, grad_pi_logits_host) = + build_gradcheck_inputs(0xC0DE_BEEF, b_size); + let h_t_d = upload_f32(&stream, &h_t_host)?; + let regime_d = upload_f32(&stream, ®ime_host)?; + let grad_pi_logits_d = upload_f32(&stream, &grad_pi_logits_host)?; + let mut pi_logits_out = stream.alloc_zeros::(b_size * N_ACTIONS)?; + + mhp.forward(&h_t_d, ®ime_d, &mut pi_logits_out)?; + mhp.backward(&grad_pi_logits_d, &h_t_d, ®ime_d, isv_dev_ptr)?; + let analytical = read_slice_d_pub(&stream, &mhp.grad_w_gate_d, k * REGIME_DIM)?; + let w_gate_baseline = read_slice_d_pub(&stream, &mhp.w_gate_d, k * REGIME_DIM)?; + + // For W_gate (K * REGIME_DIM = 2*6 = 12 params), probe ALL entries. + let n_params = k * REGIME_DIM; + let probe_idxs: Vec = (0..n_params).collect(); + + let epsilons: [f32; 5] = [1e-2, 5e-3, 1e-3, 5e-4, 1e-4]; + eprintln!("\n── ε-sweep gradcheck on grad_W_gate ──────────────────"); + eprintln!("(k = head, r = regime-dim)"); + eprintln!("| ε | max_abs_err | max_rel_err | argmax (k,r) | ana@argmax | num@argmax |"); + + for &eps in epsilons.iter() { + let mut max_abs_err = 0.0_f32; + let mut max_rel_err = 0.0_f32; + let mut argmax_idx = 0usize; + let mut ana_at = 0.0_f32; + let mut num_at = 0.0_f32; + + for &idx in &probe_idxs { + let mut w_plus = w_gate_baseline.clone(); + w_plus[idx] += eps; + write_slice_f32_d_pub(&stream, &w_plus, &mut mhp.w_gate_d)?; + mhp.forward(&h_t_d, ®ime_d, &mut pi_logits_out)?; + let pl_plus = read_slice_d_pub(&stream, &pi_logits_out, b_size * N_ACTIONS)?; + let l_plus = linear_loss(&pl_plus, &grad_pi_logits_host); + + let mut w_minus = w_gate_baseline.clone(); + w_minus[idx] -= eps; + write_slice_f32_d_pub(&stream, &w_minus, &mut mhp.w_gate_d)?; + mhp.forward(&h_t_d, ®ime_d, &mut pi_logits_out)?; + let pl_minus = read_slice_d_pub(&stream, &pi_logits_out, b_size * N_ACTIONS)?; + let l_minus = linear_loss(&pl_minus, &grad_pi_logits_host); + + let num_grad = ((l_plus - l_minus) / (2.0 * eps as f64)) as f32; + let ana_grad = analytical[idx]; + let abs_err = (num_grad - ana_grad).abs(); + let denom = num_grad.abs().max(ana_grad.abs()).max(1e-8); + let rel_err = abs_err / denom; + + if abs_err > max_abs_err { + max_abs_err = abs_err; + } + if rel_err > max_rel_err { + max_rel_err = rel_err; + argmax_idx = idx; + ana_at = ana_grad; + num_at = num_grad; + } + + write_slice_f32_d_pub(&stream, &w_gate_baseline, &mut mhp.w_gate_d)?; + } + + let r_idx = argmax_idx % REGIME_DIM; + let k_idx = argmax_idx / REGIME_DIM; + eprintln!( + "| {eps:.0e} | {max_abs_err:.3e} | {max_rel_err:.3e} | ({k_idx},{r_idx}) | {ana_at:+.4e} | {num_at:+.4e} |" + ); + } + Ok(()) +} + +#[test] +#[ignore = "requires CUDA (MlDevice::cuda(0))"] +fn backward_gradcheck_w_gate() -> Result<()> { + let Some(dev) = build_device() else { + return Ok(()); + }; + let stream = dev.cuda_stream()?.clone(); + let b_size = 2usize; + let k = 2usize; + + let isv = make_isv_with_beta(0.0)?; + let isv_dev_ptr = isv.dev_ptr; + + let mut mhp = MultiHeadPolicy::new( + &dev, + MultiHeadPolicyConfig { + k, + b_size, + seed: 0xAA_BB_CC_DD, + }, + )?; + + // Make W_gate moderately nonzero so the gate is genuinely sensitive to + // perturbations (otherwise the gating gradient flows through ~uniform + // softmax and the FD signal is too small to verify). + let w_gate_init: Vec = (0..k * REGIME_DIM) + .map(|i| ((i as f32) * 0.37).sin() * 0.5) + .collect(); + write_slice_f32_d_pub(&stream, &w_gate_init, &mut mhp.w_gate_d)?; + + let (h_t_host, regime_host, grad_pi_logits_host) = + build_gradcheck_inputs(0xC0DE_BEEF, b_size); + let h_t_d = upload_f32(&stream, &h_t_host)?; + let regime_d = upload_f32(&stream, ®ime_host)?; + let grad_pi_logits_d = upload_f32(&stream, &grad_pi_logits_host)?; + let mut pi_logits_out = stream.alloc_zeros::(b_size * N_ACTIONS)?; + + mhp.forward(&h_t_d, ®ime_d, &mut pi_logits_out)?; + mhp.backward(&grad_pi_logits_d, &h_t_d, ®ime_d, isv_dev_ptr)?; + + let analytical = read_slice_d_pub(&stream, &mhp.grad_w_gate_d, k * REGIME_DIM)?; + let w_gate_baseline = read_slice_d_pub(&stream, &mhp.w_gate_d, k * REGIME_DIM)?; + + use rand::{Rng, SeedableRng}; + use rand_chacha::ChaCha8Rng; + let mut r = ChaCha8Rng::seed_from_u64(0x9999_0000); + let n_params = k * REGIME_DIM; + let probe_idxs: Vec = (0..5).map(|_| r.gen_range(0..n_params)).collect(); + + let mut max_abs_err = 0.0_f32; + let mut max_rel_err = 0.0_f32; + + for &idx in &probe_idxs { + let mut w_plus = w_gate_baseline.clone(); + w_plus[idx] += GRADCHECK_EPS; + write_slice_f32_d_pub(&stream, &w_plus, &mut mhp.w_gate_d)?; + mhp.forward(&h_t_d, ®ime_d, &mut pi_logits_out)?; + let pl_plus = read_slice_d_pub(&stream, &pi_logits_out, b_size * N_ACTIONS)?; + let l_plus = linear_loss(&pl_plus, &grad_pi_logits_host); + + let mut w_minus = w_gate_baseline.clone(); + w_minus[idx] -= GRADCHECK_EPS; + write_slice_f32_d_pub(&stream, &w_minus, &mut mhp.w_gate_d)?; + mhp.forward(&h_t_d, ®ime_d, &mut pi_logits_out)?; + let pl_minus = read_slice_d_pub(&stream, &pi_logits_out, b_size * N_ACTIONS)?; + let l_minus = linear_loss(&pl_minus, &grad_pi_logits_host); + + let num_grad = ((l_plus - l_minus) / (2.0 * GRADCHECK_EPS as f64)) as f32; + let ana_grad = analytical[idx]; + let abs_err = (num_grad - ana_grad).abs(); + let denom = num_grad.abs().max(ana_grad.abs()).max(1e-8); + let rel_err = abs_err / denom; + + if abs_err > max_abs_err { + max_abs_err = abs_err; + } + if rel_err > max_rel_err { + max_rel_err = rel_err; + } + + let pass = rel_err < GRADCHECK_REL_TOL || abs_err < GRADCHECK_ABS_TOL; + assert!( + pass, + "grad_W_gate[{idx}]: analytical={ana_grad}, numerical={num_grad}, \ + abs_err={abs_err}, rel_err={rel_err} (tol rel={GRADCHECK_REL_TOL} abs={GRADCHECK_ABS_TOL})" + ); + + write_slice_f32_d_pub(&stream, &w_gate_baseline, &mut mhp.w_gate_d)?; + } + + eprintln!( + "PASS — grad_W_gate gradcheck (5 probes, ε={GRADCHECK_EPS}): max_abs_err={max_abs_err:.3e}, max_rel_err={max_rel_err:.3e}" + ); + Ok(()) +} + +#[test] +#[ignore = "requires CUDA (MlDevice::cuda(0))"] +fn aux_prior_grad_sums_to_zero_per_head() -> Result<()> { + let Some(dev) = build_device() else { + return Ok(()); + }; + let stream = dev.cuda_stream()?.clone(); + let b_size = 3usize; + let k = 3usize; + + // Strong aux β to make the contribution dominant + measurable. + let isv = make_isv_with_beta(1.0)?; + let isv_dev_ptr = isv.dev_ptr; + + let mut mhp = MultiHeadPolicy::new( + &dev, + MultiHeadPolicyConfig { + k, + b_size, + seed: 0x4711_4711, + }, + )?; + + // The default `W_heads` is zero, which makes per-head softmax outputs + // EXACTLY match the priors (by construction — priors mirror init biases), + // so the aux KL contribution would be identically zero and the test + // wouldn't exercise the kernel. Override `W_heads` with deterministic + // non-zero values so π diverges from the prior at this `h_t`. + use rand::{Rng, SeedableRng}; + use rand_chacha::ChaCha8Rng; + let mut r = ChaCha8Rng::seed_from_u64(0xDEAD_BEEF); + let w_heads_init: Vec = (0..k * N_ACTIONS * HIDDEN_DIM) + .map(|_| r.gen_range(-0.1..0.1)) + .collect(); + write_slice_f32_d_pub(&stream, &w_heads_init, &mut mhp.w_heads_d)?; + + let h_t_host: Vec = (0..b_size * HIDDEN_DIM) + .map(|_| r.gen_range(-0.5..0.5)) + .collect(); + let regime_host: Vec = (0..b_size * REGIME_DIM) + .map(|_| r.gen_range(-0.5..0.5)) + .collect(); + // ZERO grad_pi_logits — so the only source of `grad_pi_logits_k` is the + // aux KL kernel; the softmax-Jacobian sum-to-zero identity then must + // hold per (b, k). + let grad_pi_logits_host = vec![0.0_f32; b_size * N_ACTIONS]; + + let h_t_d = upload_f32(&stream, &h_t_host)?; + let regime_d = upload_f32(&stream, ®ime_host)?; + let grad_pi_logits_d = upload_f32(&stream, &grad_pi_logits_host)?; + let mut pi_logits_out = stream.alloc_zeros::(b_size * N_ACTIONS)?; + + mhp.forward(&h_t_d, ®ime_d, &mut pi_logits_out)?; + mhp.backward(&grad_pi_logits_d, &h_t_d, ®ime_d, isv_dev_ptr)?; + + let grad_pi_logits_k = read_slice_d_pub( + &stream, + &mhp.grad_pi_logits_k_d, + b_size * k * N_ACTIONS, + )?; + + // Mathematical identity: with β>0 and a softmax-parameterized distribution, + // Σ_a (π_a · (log(π_a/prior_a) − KL_k)) + // = Σ_a π_a · log(π_a/prior_a) − KL_k · Σ_a π_a + // = KL_k − KL_k = 0 + // The tolerance reflects fp32 accumulation noise across N_ACTIONS=11 terms. + let tol = 1e-5_f32; + let mut max_sum_abs = 0.0_f32; + for b in 0..b_size { + for head in 0..k { + let base = (b * k + head) * N_ACTIONS; + let row = &grad_pi_logits_k[base..base + N_ACTIONS]; + let s: f32 = row.iter().sum(); + if s.abs() > max_sum_abs { + max_sum_abs = s.abs(); + } + assert!( + s.abs() < tol, + "Σ_a aux_grad[b={b}, k={head}, a]={s} expected 0 (row={row:?})" + ); + } + } + + // Sanity: at least one entry must be NONzero — otherwise the aux + // kernel silently no-op'd (which would mask a real bug). + let any_nonzero = grad_pi_logits_k.iter().any(|x| x.abs() > 1e-7); + assert!( + any_nonzero, + "aux prior produced all-zero grad with β=1.0 — kernel did not fire" + ); + + eprintln!( + "PASS — aux prior softmax-Jacobian sum-to-zero identity holds per (b, k) (max |Σ|={max_sum_abs:.3e})" + ); + Ok(()) +} + +#[test] +#[ignore = "requires CUDA (MlDevice::cuda(0))"] +fn backward_is_deterministic_across_contexts() -> Result<()> { + let Some(dev) = build_device() else { + return Ok(()); + }; + let stream = dev.cuda_stream()?.clone(); + let b_size = 4usize; + let k = 3usize; + let seed = 0xCAFE_F00D; + + // Both runs use the SAME ISV (same β) for parity. + let isv = make_isv_with_beta(0.05)?; + let isv_dev_ptr = isv.dev_ptr; + + let (h_t_host, regime_host, grad_pi_logits_host) = + build_gradcheck_inputs(0x1357_2468, b_size); + let h_t_d = upload_f32(&stream, &h_t_host)?; + let regime_d = upload_f32(&stream, ®ime_host)?; + let grad_pi_logits_d = upload_f32(&stream, &grad_pi_logits_host)?; + + // Run A + let mut mhp_a = MultiHeadPolicy::new( + &dev, + MultiHeadPolicyConfig { + k, + b_size, + seed, + }, + )?; + let mut pi_logits_out_a = stream.alloc_zeros::(b_size * N_ACTIONS)?; + mhp_a.forward(&h_t_d, ®ime_d, &mut pi_logits_out_a)?; + mhp_a.backward(&grad_pi_logits_d, &h_t_d, ®ime_d, isv_dev_ptr)?; + let gw_heads_a = read_slice_d_pub(&stream, &mhp_a.grad_w_heads_d, k * N_ACTIONS * HIDDEN_DIM)?; + let gw_gate_a = read_slice_d_pub(&stream, &mhp_a.grad_w_gate_d, k * REGIME_DIM)?; + let ght_a = read_slice_d_pub(&stream, &mhp_a.grad_h_t_scratch_d, b_size * HIDDEN_DIM)?; + + // Run B — fresh instance, identical seed + ISV + let mut mhp_b = MultiHeadPolicy::new( + &dev, + MultiHeadPolicyConfig { + k, + b_size, + seed, + }, + )?; + let mut pi_logits_out_b = stream.alloc_zeros::(b_size * N_ACTIONS)?; + mhp_b.forward(&h_t_d, ®ime_d, &mut pi_logits_out_b)?; + mhp_b.backward(&grad_pi_logits_d, &h_t_d, ®ime_d, isv_dev_ptr)?; + let gw_heads_b = read_slice_d_pub(&stream, &mhp_b.grad_w_heads_d, k * N_ACTIONS * HIDDEN_DIM)?; + let gw_gate_b = read_slice_d_pub(&stream, &mhp_b.grad_w_gate_d, k * REGIME_DIM)?; + let ght_b = read_slice_d_pub(&stream, &mhp_b.grad_h_t_scratch_d, b_size * HIDDEN_DIM)?; + + // Bit-equal across all three backward outputs. + for (i, (a, b)) in gw_heads_a.iter().zip(gw_heads_b.iter()).enumerate() { + assert_eq!( + a.to_bits(), + b.to_bits(), + "grad_W_heads[{i}] non-deterministic: a={a} b={b}" + ); + } + for (i, (a, b)) in gw_gate_a.iter().zip(gw_gate_b.iter()).enumerate() { + assert_eq!( + a.to_bits(), + b.to_bits(), + "grad_W_gate[{i}] non-deterministic: a={a} b={b}" + ); + } + for (i, (a, b)) in ght_a.iter().zip(ght_b.iter()).enumerate() { + assert_eq!( + a.to_bits(), + b.to_bits(), + "grad_h_t_scratch[{i}] non-deterministic: a={a} b={b}" + ); + } + eprintln!( + "PASS — backward bit-equal across two fresh contexts with seed=0x{seed:X} \ + (grad_W_heads={}, grad_W_gate={}, grad_h_t_scratch={})", + gw_heads_a.len(), + gw_gate_a.len(), + ght_a.len() + ); + Ok(()) +}