feat(ml-alpha): Phase 2A-B — MultiHeadPolicy backward + aux KL prior
Backward pass through the K-head mixture + per-head KL prior
regularization. Components still inert — trainer integration is
Phase 2A-C.
* multi_head_policy_backward.cu: two kernels.
- backward_pi (grid=(B), block=(HIDDEN_DIM=128)): full chain rule
log_grad → mixture decomposition → per-head softmax-Jacobian
→ W_heads/b_heads/h_t grads. Per-batch grad scratch +
reduce_axis0 reduction; no atomicAdd. Sole-writer per (b,k,a,h).
- backward_gate (grid=(B), block=(K=8)): gating softmax-Jacobian +
W_gate/b_gate grad scratch + grad_regime_h (computed but NOT
routed upstream — regime_h is loader-precomputed).
* multi_head_policy_aux_prior.cu: per-head KL prior softmax-Jacobian
grad β·π·(log_ratio − KL) additively accumulated into
grad_pi_logits_k. Read-modify-write race-free (sequential same
stream after backward_pi).
* MultiHeadPolicy::backward(grad_pi_logits, h_t, regime_h, isv) —
launches backward_pi → aux_prior → backward_gate → 4× reduce_axis0
on the same stream. Priors are computed at new() as softmax of
the per-head init bias vectors, so KL=0 at initialization and
the aux term only fires once Adam moves W_heads off zero.
* 4 new GPU-oracle invariant tests (9/9 total):
- backward_gradcheck_w_heads (W_heads central difference)
- backward_gradcheck_w_gate (W_gate central difference)
- aux_prior_grad_sums_to_zero_per_head (softmax-Jacobian invariant)
- backward_is_deterministic_across_contexts
* 2 diagnostic tests (no asserts, print-only):
- backward_gradcheck_w_heads_eps_sweep
- backward_gradcheck_w_gate_eps_sweep
## Math investigation (gradcheck error chase)
Initial gradcheck at ε=1e-3 showed 1.7–1.9% relative error which
SHOULD have meant a chain-rule bug. Investigation:
1. Re-derived the analytical backward chain from first principles
(8 chain steps) and diff'd against actual kernel source line
by line. No discrepancies — math is implemented exactly as
derived. ε in log-divisor (1e-12) matches forward.
2. ε-sweep characterization (b=2, k=2):
| ε | max_rel_err W_heads | max_rel_err W_gate |
|------|---------------------|--------------------|
| 1e-2 | 1.6e-3 | 5.0e-3 |
| 5e-3 | 1.3e-3 | 1.8e-2 |
| 1e-3 | 1.9e-2 | 3.3e-2 |
| 5e-4 | 1.0e-2 | 4.3e-2 |
| 1e-4 | 1.0e-1 | 4.7e-1 |
Error INCREASES as ε shrinks — opposite of O(ε²) truncation.
Unambiguous fingerprint of fp32 catastrophic cancellation in
(L_plus − L_minus)/2ε with |L| ≈ 6.6, |grad| ≈ 1e-2:
noise ≈ ulp(L)/(2ε·|grad|) ≈ 2⁻²³·6.6/(2ε·0.01)
ε=1e-3 → 4% (matches 1.9% observed), ε=1e-2 → 0.4%, ε=1e-4 → 40%.
3. Max-error indices random (not init-bias positions, not saturated
softmax) — consistent with noise scatter, not systematic bug.
Verdict: TRUNCATION_CONFIRMED (specifically fp32 cancellation, not
chain-rule truncation). Math is correct.
Fix: ε bumped 1e-3 → 1e-2 (the cancellation/truncation sweet
spot for fp32 chained softmax). Tolerance restored to foxhunt-
standard 1e-2 / 1e-3 (was 5e-2 / 5e-4 — that was a workaround for
the noise-dominated ε, not justified). Measured error now
1.6e-3 / 5.0e-3 — comfortably under tolerance. ε-sweep diagnostic
tests added as a permanent reproducibility anchor so future
gradcheck regressions are easy to diagnose.
## Verification
* `cargo test multi_head_policy_invariants --release -- --ignored`:
11/11 PASS (9 invariants + 2 diagnostics).
* `./scripts/determinism-check.sh --quick`: exit 0
(pipeline unchanged — backward not yet wired).
* 0 atomicAdd, 0 raw memcpy_htod/dtoh, 0 TODO/FIXME.
## Linked
* Phase 2A-A: 0b3e40150
* Plan: docs/superpowers/plans/2026-06-03-multi-head-policy-implementation.md
* Spec ADDENDUM: docs/superpowers/specs/2026-06-02-multi-head-policy-with-r-multiple.md
Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
This commit is contained in:
jgrusewski
@@ -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
|
||||
|
||||
91
crates/ml-alpha/cuda/multi_head_policy_aux_prior.cu
Normal file
91
crates/ml-alpha/cuda/multi_head_policy_aux_prior.cu
Normal file
@@ -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 <stdint.h>
|
||||
|
||||
#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;
|
||||
}
|
||||
300
crates/ml-alpha/cuda/multi_head_policy_backward.cu
Normal file
300
crates/ml-alpha/cuda/multi_head_policy_backward.cu
Normal file
@@ -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 <stdint.h>
|
||||
|
||||
#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<K * N_ACTIONS additionally write grad_b_heads_pb[b, k, a].
|
||||
if (c < HIDDEN_DIM) {
|
||||
const float h_bc = h_t[b * HIDDEN_DIM + c];
|
||||
float gh_acc = 0.0f;
|
||||
for (int k = 0; k < K; ++k) {
|
||||
for (int a = 0; a < N_ACTIONS; ++a) {
|
||||
const int kak = k * N_ACTIONS + a;
|
||||
const float g = s_grad_pi_logits_k[kak];
|
||||
// grad_W_heads_pb[b, k, a, c] = g * h[b, c]
|
||||
grad_W_heads_pb[(long long)((b * K + k) * N_ACTIONS + a) * HIDDEN_DIM + c]
|
||||
= g * h_bc;
|
||||
// Accumulate grad_h_t[b, c] = Σ_{k,a} g · W_heads[k, a, c]
|
||||
gh_acc += g * W_heads[(long long)((k * N_ACTIONS) + a) * HIDDEN_DIM + c];
|
||||
}
|
||||
}
|
||||
// OVERWRITE — caller folds via grad_h_accumulate_scaled.
|
||||
grad_h_t_scratch[b * HIDDEN_DIM + c] = gh_acc;
|
||||
}
|
||||
|
||||
// Bias grads — sole-writer per (b, k, a). Use the first K*N_ACTIONS
|
||||
// threads to cover the per-(k, a) slots in one pass.
|
||||
if (c < K * N_ACTIONS) {
|
||||
const int k = c / N_ACTIONS;
|
||||
const int a = c % N_ACTIONS;
|
||||
grad_b_heads_pb[(b * K + k) * N_ACTIONS + a] =
|
||||
s_grad_pi_logits_k[k * N_ACTIONS + a];
|
||||
}
|
||||
}
|
||||
|
||||
// ─────────────────────────────────────────────────────────────────────
|
||||
// Kernel 2: backward through the gating softmax.
|
||||
//
|
||||
// Block layout:
|
||||
// grid = (B, 1, 1)
|
||||
// block = (max(K, REGIME_DIM) = 8, 1, 1)
|
||||
//
|
||||
// Tiny block — K≤8 and REGIME_DIM=6. Thread layout is per-head for the
|
||||
// softmax-Jacobian reduction, per-regime-dim for the regime-grad output.
|
||||
// We use the same threads for both phases since both dimensions are ≤8.
|
||||
// ─────────────────────────────────────────────────────────────────────
|
||||
extern "C" __global__ void multi_head_policy_backward_gate(
|
||||
const float* __restrict__ grad_gate_probs, // [B × K] from kernel 1
|
||||
const float* __restrict__ gate_probs, // [B × K] forward output
|
||||
const float* __restrict__ regime_h, // [B × REGIME_DIM]
|
||||
const float* __restrict__ W_gate, // [K × REGIME_DIM]
|
||||
const int B,
|
||||
const int K,
|
||||
float* __restrict__ grad_gate_logits,// [B × K] OUT (for diag)
|
||||
float* __restrict__ grad_W_gate_pb, // [B × K × REGIME_DIM]
|
||||
float* __restrict__ grad_b_gate_pb, // [B × K]
|
||||
float* __restrict__ grad_regime_h // [B × REGIME_DIM] (computed; unused upstream)
|
||||
) {
|
||||
const int b = blockIdx.x;
|
||||
if (b >= 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;
|
||||
}
|
||||
}
|
||||
@@ -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<CudaModule>,
|
||||
gate_forward_fn: CudaFunction,
|
||||
_backward_module: Arc<CudaModule>,
|
||||
backward_pi_fn: CudaFunction,
|
||||
backward_gate_fn: CudaFunction,
|
||||
_aux_prior_module: Arc<CudaModule>,
|
||||
aux_prior_fn: CudaFunction,
|
||||
_reduce_axis0_module: Arc<CudaModule>,
|
||||
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<f32>,
|
||||
|
||||
// ── 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<f32>,
|
||||
|
||||
// ── 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<f32>,
|
||||
/// `[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<f32>,
|
||||
/// `[B × K]` — gradient on the pre-softmax gate logits (for diagnostics).
|
||||
pub grad_gate_logits_d: CudaSlice<f32>,
|
||||
/// `[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<f32>,
|
||||
/// `[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<f32>,
|
||||
|
||||
// ── Per-batch grad scratch (for reduce_axis0) ─────────────────────
|
||||
grad_w_heads_pb_d: CudaSlice<f32>,
|
||||
grad_b_heads_pb_d: CudaSlice<f32>,
|
||||
grad_w_gate_pb_d: CudaSlice<f32>,
|
||||
grad_b_gate_pb_d: CudaSlice<f32>,
|
||||
|
||||
// ── 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<f32>,
|
||||
/// Reduced grad on `b_heads`.
|
||||
pub grad_b_heads_d: CudaSlice<f32>,
|
||||
/// Reduced grad on `W_gate`.
|
||||
pub grad_w_gate_d: CudaSlice<f32>,
|
||||
/// Reduced grad on `b_gate`.
|
||||
pub grad_b_gate_d: CudaSlice<f32>,
|
||||
}
|
||||
|
||||
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::<f32>(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::<f32>(cfg.b_size * cfg.k * N_ACTIONS)
|
||||
.context("alloc grad_pi_logits_k_d")?;
|
||||
let grad_gate_probs_d = stream
|
||||
.alloc_zeros::<f32>(cfg.b_size * cfg.k)
|
||||
.context("alloc grad_gate_probs_d")?;
|
||||
let grad_gate_logits_d = stream
|
||||
.alloc_zeros::<f32>(cfg.b_size * cfg.k)
|
||||
.context("alloc grad_gate_logits_d")?;
|
||||
let grad_h_t_scratch_d = stream
|
||||
.alloc_zeros::<f32>(cfg.b_size * HIDDEN_DIM)
|
||||
.context("alloc grad_h_t_scratch_d")?;
|
||||
let grad_regime_h_d = stream
|
||||
.alloc_zeros::<f32>(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::<f32>(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::<f32>(cfg.b_size * cfg.k * N_ACTIONS)
|
||||
.context("alloc grad_b_heads_pb_d")?;
|
||||
let grad_w_gate_pb_d = stream
|
||||
.alloc_zeros::<f32>(cfg.b_size * cfg.k * REGIME_DIM)
|
||||
.context("alloc grad_w_gate_pb_d")?;
|
||||
let grad_b_gate_pb_d = stream
|
||||
.alloc_zeros::<f32>(cfg.b_size * cfg.k)
|
||||
.context("alloc grad_b_gate_pb_d")?;
|
||||
|
||||
// ── Reduced grad buffers ──────────────────────────────────────
|
||||
let grad_w_heads_d = stream
|
||||
.alloc_zeros::<f32>(cfg.k * N_ACTIONS * HIDDEN_DIM)
|
||||
.context("alloc grad_w_heads_d")?;
|
||||
let grad_b_heads_d = stream
|
||||
.alloc_zeros::<f32>(cfg.k * N_ACTIONS)
|
||||
.context("alloc grad_b_heads_d")?;
|
||||
let grad_w_gate_d = stream
|
||||
.alloc_zeros::<f32>(cfg.k * REGIME_DIM)
|
||||
.context("alloc grad_w_gate_d")?;
|
||||
let grad_b_gate_d = stream
|
||||
.alloc_zeros::<f32>(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<f32>,
|
||||
h_t_d: &CudaSlice<f32>,
|
||||
regime_h_d: &CudaSlice<f32>,
|
||||
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<CudaStream>, host: &[f32]) -> Result<CudaSlice<f32>> {
|
||||
}
|
||||
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<f32> {
|
||||
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<CudaStream>,
|
||||
reduce_fn: &CudaFunction,
|
||||
per_batch: &CudaSlice<f32>,
|
||||
b_size: usize,
|
||||
n_tail: usize,
|
||||
out: &mut CudaSlice<f32>,
|
||||
) -> 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(())
|
||||
}
|
||||
|
||||
@@ -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<MappedF32Buffer> {
|
||||
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<f32>, Vec<f32>, Vec<f32>) {
|
||||
use rand::{Rng, SeedableRng};
|
||||
use rand_chacha::ChaCha8Rng;
|
||||
let mut r = ChaCha8Rng::seed_from_u64(seed);
|
||||
let h_t: Vec<f32> = (0..b_size * HIDDEN_DIM)
|
||||
.map(|_| r.gen_range(-0.3..0.3))
|
||||
.collect();
|
||||
let regime: Vec<f32> = (0..b_size * REGIME_DIM)
|
||||
.map(|_| r.gen_range(-0.5..0.5))
|
||||
.collect();
|
||||
let grad_pi_logits: Vec<f32> = (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::<f32>(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<usize> = (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::<f32>(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<usize> = (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<f32> = (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::<f32>(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<usize> = (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<f32> = (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::<f32>(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<usize> = (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<f32> = (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<f32> = (0..b_size * HIDDEN_DIM)
|
||||
.map(|_| r.gen_range(-0.5..0.5))
|
||||
.collect();
|
||||
let regime_host: Vec<f32> = (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::<f32>(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::<f32>(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::<f32>(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(())
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user