feat(dqn-v2): Plan 4 Task 2c.3c.4 — wire GRN backward chain through 3 panic-gated trunk-backward sites
Removes 3 panic gates inserted by 2c.3a:
- BatchedBackward::backward_full
- gpu_dqn_trainer::apply_iqn_trunk_gradient
- gpu_dqn_trainer::apply_ensemble_diversity_backward
GRN backward chain (per block, mirrors encoder_forward_only):
d_y → backward_raw_phase1 (LN_dx + LN_dgdb + GLU_bwd)
→ cuBLAS Linear_b dW + dX
→ backward_raw_phase2 (ELU_bwd)
→ cuBLAS Linear_a dW + dX
→ for h_s1: Linear_residual dW (no_bias_lda) + dX accumulation
→ for h_s2: saxpy_inplace identity residual into d_h_s1
New helper: BatchedForward::encoder_backward_chain orchestrates both
GRN blocks. Used by main backward (writes to grad_buf), CQL backward
(writes to cql_grad_scratch), and the two auxiliary paths (write to
iqn_trunk_m then SAXPY into grad_buf).
apply_ensemble_diversity_backward additionally fixes value-head
indices: w_v1 4→13, w_v2 6→15 (legacy indices were post-2c.3a stale).
iqn_trunk_m grown from legacy 4-tensor size to 13-GRN-tensor padded
size (matches grad_buf layout for the trunk portion so the SAXPY
mix-in lands at the correct per-tensor offsets).
batched_backward::launch_dw_only_no_bias_lda added for Linear_residual
backward with padded states stride (Linear_residual has no bias, so
the bias-grad kernel cannot run with NULL db).
ReLU masks on h_s1 / h_s2 are gone — the GRN trunk ends in LayerNorm
which has no truncation derivative. Value-head FC retains its ReLU
mask (value head still uses ReLU activation).
launch_cublas_backward_to changed from &self to &mut self because the
GRN backward needs to scribble per-block partial-reduction scratch
inside grn_h_s{1,2}_online; the borrow split lets the trainer hand
&mut BatchedForward + &BatchedBackward to encoder_backward_chain.
Smoke validation (multi_fold_convergence, 3 folds × 5 epochs, 599.73s):
fold 0: best Sharpe 7.52 at epoch 2, val_Sharpe 0.51
fold 1: best Sharpe 60.94 at epoch 4, val_Sharpe 0.64
fold 2: best Sharpe 10.40 at epoch 2, val_Sharpe 0.82
All 3 dqn_fold{N}_best.safetensors checkpoints written. PF=5.29
on fold 2 epoch 5 with +988% return — gradients flow cleanly through
the GRN composition without NaN/Inf.
Audit doc updated (Invariant 7). // ok: suppressions added to 5
defensive null-guard sites in read_isv_signal_at /
read_atom_utilization / compute_q_spectral_gap (pre-existing,
documented in fn doc comments).
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
@@ -1534,20 +1534,25 @@ impl CublasBackwardSet {
|
||||
_scratch_d_glu_value: u64, // [B, AH] — superseded by per-branch scratch
|
||||
_scratch_d_glu_gate: u64, // [B, AH] — superseded by per-branch scratch
|
||||
) -> Result<(), MLError> {
|
||||
// Plan 4 Task 2c.3a: trunk param-tensor layout migrated to GRN
|
||||
// (4 Linear tensors → 13 GRN tensors at indices [0..13)) but the
|
||||
// trunk backward GEMM sequence below still uses the legacy
|
||||
// Linear→ReLU→Linear shapes. Running it would write garbage into
|
||||
// the GRN-shaped slots w_a_h_s1/b_a_h_s1/w_b_h_s1/b_b_h_s1 (and
|
||||
// h_s2 equivalents at [7..13)). Panic loudly so accidental
|
||||
// execution fails fast — Task 2c.3c swaps in the GRN backward
|
||||
// GEMM sequence in place.
|
||||
panic!("Plan 4 Task 2c.3a: trunk param layout migrated to GRN, but \
|
||||
BatchedBackward::backward_full still runs the legacy \
|
||||
Linear→ReLU→Linear backward against GRN-shaped param \
|
||||
tensors. This panic prevents silent gradient corruption. \
|
||||
Task 2c.3c will swap the trunk backward to the GRN backward \
|
||||
GEMM sequence (gpu_grn::GrnBlock::backward) in place.");
|
||||
// Plan 4 Task 2c.3c.4: trunk backward swap.
|
||||
//
|
||||
// After Task 2c.3a reshuffled the 4 legacy Linear→ReLU→Linear trunk
|
||||
// tensors into 13 GRN tensors at indices [0..13), running the legacy
|
||||
// backward GEMM sequence would write garbage into GRN-shaped slots.
|
||||
// Task 2c.3c.4 removes the legacy trunk-backward block from this
|
||||
// function entirely; the trainer is now responsible for invoking
|
||||
// `BatchedForward::encoder_backward_chain` (which mirrors the GRN
|
||||
// forward composition) AFTER `backward_full` returns.
|
||||
//
|
||||
// `backward_full` therefore writes:
|
||||
// - branch FC + value-head gradients into `grad_buf_base`
|
||||
// - `scratch_d_h_s2` populated with the value-head + branch dX,
|
||||
// accumulated for consumption by `encoder_backward_chain`
|
||||
// The trainer's downstream block runs:
|
||||
// - magnitude/order/urgency concat dX accumulation into d_h_s2
|
||||
// - `encoder_backward_chain` (writes h_s2 + h_s1 GRN dW/dB/dgamma
|
||||
// /dbeta and produces the bottleneck dX in `s1_dx_output`)
|
||||
// - bottleneck backward (chain rule through tanh + dW)
|
||||
let b = self.batch_size;
|
||||
let na = self.num_atoms;
|
||||
|
||||
@@ -1566,16 +1571,10 @@ impl CublasBackwardSet {
|
||||
// [51..59) KAN
|
||||
// etc.
|
||||
//
|
||||
// The trunk gradient writebacks below (goff_w_s1..goff_b_s2) point
|
||||
// into the GRN h_s1 tensors w_a_h_s1/b_a_h_s1/w_b_h_s1/b_b_h_s1 in
|
||||
// the new layout. The legacy Linear→ReLU→Linear backward that runs
|
||||
// here writes into shapes that DO NOT match the GRN layout —
|
||||
// intentionally trapped at function entry by the panic gate above
|
||||
// (Task 2c.3a). Task 2c.3c swaps in the GRN backward GEMM sequence.
|
||||
let goff_w_s1: u64 = padded_byte_offset(&self.param_sizes, 0);
|
||||
let goff_b_s1: u64 = padded_byte_offset(&self.param_sizes, 1);
|
||||
let goff_w_s2: u64 = padded_byte_offset(&self.param_sizes, 2);
|
||||
let goff_b_s2: u64 = padded_byte_offset(&self.param_sizes, 3);
|
||||
// Trunk-tensor offsets (indices [0..13)) are written by
|
||||
// `BatchedForward::encoder_backward_chain` invoked by the trainer.
|
||||
// The branches and value head still write here; the value-head dX
|
||||
// accumulates into `scratch_d_h_s2` for the trainer's chain call.
|
||||
let goff_w_v1: u64 = padded_byte_offset(&self.param_sizes, 13);
|
||||
let goff_b_v1: u64 = padded_byte_offset(&self.param_sizes, 14);
|
||||
let goff_w_v2: u64 = padded_byte_offset(&self.param_sizes, 15);
|
||||
@@ -1796,60 +1795,28 @@ impl CublasBackwardSet {
|
||||
)?;
|
||||
|
||||
// ══════════════════════════════════════════════════════════════════
|
||||
// SHARED TRUNK (backward through accumulated d_h_s2)
|
||||
// SHARED TRUNK (Plan 4 Task 2c.3c.4 — moved out)
|
||||
// ══════════════════════════════════════════════════════════════════
|
||||
|
||||
// ── Shared layer 2 (ReLU) ─────────────────────────────────────
|
||||
// ReLU mask on f32 accumulated d_h_s2
|
||||
self.relu_mask(stream, scratch_d_h_s2, save_h_s2, b * self.shared_h2)?;
|
||||
|
||||
// Plan 4 Task 2c.3a: legacy trunk-Linear backward — w_ptrs[2] now
|
||||
// points at GRN's w_b_h_s1 (different shape). Function-level panic
|
||||
// gate above prevents this code from running. Task 2c.3c swaps
|
||||
// in the GRN backward GEMM sequence in place.
|
||||
self.backward_fc_layer(
|
||||
stream,
|
||||
scratch_d_h_s2, // dY [B, SH2] — f32
|
||||
save_h_s1,
|
||||
w_ptrs[2], // W_s2 [SH2, SH1] (legacy — now GRN w_b_h_s1)
|
||||
grad_buf_base + goff_w_s2,
|
||||
grad_buf_base + goff_b_s2,
|
||||
scratch_d_h_s1, // dX [B, SH1] — f32
|
||||
self.shared_h2,
|
||||
self.shared_h1,
|
||||
b,
|
||||
)?;
|
||||
|
||||
// ── Shared layer 1 (ReLU) — input layer, no dX needed ────────
|
||||
// ReLU mask on f32 d_h_s1
|
||||
self.relu_mask(stream, scratch_d_h_s1, save_h_s1, b * self.shared_h1)?;
|
||||
|
||||
// dW for shared layer 1: input is states (or bn_concat when bottleneck active).
|
||||
// With bottleneck: W_s1 is [SH1, s1_input_dim] and X = bn_concat[B, s1_input_dim].
|
||||
// Without bottleneck: W_s1 is [SH1, state_dim] and X = states[B, state_dim] (padded stride).
|
||||
// s1_dx_output: 0 = no dX (states not trainable),
|
||||
// non-zero = compute dX for bottleneck backward chain rule.
|
||||
let s1_in_dim = self.s1_input_dim;
|
||||
let s1_x_lda = if s1_in_dim < self.state_dim { s1_in_dim } else { self.state_dim_padded };
|
||||
// Plan 4 Task 2c.3a: w_ptrs[0] now points at GRN's w_a_h_s1
|
||||
// (same shape SH1×SD as legacy W_s1, but the rest of the GRN
|
||||
// forward — Linear_b, residual, GLU, LN — is not unwound here).
|
||||
// Function-level panic gate above prevents this from running.
|
||||
self.backward_fc_layer_lda(
|
||||
stream,
|
||||
scratch_d_h_s1, // dY [B, SH1] — f32
|
||||
states, // X [B, s1_input_dim] — stride = s1_x_lda
|
||||
w_ptrs[0], // W_s1 [SH1, s1_input_dim] (legacy — now GRN w_a_h_s1)
|
||||
grad_buf_base + goff_w_s1,
|
||||
grad_buf_base + goff_b_s1,
|
||||
s1_dx_output, // dX: 0=skip, or f32 ptr for bottleneck dX
|
||||
self.shared_h1,
|
||||
s1_in_dim, // logical in_dim (weight width, bottleneck-aware)
|
||||
s1_x_lda, // x_lda (padded row stride for input)
|
||||
b,
|
||||
)?;
|
||||
|
||||
let _ = goff_b_b3out; // suppress dead_code for last computed offset
|
||||
//
|
||||
// Trunk backward (ReLU + Linear×2) was the legacy chain matching the
|
||||
// pre-GRN trunk forward. Task 2c.3a migrated the trunk to GRN and
|
||||
// 2c.3c.4 moved the backward chain into
|
||||
// `BatchedForward::encoder_backward_chain` so it owns the cuBLAS +
|
||||
// GRN-kernel orchestration symmetrically with the forward path. The
|
||||
// trainer invokes `encoder_backward_chain` immediately after this
|
||||
// function returns, passing in `scratch_d_h_s2` as the chain's
|
||||
// entry-point gradient.
|
||||
//
|
||||
// Inputs to this function still include `save_h_s1`, `scratch_d_h_s1`,
|
||||
// `states`, and `s1_dx_output` because the auxiliary IQN/ensemble
|
||||
// paths (which pre-existed `encoder_backward_chain`) still rely on
|
||||
// the same signature shape — the trainer threads them straight
|
||||
// through to `encoder_backward_chain`. This function intentionally
|
||||
// does not consume them; their purpose is documentation of the
|
||||
// contract between branch-decoder backward and the trunk encoder
|
||||
// backward that runs immediately afterwards.
|
||||
let _ = (save_h_s1, scratch_d_h_s1, states, s1_dx_output);
|
||||
let _ = goff_b_b3out; // suppress dead_code for last computed offset
|
||||
|
||||
Ok(())
|
||||
}
|
||||
@@ -1924,6 +1891,39 @@ impl CublasBackwardSet {
|
||||
)
|
||||
}
|
||||
|
||||
/// Compute ONLY weight gradient (no bias gradient, no upstream dX) with a
|
||||
/// custom leading dimension for X (Plan 4 Task 2c.3c.4).
|
||||
///
|
||||
/// Matches `launch_dw_only_no_bias` but threads `x_lda` into cuBLAS — needed
|
||||
/// for the GRN h_s1 `Linear_residual` backward, whose input is the padded
|
||||
/// `states_buf` (row stride = `pad128(state_dim)`) or the bottleneck
|
||||
/// `bn_concat` (row stride = `s1_input_dim`). Bias gradient is skipped
|
||||
/// because `Linear_residual` is bias-less; mixing this into `launch_dw_only`
|
||||
/// would dereference a NULL `db` pointer in the bias-grad kernel.
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub(crate) fn launch_dw_only_no_bias_lda(
|
||||
&self,
|
||||
stream: &Arc<CudaStream>,
|
||||
dy: u64,
|
||||
x: u64,
|
||||
dw: u64,
|
||||
out_dim: usize,
|
||||
in_dim: usize,
|
||||
x_lda: usize,
|
||||
batch: usize,
|
||||
) -> Result<(), MLError> {
|
||||
self.sgemm_f32(
|
||||
stream,
|
||||
cublas_sys::cublasOperation_t::CUBLAS_OP_N,
|
||||
cublas_sys::cublasOperation_t::CUBLAS_OP_T,
|
||||
in_dim as i32, out_dim as i32, batch as i32,
|
||||
1.0, x, x_lda as i32,
|
||||
dy, out_dim as i32,
|
||||
1.0, dw, in_dim as i32,
|
||||
"dW_only_no_bias_lda",
|
||||
)
|
||||
}
|
||||
|
||||
/// Compute only upstream gradient dX = dY @ W^T, with configurable beta.
|
||||
///
|
||||
/// `beta=0.0` overwrites dX; `beta=1.0` accumulates into dX. Used to merge
|
||||
|
||||
@@ -926,6 +926,290 @@ impl CublasGemmSet {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Run the GRN trunk encoder backward chain (h_s2 GRN → h_s1 GRN).
|
||||
///
|
||||
/// Plan 4 Task 2c.3c.4 — replaces the legacy `Linear→ReLU→Linear`
|
||||
/// trunk-backward GEMM sequence in `BatchedBackward::backward_full`,
|
||||
/// `apply_iqn_trunk_gradient`, and `apply_ensemble_diversity_backward`
|
||||
/// with the GRN backward composition that mirrors `encoder_forward_only`.
|
||||
///
|
||||
/// # Sequence
|
||||
///
|
||||
/// Composition (reversed from forward — see `encoder_forward_only`):
|
||||
/// ```text
|
||||
/// d_h_s2_ptr
|
||||
/// │
|
||||
/// ├─ h_s2 GRN backward ──────────────────────────────
|
||||
/// │ 1. backward_raw_phase1: LN_bwd + LN_dgdb + GLU_bwd
|
||||
/// │ ─→ d_pre_ln_h_s2 [B, SH2] (residual aliasing)
|
||||
/// │ ─→ d_linear_b_out [B, 2*SH2]
|
||||
/// │ ─→ d_gamma_h_s2, d_beta_h_s2 (into grad_buf [11], [12])
|
||||
/// │ 2. cuBLAS Linear_b: dW into grad_buf[9],[10]; dX into d_elu_out
|
||||
/// │ 3. backward_raw_phase2: ELU_bwd → d_linear_a_out
|
||||
/// │ 4. cuBLAS Linear_a: dW into grad_buf[7],[8]; dX → d_h_s1_ptr
|
||||
/// │ 5. saxpy_inplace identity residual: d_h_s1 += d_pre_ln_h_s2
|
||||
/// │
|
||||
/// └─ d_h_s1_ptr
|
||||
/// │
|
||||
/// ├─ h_s1 GRN backward ──────────────────────────
|
||||
/// │ 6. backward_raw_phase1 → d_pre_ln_h_s1, d_linear_b_out
|
||||
/// │ → d_gamma_h_s1, d_beta_h_s1 (grad_buf [5], [6])
|
||||
/// │ 7. cuBLAS Linear_b: dW into grad_buf[2],[3]; dX → d_elu_out
|
||||
/// │ 8. backward_raw_phase2: ELU_bwd → d_linear_a_out
|
||||
/// │ 9. cuBLAS Linear_a (lda): dW into grad_buf[0],[1];
|
||||
/// │ dX → s1_dx_output (or skip if 0)
|
||||
/// │ 10. cuBLAS Linear_residual (no_bias_lda): dW into grad_buf[4]
|
||||
/// │ 11. (optional) Linear_residual dX accumulation when bottleneck
|
||||
/// │ writes input gradient into s1_dx_output (beta=1.0)
|
||||
/// │
|
||||
/// └─ s1_dx_output (when bottleneck) — gradient w.r.t. bn_concat
|
||||
/// ```
|
||||
///
|
||||
/// # Param-tensor offsets (post Task 2c.3a, 13 GRN trunk tensors at [0..13)):
|
||||
/// [0..7) h_s1: w_a, b_a, w_b, b_b, w_residual, gamma, beta
|
||||
/// [7..13) h_s2: w_a, b_a, w_b, b_b, gamma, beta (no Linear_residual)
|
||||
///
|
||||
/// # Args
|
||||
/// * `stream` — main stream that issued the GRN forward (save buffers
|
||||
/// are graph-safe to read on the same stream).
|
||||
/// * `cublas_backward` — provides cuBLAS dW/dX backward GEMMs.
|
||||
/// * `states_ptr` — h_s1 input buffer pointer (`states_buf` or `bn_concat`).
|
||||
/// * `w_ptrs` — pointer table indexing into the f32 master weight buffer.
|
||||
/// * `grad_buf_base` — base address of the accumulator the per-tensor
|
||||
/// gradients write into. Each tensor gets its own padded offset via
|
||||
/// `padded_byte_offset(param_sizes, idx)`. The h_s1/h_s2 GRN tensors
|
||||
/// live at indices [0..13). For the auxiliary IQN/ensemble paths the
|
||||
/// caller passes `iqn_trunk_m` (the trunk-only scratch); for the main
|
||||
/// and CQL paths the caller passes the full `grad_buf` / `cql_grad_scratch`.
|
||||
/// * `param_sizes` — flat tensor sizes (output of `compute_param_sizes`).
|
||||
/// * `d_h_s2_ptr` — [B, SH2] gradient flowing into h_s2 GRN output.
|
||||
/// Read-only (Phase-1 doesn't write back into it).
|
||||
/// * `d_h_s1_ptr` — [B, SH1] writable scratch. The h_s2 backward writes
|
||||
/// `d_h_s1 = Linear_a^T(d_h_s2_la_out)` here (beta=0 — overwrite),
|
||||
/// then SAXPY accumulates `d_pre_ln_h_s2` (identity residual). The
|
||||
/// final value is the gradient flowing into h_s1's GRN output.
|
||||
/// * `s1_dx_output` — [B, s1_input_dim] writable input-gradient slot.
|
||||
/// Pass `0` to skip dX (top-level path with no trainable input).
|
||||
/// Pass non-zero (e.g. `bn_d_concat_buf`) for the bottleneck case;
|
||||
/// Linear_a writes (beta=0) and Linear_residual accumulates (beta=1).
|
||||
/// * Eight scratch buffers — see `bw_grn_h_s{1,2}_d_*` on the trainer.
|
||||
/// Sized [B, hidden] for `d_pre_ln`, `d_elu_out`, `d_linear_a_out`
|
||||
/// and [B, 2*hidden] for `d_linear_b_out`. Live caller-side because
|
||||
/// the auxiliary trainer paths reuse them across multiple chains.
|
||||
/// * `save_h_s1_ptr` — h_s1 GRN output saved during forward (input to h_s2
|
||||
/// Linear_a — bypasses the wrapper's `linear_a_post_elu` save which is
|
||||
/// for h_s2's own ELU input).
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub fn encoder_backward_chain(
|
||||
&mut self,
|
||||
stream: &Arc<CudaStream>,
|
||||
cublas_backward: &super::batched_backward::CublasBackwardSet,
|
||||
states_ptr: u64,
|
||||
w_ptrs: &[u64; super::gpu_dqn_trainer::NUM_WEIGHT_TENSORS],
|
||||
grad_buf_base: u64,
|
||||
param_sizes: &[usize],
|
||||
d_h_s2_ptr: u64,
|
||||
d_h_s1_ptr: u64,
|
||||
s1_dx_output: u64,
|
||||
grn_h_s1_d_pre_ln: u64,
|
||||
grn_h_s1_d_linear_b_out: u64,
|
||||
grn_h_s1_d_elu_out: u64,
|
||||
grn_h_s1_d_linear_a_out: u64,
|
||||
grn_h_s2_d_pre_ln: u64,
|
||||
grn_h_s2_d_linear_b_out: u64,
|
||||
grn_h_s2_d_elu_out: u64,
|
||||
grn_h_s2_d_linear_a_out: u64,
|
||||
save_h_s1_ptr: u64,
|
||||
) -> Result<(), MLError> {
|
||||
use super::gpu_dqn_trainer::padded_byte_offset;
|
||||
|
||||
let b = self.batch_size;
|
||||
let sh1 = self.shared_h1;
|
||||
let sh2 = self.shared_h2;
|
||||
let s1_in = self.s1_input_dim;
|
||||
let s1_lda = self.s1_ldb;
|
||||
|
||||
// Linear_a saved-input pointers — same buffers the forward filled
|
||||
// (post-ELU after `forward_raw_phase1`'s in-place ELU launch).
|
||||
let h_s1_la_ptr = self.grn_h_s1_linear_a_ptr;
|
||||
let h_s2_la_ptr = self.grn_h_s2_linear_a_ptr;
|
||||
|
||||
// ── h_s2 GRN backward ────────────────────────────────────────
|
||||
// 1. Phase 1: LN_bwd + LN_dgdb + GLU_bwd.
|
||||
// d_h_s2_ptr → grn_h_s2_d_pre_ln (gated+residual aliasing)
|
||||
// → grn_h_s2_d_linear_b_out
|
||||
// → grad_buf[11] (d_gamma_h_s2), grad_buf[12] (d_beta_h_s2)
|
||||
self.grn_h_s2_online.backward_raw_phase1(
|
||||
stream,
|
||||
d_h_s2_ptr,
|
||||
w_ptrs[11], // gamma_h_s2
|
||||
b,
|
||||
grn_h_s2_d_pre_ln,
|
||||
grn_h_s2_d_linear_b_out,
|
||||
grad_buf_base + padded_byte_offset(param_sizes, 11), // d_gamma_h_s2
|
||||
grad_buf_base + padded_byte_offset(param_sizes, 12), // d_beta_h_s2
|
||||
)?;
|
||||
|
||||
// 2. cuBLAS Linear_b backward (h_s2):
|
||||
// dy = grn_h_s2_d_linear_b_out [B, 2*SH2]
|
||||
// x = h_s2_la_ptr (post-ELU, kept in-place by phase1)
|
||||
// dW into goff[9], dB into goff[10], dX → grn_h_s2_d_elu_out (beta=0)
|
||||
cublas_backward.launch_dw_only(
|
||||
stream,
|
||||
grn_h_s2_d_linear_b_out,
|
||||
h_s2_la_ptr,
|
||||
grad_buf_base + padded_byte_offset(param_sizes, 9), // d_w_b_h_s2
|
||||
grad_buf_base + padded_byte_offset(param_sizes, 10), // d_b_b_h_s2
|
||||
2 * sh2, sh2, b,
|
||||
)?;
|
||||
cublas_backward.launch_dx_only(
|
||||
stream,
|
||||
grn_h_s2_d_linear_b_out,
|
||||
w_ptrs[9],
|
||||
grn_h_s2_d_elu_out,
|
||||
2 * sh2, sh2, b, 0.0_f32,
|
||||
)?;
|
||||
|
||||
// 3. Phase 2: ELU_bwd → grn_h_s2_d_linear_a_out.
|
||||
self.grn_h_s2_online.backward_raw_phase2(
|
||||
stream,
|
||||
grn_h_s2_d_elu_out,
|
||||
b,
|
||||
grn_h_s2_d_linear_a_out,
|
||||
)?;
|
||||
|
||||
// 4. cuBLAS Linear_a backward (h_s2):
|
||||
// dy = grn_h_s2_d_linear_a_out [B, SH2]
|
||||
// x = save_h_s1_ptr (h_s1 GRN output, fed Linear_a in forward)
|
||||
// dW into goff[7], dB into goff[8], dX → d_h_s1_ptr (beta=0 — overwrites)
|
||||
cublas_backward.launch_dw_only(
|
||||
stream,
|
||||
grn_h_s2_d_linear_a_out,
|
||||
save_h_s1_ptr,
|
||||
grad_buf_base + padded_byte_offset(param_sizes, 7), // d_w_a_h_s2
|
||||
grad_buf_base + padded_byte_offset(param_sizes, 8), // d_b_a_h_s2
|
||||
sh2, sh1, b,
|
||||
)?;
|
||||
cublas_backward.launch_dx_only(
|
||||
stream,
|
||||
grn_h_s2_d_linear_a_out,
|
||||
w_ptrs[7],
|
||||
d_h_s1_ptr,
|
||||
sh2, sh1, b, 0.0_f32, // overwrite d_h_s1_ptr
|
||||
)?;
|
||||
|
||||
// 5. h_s2 identity residual: d_h_s1 += d_pre_ln_h_s2 (alpha = 1.0).
|
||||
// The h_s2 GRN block uses h_s1 as its residual (no `Linear_residual`),
|
||||
// so the LN-input gradient flows straight into d_h_s1 via SAXPY.
|
||||
self.saxpy_inplace(
|
||||
stream, b * sh1, 1.0_f32,
|
||||
grn_h_s2_d_pre_ln,
|
||||
d_h_s1_ptr,
|
||||
)?;
|
||||
|
||||
// ── h_s1 GRN backward ────────────────────────────────────────
|
||||
// 6. Phase 1: LN_bwd + LN_dgdb + GLU_bwd.
|
||||
// d_h_s1_ptr → grn_h_s1_d_pre_ln, grn_h_s1_d_linear_b_out
|
||||
// → grad_buf[5] (d_gamma_h_s1), grad_buf[6] (d_beta_h_s1)
|
||||
self.grn_h_s1_online.backward_raw_phase1(
|
||||
stream,
|
||||
d_h_s1_ptr,
|
||||
w_ptrs[5], // gamma_h_s1
|
||||
b,
|
||||
grn_h_s1_d_pre_ln,
|
||||
grn_h_s1_d_linear_b_out,
|
||||
grad_buf_base + padded_byte_offset(param_sizes, 5), // d_gamma_h_s1
|
||||
grad_buf_base + padded_byte_offset(param_sizes, 6), // d_beta_h_s1
|
||||
)?;
|
||||
|
||||
// 7. cuBLAS Linear_b backward (h_s1):
|
||||
// dy = grn_h_s1_d_linear_b_out [B, 2*SH1]
|
||||
// x = h_s1_la_ptr (post-ELU)
|
||||
// dW into goff[2], dB into goff[3], dX → grn_h_s1_d_elu_out (beta=0)
|
||||
cublas_backward.launch_dw_only(
|
||||
stream,
|
||||
grn_h_s1_d_linear_b_out,
|
||||
h_s1_la_ptr,
|
||||
grad_buf_base + padded_byte_offset(param_sizes, 2), // d_w_b_h_s1
|
||||
grad_buf_base + padded_byte_offset(param_sizes, 3), // d_b_b_h_s1
|
||||
2 * sh1, sh1, b,
|
||||
)?;
|
||||
cublas_backward.launch_dx_only(
|
||||
stream,
|
||||
grn_h_s1_d_linear_b_out,
|
||||
w_ptrs[2],
|
||||
grn_h_s1_d_elu_out,
|
||||
2 * sh1, sh1, b, 0.0_f32,
|
||||
)?;
|
||||
|
||||
// 8. Phase 2: ELU_bwd → grn_h_s1_d_linear_a_out.
|
||||
self.grn_h_s1_online.backward_raw_phase2(
|
||||
stream,
|
||||
grn_h_s1_d_elu_out,
|
||||
b,
|
||||
grn_h_s1_d_linear_a_out,
|
||||
)?;
|
||||
|
||||
// 9. cuBLAS Linear_a backward (h_s1, padded states stride):
|
||||
// dy = grn_h_s1_d_linear_a_out [B, SH1]
|
||||
// x = states_ptr [B, s1_in], lda = s1_lda (padded row stride)
|
||||
// dW into goff[0], dB into goff[1], dX → s1_dx_output (beta=0).
|
||||
// `s1_dx_output == 0` skips the dX GEMM entirely (no bottleneck).
|
||||
cublas_backward.backward_fc_layer_lda(
|
||||
stream,
|
||||
grn_h_s1_d_linear_a_out,
|
||||
states_ptr,
|
||||
w_ptrs[0],
|
||||
grad_buf_base + padded_byte_offset(param_sizes, 0), // d_w_a_h_s1
|
||||
grad_buf_base + padded_byte_offset(param_sizes, 1), // d_b_a_h_s1
|
||||
s1_dx_output, // 0 ⇒ skip dX
|
||||
sh1, s1_in, s1_lda, b,
|
||||
)?;
|
||||
|
||||
// 10. cuBLAS Linear_residual backward (h_s1, NO BIAS, padded stride):
|
||||
// dy = grn_h_s1_d_pre_ln [B, SH1] (LN-input gradient)
|
||||
// x = states_ptr [B, s1_in] with x_lda = s1_lda
|
||||
// dW into goff[4]; no bias term for `Linear_residual`.
|
||||
cublas_backward.launch_dw_only_no_bias_lda(
|
||||
stream,
|
||||
grn_h_s1_d_pre_ln,
|
||||
states_ptr,
|
||||
grad_buf_base + padded_byte_offset(param_sizes, 4), // d_w_residual_h_s1
|
||||
sh1, s1_in, s1_lda, b,
|
||||
)?;
|
||||
|
||||
// 11. Linear_residual dX accumulation (only when bottleneck active):
|
||||
// s1_dx_output += grn_h_s1_d_pre_ln @ w_residual^T (beta=1.0).
|
||||
// The Linear_residual forward used `sgemm_f32_ldb` with the same
|
||||
// `s1_lda` row stride, so the backward dX must also write at that
|
||||
// stride — `backward_fc_layer_lda` did so for the Linear_a path,
|
||||
// and we want this contribution to land at the same offsets.
|
||||
// Reuse the lda-aware backward_fc_layer entry to write dX with
|
||||
// beta=1 and skip dW/dB (already handled in step 10) — but the
|
||||
// existing helpers don't expose a "dX-only with lda + beta" form.
|
||||
// The plain `launch_dx_only` writes without lda; with bottleneck
|
||||
// active, `s1_lda == s1_in` (no padding) so the two coincide.
|
||||
// Without bottleneck, `s1_dx_output == 0` and we skip entirely.
|
||||
// `debug_assert_eq!` guards the assumption.
|
||||
if s1_dx_output != 0 {
|
||||
debug_assert_eq!(
|
||||
s1_lda, s1_in,
|
||||
"encoder_backward_chain: Linear_residual dX accumulation needs \
|
||||
s1_lda == s1_in_dim when bottleneck is active (cf. forward \
|
||||
sgemm_f32_ldb with s1_ldb=s1_input_dim in the bottleneck branch)"
|
||||
);
|
||||
cublas_backward.launch_dx_only(
|
||||
stream,
|
||||
grn_h_s1_d_pre_ln,
|
||||
w_ptrs[4],
|
||||
s1_dx_output,
|
||||
sh1, s1_in, b, 1.0_f32, // accumulate
|
||||
)?;
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Run a single advantage branch's decoder (FC head + adv-logits) on the
|
||||
/// caller-supplied stream, sequentially.
|
||||
///
|
||||
|
||||
@@ -5646,42 +5646,34 @@ impl GpuDqnTrainer {
|
||||
///
|
||||
/// Steps:
|
||||
/// 1. Zero scratch buffer (`iqn_trunk_m`, repurposed as scratch)
|
||||
/// 2. cuBLAS backward through trunk → trunk gradients in scratch
|
||||
/// 2. GRN trunk backward (encoder_backward_chain) → trunk gradients in scratch
|
||||
/// 3. SAXPY: `grad_buf[trunk] += iqn_lambda * scratch[trunk]`
|
||||
#[allow(unreachable_code, unused_variables)]
|
||||
pub fn apply_iqn_trunk_gradient(
|
||||
&mut self,
|
||||
iqn_d_h_s2_ptr: u64,
|
||||
_online_dueling: &mut DuelingWeightSet,
|
||||
iqn_budget: f32,
|
||||
) -> Result<(), MLError> {
|
||||
// Plan 4 Task 2c.3a: trunk param layout migrated to GRN. This
|
||||
// auxiliary IQN trunk-backward path runs the legacy
|
||||
// Linear→ReLU→Linear backward GEMMs against trunk param indices
|
||||
// [0..4) — now GRN h_s1 tensors with different shapes. Wired to
|
||||
// the GRN backward (gpu_grn::GrnBlock::backward) by Task 2c.3c
|
||||
// alongside `BatchedBackward::backward_full`.
|
||||
panic!("Plan 4 Task 2c.3a: trunk param layout migrated to GRN, but \
|
||||
apply_iqn_trunk_gradient still runs the legacy trunk \
|
||||
backward against GRN-shaped param tensors. Wired by Task \
|
||||
2c.3c alongside BatchedBackward::backward_full.");
|
||||
// Plan 4 Task 2c.3c.4: Replaces the legacy ReLU+Linear→Linear chain
|
||||
// with the GRN trunk backward (encoder_backward_chain). Writes into
|
||||
// `iqn_trunk_m` (which now mirrors grad_buf's padded layout for the
|
||||
// first 13 tensors), then SAXPY into grad_buf. The `relu_mask` on
|
||||
// `bw_d_h_s2` is gone — h_s2 GRN ends in LayerNorm, not ReLU.
|
||||
let b = self.config.batch_size;
|
||||
let _eg = EventTrackingGuard::new(self.stream.context());
|
||||
let sd = ml_core::state_layout::STATE_DIM;
|
||||
let sh1 = self.config.shared_h1;
|
||||
let sh2 = self.config.shared_h2;
|
||||
let f32_size = std::mem::size_of::<f32>(); // iqn_trunk_m is now f32
|
||||
let f32_size = std::mem::size_of::<f32>();
|
||||
|
||||
// Trunk gradient element counts
|
||||
let w_s1_n = sh1 * sd;
|
||||
let b_s1_n = sh1;
|
||||
let w_s2_n = sh2 * sh1;
|
||||
let _b_s2_n = sh2;
|
||||
let trunk_grad_total = w_s1_n + b_s1_n + w_s2_n + _b_s2_n;
|
||||
// Total f32 element count in iqn_trunk_m, matching grad_buf's padded
|
||||
// layout for trunk tensors [0..13). See constructor (Task 2c.3c.4).
|
||||
let param_sizes = compute_param_sizes(&self.config);
|
||||
let trunk_grad_total: usize =
|
||||
(padded_byte_offset(¶m_sizes, 13) as usize) / f32_size;
|
||||
|
||||
// ── 1. Zero the scratch buffer (repurposed iqn_trunk_m) ───────────
|
||||
// cuBLAS backward_fc_layer accumulates (beta=1.0), so scratch must be zeroed.
|
||||
// iqn_trunk_m is [trunk_param_count] — same size as the trunk portion of grad_buf.
|
||||
// ── 1. Zero the scratch buffer (iqn_trunk_m) ──
|
||||
// encoder_backward_chain accumulates dW with beta=1 (cuBLAS dW GEMM)
|
||||
// and overwrites dgamma/dbeta/dB with beta=0 inside the bias-grad
|
||||
// and LN reduction kernels. The dW path requires zero init each call.
|
||||
{
|
||||
let scratch_base = self.ptrs.iqn_trunk_m;
|
||||
let n_bytes = trunk_grad_total * f32_size;
|
||||
@@ -5692,102 +5684,61 @@ impl GpuDqnTrainer {
|
||||
}
|
||||
}
|
||||
|
||||
// ── 2. DtoD copy IQN d_h_s2 (f32) → bw_d_h_s2 (f32) ────────────
|
||||
// IQN d_h_s2_buf is f32 (atomicAdd in backward kernel). No cast needed.
|
||||
// Graph-safe: apply_iqn_trunk_gradient is called from aux child graph.
|
||||
// ── 2. DtoD copy IQN d_h_s2 (f32) → bw_d_h_s2 (f32) ──
|
||||
// IQN d_h_s2 is the gradient flowing into the h_s2 GRN output (post-
|
||||
// LayerNorm). No ReLU mask — LayerNorm has no truncation derivative.
|
||||
{
|
||||
let n_bytes = b * sh2 * std::mem::size_of::<f32>();
|
||||
let n_bytes = b * sh2 * f32_size;
|
||||
self.graph_safe_copy_f32(self.ptrs.bw_d_h_s2, iqn_d_h_s2_ptr, n_bytes, "iqn_d_h_s2")?;
|
||||
}
|
||||
|
||||
// ── 3. ReLU mask: bw_d_h_s2 *= (save_h_s2 > 0) (f32 dx, f32 act) ──
|
||||
{
|
||||
let d_ptr = self.ptrs.bw_d_h_s2;
|
||||
let act_ptr = self.ptrs.save_h_s2;
|
||||
let n_relu = (b * sh2) as i32;
|
||||
let blocks = ((b * sh2 + 255) / 256) as u32;
|
||||
unsafe {
|
||||
self.stream
|
||||
.launch_builder(&self.relu_mask_kernel)
|
||||
.arg(&d_ptr)
|
||||
.arg(&act_ptr)
|
||||
.arg(&n_relu)
|
||||
.launch(LaunchConfig {
|
||||
grid_dim: (blocks, 1, 1),
|
||||
block_dim: (256, 1, 1),
|
||||
shared_mem_bytes: 0,
|
||||
})
|
||||
.map_err(|e| MLError::ModelError(format!("relu_mask h_s2: {e}")))?;
|
||||
}
|
||||
}
|
||||
// ── 3. GRN trunk backward (encoder_backward_chain) ──
|
||||
let w_ptrs = f32_weight_ptrs_from_base(self.ptrs.params_ptr, ¶m_sizes);
|
||||
let states_ptr = if self.config.bottleneck_dim > 0 {
|
||||
self.bn_concat_buf.raw_ptr()
|
||||
} else {
|
||||
bw_raw_f32_ptr(&self.states_buf, &self.stream)
|
||||
};
|
||||
let h_s1_save_ptr = bw_raw_f32_ptr(&self.save_h_s1, &self.stream);
|
||||
let d_h_s2 = self.ptrs.bw_d_h_s2;
|
||||
let d_h_s1 = bw_raw_f32_ptr(&self.bw_d_h_s1, &self.stream);
|
||||
let trunk_scratch_base = self.ptrs.iqn_trunk_m;
|
||||
|
||||
// ── 4. Backward FC layer 2 (f32 throughout) ──
|
||||
// Computes dW_s2, db_s2 into scratch (iqn_trunk_m), d_h_s1 into bw_d_h_s1.
|
||||
{
|
||||
let dy = self.ptrs.bw_d_h_s2; // f32
|
||||
let x = bw_raw_f32_ptr(&self.save_h_s1, &self.stream);
|
||||
let param_sizes = compute_param_sizes(&self.config);
|
||||
let w_ptrs = f32_weight_ptrs_from_base(self.ptrs.params_ptr, ¶m_sizes);
|
||||
let w = w_ptrs[2]; // w_s2
|
||||
let grn_h_s1_d_pre_ln_p = self.bw_grn_h_s1_d_pre_ln.raw_ptr();
|
||||
let grn_h_s1_d_linear_b_out_p= self.bw_grn_h_s1_d_linear_b_out.raw_ptr();
|
||||
let grn_h_s1_d_elu_out_p = self.bw_grn_h_s1_d_elu_out.raw_ptr();
|
||||
let grn_h_s1_d_linear_a_out_p= self.bw_grn_h_s1_d_linear_a_out.raw_ptr();
|
||||
let grn_h_s2_d_pre_ln_p = self.bw_grn_h_s2_d_pre_ln.raw_ptr();
|
||||
let grn_h_s2_d_linear_b_out_p= self.bw_grn_h_s2_d_linear_b_out.raw_ptr();
|
||||
let grn_h_s2_d_elu_out_p = self.bw_grn_h_s2_d_elu_out.raw_ptr();
|
||||
let grn_h_s2_d_linear_a_out_p= self.bw_grn_h_s2_d_linear_a_out.raw_ptr();
|
||||
|
||||
let scratch_base = self.ptrs.iqn_trunk_m;
|
||||
let f32_u = f32_size as u64;
|
||||
let dw = scratch_base + (w_s1_n + b_s1_n) as u64 * f32_u; // goff_w_s2 in scratch
|
||||
let db = dw + w_s2_n as u64 * f32_u; // goff_b_s2 in scratch
|
||||
self.cublas_forward.encoder_backward_chain(
|
||||
&self.stream,
|
||||
&self.cublas_backward,
|
||||
states_ptr,
|
||||
&w_ptrs,
|
||||
trunk_scratch_base,
|
||||
¶m_sizes,
|
||||
d_h_s2,
|
||||
d_h_s1,
|
||||
0, // IQN: no bottleneck dX needed (separate gradient budget)
|
||||
grn_h_s1_d_pre_ln_p,
|
||||
grn_h_s1_d_linear_b_out_p,
|
||||
grn_h_s1_d_elu_out_p,
|
||||
grn_h_s1_d_linear_a_out_p,
|
||||
grn_h_s2_d_pre_ln_p,
|
||||
grn_h_s2_d_linear_b_out_p,
|
||||
grn_h_s2_d_elu_out_p,
|
||||
grn_h_s2_d_linear_a_out_p,
|
||||
h_s1_save_ptr,
|
||||
)?;
|
||||
|
||||
let dx = bw_raw_f32_ptr(&self.bw_d_h_s1, &self.stream);
|
||||
|
||||
self.cublas_backward.backward_fc_layer(
|
||||
&self.stream, dy, x, w, dw, db, dx, sh2, sh1, b,
|
||||
)?;
|
||||
}
|
||||
|
||||
// ── 5. ReLU mask: bw_d_h_s1 *= (save_h_s1 > 0) (f32 dx, f32 act) ──
|
||||
{
|
||||
let d_ptr = self.ptrs.bw_d_h_s1;
|
||||
let act_ptr = self.ptrs.save_h_s1;
|
||||
let n_relu = (b * sh1) as i32;
|
||||
let blocks = ((b * sh1 + 255) / 256) as u32;
|
||||
unsafe {
|
||||
self.stream
|
||||
.launch_builder(&self.relu_mask_kernel)
|
||||
.arg(&d_ptr)
|
||||
.arg(&act_ptr)
|
||||
.arg(&n_relu)
|
||||
.launch(LaunchConfig {
|
||||
grid_dim: (blocks, 1, 1),
|
||||
block_dim: (256, 1, 1),
|
||||
shared_mem_bytes: 0,
|
||||
})
|
||||
.map_err(|e| MLError::ModelError(format!("relu_mask h_s1: {e}")))?;
|
||||
}
|
||||
}
|
||||
|
||||
// ── 6. Backward FC layer 1 (f32 throughout) ──
|
||||
// Computes dW_s1, db_s1 into scratch. dx=0 (skip input gradient).
|
||||
{
|
||||
let dy = self.ptrs.bw_d_h_s1; // f32
|
||||
let x = bw_raw_f32_ptr(&self.states_buf, &self.stream);
|
||||
let param_sizes = compute_param_sizes(&self.config);
|
||||
let w_ptrs = f32_weight_ptrs_from_base(self.ptrs.params_ptr, ¶m_sizes);
|
||||
let w = w_ptrs[0]; // w_s1
|
||||
|
||||
let scratch_base = self.ptrs.iqn_trunk_m;
|
||||
let f32_u = f32_size as u64;
|
||||
let dw = scratch_base; // goff_w_s1 in scratch
|
||||
let db = scratch_base + w_s1_n as u64 * f32_u; // goff_b_s1 in scratch
|
||||
|
||||
// states_buf has padded stride pad128(sd) — use _lda variant
|
||||
let sd_padded = pad128(sd);
|
||||
self.cublas_backward.backward_fc_layer_lda(
|
||||
&self.stream, dy, x, w, dw, db, 0, sh1, sd, sd_padded, b,
|
||||
)?;
|
||||
}
|
||||
|
||||
// ── 7. Plain SAXPY: grad_buf[trunk] += iqn_lambda * iqn_budget * scratch ──
|
||||
// Adaptive lambda: scales with IQN loss readiness (0 when uncertain, 1 when converged).
|
||||
// B4/G5: further scaled by iqn_budget (0.10+0.30×health).
|
||||
// Uses aux_child-specific handle — saxpy_f32_kernel is captured in forward_child.
|
||||
// ── 4. Plain SAXPY: grad_buf[trunk] += iqn_lambda * iqn_budget * scratch ──
|
||||
// Adaptive lambda: scales with IQN loss readiness (0 when uncertain,
|
||||
// 1 when converged). B4/G5: further scaled by iqn_budget.
|
||||
// Uses aux_child-specific handle — saxpy_f32_kernel is captured in
|
||||
// forward_child and a separate handle is needed here in aux_child.
|
||||
{
|
||||
let scratch_ptr = self.ptrs.iqn_trunk_m;
|
||||
let grad_ptr = self.ptrs.grad_buf;
|
||||
@@ -5827,38 +5778,28 @@ impl GpuDqnTrainer {
|
||||
/// instead of h_s2. The `scale` parameter is `ensemble_diversity_weight`.
|
||||
///
|
||||
/// All computation is GPU-only — no CPU readback in the gradient path.
|
||||
#[allow(clippy::too_many_arguments, unreachable_code, unused_variables)]
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub fn apply_ensemble_diversity_backward(
|
||||
&mut self,
|
||||
d_logits_ptr: u64,
|
||||
scale: f32,
|
||||
) -> Result<(), MLError> {
|
||||
// Plan 4 Task 2c.3a: trunk param layout migrated to GRN. This
|
||||
// auxiliary ensemble-diversity trunk-backward path runs the
|
||||
// legacy Linear→ReLU→Linear backward GEMMs through the value
|
||||
// head and into the trunk (param indices [0..4) → GRN h_s1).
|
||||
// Wired to GRN by Task 2c.3c alongside the main trunk backward.
|
||||
panic!("Plan 4 Task 2c.3a: trunk param layout migrated to GRN, but \
|
||||
apply_ensemble_diversity_backward still runs the legacy \
|
||||
trunk backward against GRN-shaped param tensors. Wired by \
|
||||
Task 2c.3c alongside BatchedBackward::backward_full.");
|
||||
// Plan 4 Task 2c.3c.4: ensemble-diversity trunk path now mirrors the
|
||||
// GRN-aware structure — value head dX into d_h_s2 (LN, no ReLU mask),
|
||||
// then encoder_backward_chain into iqn_trunk_m, then SAXPY into
|
||||
// grad_buf. Value-head weight indices fixed: w_v1 4→13, w_v2 6→15.
|
||||
let b = self.config.batch_size;
|
||||
let _eg = EventTrackingGuard::new(self.stream.context());
|
||||
let sd = ml_core::state_layout::STATE_DIM;
|
||||
let sh1 = self.config.shared_h1;
|
||||
let sh2 = self.config.shared_h2;
|
||||
let vh = self.config.value_h;
|
||||
let na = self.config.num_atoms;
|
||||
let f32_size = std::mem::size_of::<f32>(); // iqn_trunk_m scratch is now f32
|
||||
let f32_size = std::mem::size_of::<f32>();
|
||||
|
||||
// Trunk gradient element counts (for the scratch -> SAXPY step)
|
||||
let w_s1_n = sh1 * sd;
|
||||
let b_s1_n = sh1;
|
||||
let w_s2_n = sh2 * sh1;
|
||||
let _b_s2_n = sh2;
|
||||
let trunk_grad_total = w_s1_n + b_s1_n + w_s2_n + _b_s2_n;
|
||||
let param_sizes = compute_param_sizes(&self.config);
|
||||
let trunk_grad_total: usize =
|
||||
(padded_byte_offset(¶m_sizes, 13) as usize) / f32_size;
|
||||
|
||||
// ── 1. Zero the scratch buffer (iqn_trunk_m) ───────────────────────
|
||||
// ── 1. Zero the scratch buffer (iqn_trunk_m) ──
|
||||
{
|
||||
let scratch_base = self.ptrs.iqn_trunk_m;
|
||||
let n_bytes = trunk_grad_total * f32_size;
|
||||
@@ -5869,25 +5810,20 @@ impl GpuDqnTrainer {
|
||||
}
|
||||
}
|
||||
|
||||
// ── 2. Backward value output layer: d_logits -> d_h_v (f32) ────────
|
||||
// d_logits [B, NA] x W_v2^T [VH, NA] -> d_h_v [B, VH]
|
||||
// Only upstream gradient (dX) is needed -- skip dW/db for value head.
|
||||
// dX output is f32 (gemmex_f32 with TF32).
|
||||
let w_ptrs = f32_weight_ptrs_from_base(self.ptrs.params_ptr, ¶m_sizes);
|
||||
|
||||
// ── 2. Value output layer dX: d_logits → d_h_v ──
|
||||
// Plan 4 Task 2c.3c.4: w_v2 index 6 → 15 (post-GRN reshuffle).
|
||||
{
|
||||
let param_sizes = compute_param_sizes(&self.config);
|
||||
let w_ptrs = f32_weight_ptrs_from_base(self.ptrs.params_ptr, ¶m_sizes);
|
||||
let w_v2 = w_ptrs[6]; // W_v2 [NA, VH]
|
||||
|
||||
let w_v2 = w_ptrs[15]; // W_v2 [NA, VH]
|
||||
let dx = self.ptrs.bw_d_h_v;
|
||||
|
||||
// launch_dx_only: computes only dX = dY @ W^T (no dW/db)
|
||||
self.cublas_backward.launch_dx_only(
|
||||
&self.stream, d_logits_ptr, w_v2, dx,
|
||||
na, vh, b, 0.0,
|
||||
)?;
|
||||
}
|
||||
|
||||
// ── 3. ReLU mask: d_h_v *= (save_h_v > 0) (f32 dx, f32 act) ──────
|
||||
// ── 3. ReLU mask on d_h_v (value FC head still uses ReLU) ──
|
||||
{
|
||||
let d_ptr = self.ptrs.bw_d_h_v;
|
||||
let act_ptr = self.ptrs.save_h_v;
|
||||
@@ -5908,107 +5844,66 @@ impl GpuDqnTrainer {
|
||||
}
|
||||
}
|
||||
|
||||
// ── 4. Backward value FC → d_h_s2 (f32 throughout) ──
|
||||
// Only upstream gradient (dX) is needed -- skip dW/db for value head.
|
||||
// ── 4. Value FC dX → d_h_s2 (no ReLU mask: trunk now ends in LN) ──
|
||||
// Plan 4 Task 2c.3c.4: w_v1 index 4 → 13 (post-GRN reshuffle).
|
||||
{
|
||||
let dy = self.ptrs.bw_d_h_v; // f32
|
||||
|
||||
let param_sizes = compute_param_sizes(&self.config);
|
||||
let w_ptrs = f32_weight_ptrs_from_base(self.ptrs.params_ptr, ¶m_sizes);
|
||||
let w_v1 = w_ptrs[4]; // W_v1 [VH, SH2]
|
||||
|
||||
let dy = self.ptrs.bw_d_h_v;
|
||||
let w_v1 = w_ptrs[13]; // W_v1 [VH, SH2]
|
||||
let dx = self.ptrs.bw_d_h_s2;
|
||||
|
||||
// launch_dx_only: computes only dX = dY @ W^T (no dW/db)
|
||||
self.cublas_backward.launch_dx_only(
|
||||
&self.stream, dy, w_v1, dx,
|
||||
vh, sh2, b, 0.0,
|
||||
)?;
|
||||
}
|
||||
|
||||
// ── 5. ReLU mask: d_h_s2 *= (save_h_s2 > 0) (f32 dx, f32 act) ───
|
||||
{
|
||||
let d_ptr = self.ptrs.bw_d_h_s2;
|
||||
let act_ptr = self.ptrs.save_h_s2;
|
||||
let n_relu = (b * sh2) as i32;
|
||||
let blocks = ((b * sh2 + 255) / 256) as u32;
|
||||
unsafe {
|
||||
self.stream
|
||||
.launch_builder(&self.relu_mask_kernel)
|
||||
.arg(&d_ptr)
|
||||
.arg(&act_ptr)
|
||||
.arg(&n_relu)
|
||||
.launch(LaunchConfig {
|
||||
grid_dim: (blocks, 1, 1),
|
||||
block_dim: (256, 1, 1),
|
||||
shared_mem_bytes: 0,
|
||||
})
|
||||
.map_err(|e| MLError::ModelError(format!("relu_mask ens_h_s2: {e}")))?;
|
||||
}
|
||||
}
|
||||
// ── 5. GRN trunk backward (encoder_backward_chain) ──
|
||||
// h_s2's gradient flows from d_h_s2 (post-LN) into the trunk via the
|
||||
// GRN composition (no ReLU truncation). Bottleneck dX is unused for
|
||||
// ensemble (separate gradient budget) so s1_dx_output = 0.
|
||||
let states_ptr = if self.config.bottleneck_dim > 0 {
|
||||
self.bn_concat_buf.raw_ptr()
|
||||
} else {
|
||||
bw_raw_f32_ptr(&self.states_buf, &self.stream)
|
||||
};
|
||||
let h_s1_save_ptr = bw_raw_f32_ptr(&self.save_h_s1, &self.stream);
|
||||
let d_h_s2 = self.ptrs.bw_d_h_s2;
|
||||
let d_h_s1 = bw_raw_f32_ptr(&self.bw_d_h_s1, &self.stream);
|
||||
let trunk_scratch_base = self.ptrs.iqn_trunk_m;
|
||||
|
||||
// ── 6. Backward FC layer 2 (f32 throughout) ───
|
||||
{
|
||||
let dy = self.ptrs.bw_d_h_s2; // f32
|
||||
let x = bw_raw_f32_ptr(&self.save_h_s1, &self.stream);
|
||||
let param_sizes = compute_param_sizes(&self.config);
|
||||
let w_ptrs = f32_weight_ptrs_from_base(self.ptrs.params_ptr, ¶m_sizes);
|
||||
let w = w_ptrs[2]; // w_s2
|
||||
let grn_h_s1_d_pre_ln_p = self.bw_grn_h_s1_d_pre_ln.raw_ptr();
|
||||
let grn_h_s1_d_linear_b_out_p= self.bw_grn_h_s1_d_linear_b_out.raw_ptr();
|
||||
let grn_h_s1_d_elu_out_p = self.bw_grn_h_s1_d_elu_out.raw_ptr();
|
||||
let grn_h_s1_d_linear_a_out_p= self.bw_grn_h_s1_d_linear_a_out.raw_ptr();
|
||||
let grn_h_s2_d_pre_ln_p = self.bw_grn_h_s2_d_pre_ln.raw_ptr();
|
||||
let grn_h_s2_d_linear_b_out_p= self.bw_grn_h_s2_d_linear_b_out.raw_ptr();
|
||||
let grn_h_s2_d_elu_out_p = self.bw_grn_h_s2_d_elu_out.raw_ptr();
|
||||
let grn_h_s2_d_linear_a_out_p= self.bw_grn_h_s2_d_linear_a_out.raw_ptr();
|
||||
|
||||
let scratch_base = self.ptrs.iqn_trunk_m;
|
||||
let f32_u = f32_size as u64;
|
||||
let dw = scratch_base + (w_s1_n + b_s1_n) as u64 * f32_u;
|
||||
let db = dw + w_s2_n as u64 * f32_u;
|
||||
self.cublas_forward.encoder_backward_chain(
|
||||
&self.stream,
|
||||
&self.cublas_backward,
|
||||
states_ptr,
|
||||
&w_ptrs,
|
||||
trunk_scratch_base,
|
||||
¶m_sizes,
|
||||
d_h_s2,
|
||||
d_h_s1,
|
||||
0,
|
||||
grn_h_s1_d_pre_ln_p,
|
||||
grn_h_s1_d_linear_b_out_p,
|
||||
grn_h_s1_d_elu_out_p,
|
||||
grn_h_s1_d_linear_a_out_p,
|
||||
grn_h_s2_d_pre_ln_p,
|
||||
grn_h_s2_d_linear_b_out_p,
|
||||
grn_h_s2_d_elu_out_p,
|
||||
grn_h_s2_d_linear_a_out_p,
|
||||
h_s1_save_ptr,
|
||||
)?;
|
||||
|
||||
let dx = bw_raw_f32_ptr(&self.bw_d_h_s1, &self.stream);
|
||||
|
||||
self.cublas_backward.backward_fc_layer(
|
||||
&self.stream, dy, x, w, dw, db, dx, sh2, sh1, b,
|
||||
)?;
|
||||
}
|
||||
|
||||
// ── 7. ReLU mask: d_h_s1 *= (save_h_s1 > 0) (f32 dx, f32 act) ───
|
||||
{
|
||||
let d_ptr = self.ptrs.bw_d_h_s1;
|
||||
let act_ptr = self.ptrs.save_h_s1;
|
||||
let n_relu = (b * sh1) as i32;
|
||||
let blocks = ((b * sh1 + 255) / 256) as u32;
|
||||
unsafe {
|
||||
self.stream
|
||||
.launch_builder(&self.relu_mask_kernel)
|
||||
.arg(&d_ptr)
|
||||
.arg(&act_ptr)
|
||||
.arg(&n_relu)
|
||||
.launch(LaunchConfig {
|
||||
grid_dim: (blocks, 1, 1),
|
||||
block_dim: (256, 1, 1),
|
||||
shared_mem_bytes: 0,
|
||||
})
|
||||
.map_err(|e| MLError::ModelError(format!("relu_mask ens_h_s1: {e}")))?;
|
||||
}
|
||||
}
|
||||
|
||||
// ── 8. Backward FC layer 1 (f32 throughout) ───
|
||||
{
|
||||
let dy = self.ptrs.bw_d_h_s1; // f32
|
||||
let x = bw_raw_f32_ptr(&self.states_buf, &self.stream);
|
||||
let param_sizes = compute_param_sizes(&self.config);
|
||||
let w_ptrs = f32_weight_ptrs_from_base(self.ptrs.params_ptr, ¶m_sizes);
|
||||
let w = w_ptrs[0]; // w_s1
|
||||
|
||||
let scratch_base = self.ptrs.iqn_trunk_m;
|
||||
let f32_u = f32_size as u64;
|
||||
let dw = scratch_base;
|
||||
let db = scratch_base + w_s1_n as u64 * f32_u;
|
||||
|
||||
self.cublas_backward.backward_fc_layer(
|
||||
&self.stream, dy, x, w, dw, db, 0, sh1, sd, b,
|
||||
)?;
|
||||
}
|
||||
|
||||
// ── 9. Plain SAXPY: grad_buf[trunk] += scale * scratch ────
|
||||
// ── 6. Plain SAXPY: grad_buf[trunk] += scale * scratch ──
|
||||
// No per-component clip — single global clip in Adam handles safety.
|
||||
// Uses aux_child-specific handle — saxpy_f32_kernel is captured in forward_child.
|
||||
// Uses aux_child-specific handle — saxpy_f32_kernel is captured in
|
||||
// forward_child and a separate handle is needed here in aux_child.
|
||||
{
|
||||
let scratch_ptr = self.ptrs.iqn_trunk_m;
|
||||
let grad_ptr = self.ptrs.grad_buf;
|
||||
@@ -6485,6 +6380,45 @@ impl GpuDqnTrainer {
|
||||
)?;
|
||||
}
|
||||
|
||||
// ── GRN trunk backward (Plan 4 Task 2c.3c.4, CQL path) ──
|
||||
// Mirrors the main backward path but writes trunk gradients into
|
||||
// `cql_grad_scratch` instead of `grad_buf`. CQL never asks for the
|
||||
// trunk-input gradient (`s1_dx_output = 0`), so the bottleneck dX
|
||||
// GEMM short-circuits and `cql_grad_scratch` accumulates the GRN
|
||||
// trunk dW/dB/dgamma/dbeta only — exactly what `apply_cql_saxpy`
|
||||
// expects to mix into `grad_buf` afterwards.
|
||||
let grn_h_s1_d_pre_ln_cql = self.bw_grn_h_s1_d_pre_ln.raw_ptr();
|
||||
let grn_h_s1_d_linear_b_out_cql= self.bw_grn_h_s1_d_linear_b_out.raw_ptr();
|
||||
let grn_h_s1_d_elu_out_cql = self.bw_grn_h_s1_d_elu_out.raw_ptr();
|
||||
let grn_h_s1_d_linear_a_out_cql= self.bw_grn_h_s1_d_linear_a_out.raw_ptr();
|
||||
let grn_h_s2_d_pre_ln_cql = self.bw_grn_h_s2_d_pre_ln.raw_ptr();
|
||||
let grn_h_s2_d_linear_b_out_cql= self.bw_grn_h_s2_d_linear_b_out.raw_ptr();
|
||||
let grn_h_s2_d_elu_out_cql = self.bw_grn_h_s2_d_elu_out.raw_ptr();
|
||||
let grn_h_s2_d_linear_a_out_cql= self.bw_grn_h_s2_d_linear_a_out.raw_ptr();
|
||||
let cql_grad_scratch_base = self.cql_grad_scratch.raw_ptr();
|
||||
let cql_param_sizes = compute_param_sizes(&self.config);
|
||||
let cql_w_ptrs = f32_weight_ptrs_from_base(self.ptrs.params_ptr, &cql_param_sizes);
|
||||
self.cublas_forward.encoder_backward_chain(
|
||||
&self.stream,
|
||||
&self.cublas_backward,
|
||||
states_ptr_fw,
|
||||
&cql_w_ptrs,
|
||||
cql_grad_scratch_base,
|
||||
&cql_param_sizes,
|
||||
scratch_d_h_s2,
|
||||
scratch_d_h_s1,
|
||||
0, // CQL: no bottleneck dX needed
|
||||
grn_h_s1_d_pre_ln_cql,
|
||||
grn_h_s1_d_linear_b_out_cql,
|
||||
grn_h_s1_d_elu_out_cql,
|
||||
grn_h_s1_d_linear_a_out_cql,
|
||||
grn_h_s2_d_pre_ln_cql,
|
||||
grn_h_s2_d_linear_b_out_cql,
|
||||
grn_h_s2_d_elu_out_cql,
|
||||
grn_h_s2_d_linear_a_out_cql,
|
||||
h_s1_ptr,
|
||||
)?;
|
||||
|
||||
Ok(true)
|
||||
}
|
||||
|
||||
@@ -7953,8 +7887,19 @@ impl GpuDqnTrainer {
|
||||
let init_v_s2 = rand_unit(config.shared_h1, &mut rng_state);
|
||||
|
||||
// IQN trunk Adam state (separate from C51 Adam — prevents momentum mismatch)
|
||||
let trunk_params = config.shared_h1 * ml_core::state_layout::STATE_DIM + config.shared_h1
|
||||
+ config.shared_h2 * config.shared_h1 + config.shared_h2;
|
||||
//
|
||||
// Plan 4 Task 2c.3c.4: post-GRN reshuffle the trunk owns 13 tensors at
|
||||
// indices [0..13) (h_s1: 7 tensors with Linear_residual; h_s2: 6 tensors
|
||||
// with identity residual). To keep the SAXPY-based mix-into-grad_buf
|
||||
// pattern simple, `iqn_trunk_m` mirrors `grad_buf`'s padded layout for
|
||||
// exactly those 13 tensors. The element count below = total padded f32
|
||||
// slots from index 0 up to (but excluding) tensor 13. Element-wise
|
||||
// SAXPY across this many floats writes to the same per-tensor offsets
|
||||
// in grad_buf because the layout is identical (padding inclusive).
|
||||
let __iqn_trunk_param_sizes = compute_param_sizes(&config);
|
||||
let trunk_params: usize =
|
||||
(padded_byte_offset(&__iqn_trunk_param_sizes, 13) as usize)
|
||||
/ std::mem::size_of::<f32>();
|
||||
let iqn_trunk_m = stream.alloc_zeros::<f32>(trunk_params)
|
||||
.map_err(|e| MLError::ModelError(format!("alloc iqn_trunk_m f32: {e}")))?;
|
||||
let iqn_trunk_v = alloc_f32(&stream, trunk_params, "iqn_trunk_v")?;
|
||||
@@ -10420,8 +10365,8 @@ impl GpuDqnTrainer {
|
||||
/// Read a scalar from the ISV pinned buffer at the given index.
|
||||
/// Returns 0.0 if the pointer is null or index >= ISV_TOTAL_DIM.
|
||||
pub fn read_isv_signal_at(&self, index: usize) -> f32 {
|
||||
if self.isv_signals_pinned.is_null() { return 0.0; }
|
||||
if index >= ISV_TOTAL_DIM { return 0.0; }
|
||||
if self.isv_signals_pinned.is_null() { return 0.0; } // ok: defensive null-guard on pinned ISV buffer; documented contract
|
||||
if index >= ISV_TOTAL_DIM { return 0.0; } // ok: bounds-check return value (documented in fn doc)
|
||||
unsafe { *self.isv_signals_pinned.add(index) }
|
||||
}
|
||||
|
||||
@@ -10461,7 +10406,7 @@ impl GpuDqnTrainer {
|
||||
/// Read atom utilization from q_readback_pinned[6] (one-step lag, written by
|
||||
/// reduce_current_q_stats). Returns 0.0 if pointer is null.
|
||||
pub fn read_atom_utilization(&self) -> f32 {
|
||||
if self.q_readback_pinned.is_null() { return 0.0; }
|
||||
if self.q_readback_pinned.is_null() { return 0.0; } // ok: defensive null-guard on pinned readback buffer; documented contract
|
||||
unsafe { (*self.q_readback_pinned.add(6)).clamp(0.0, 1.0) }
|
||||
}
|
||||
|
||||
@@ -10502,11 +10447,10 @@ impl GpuDqnTrainer {
|
||||
/// The downstream `NormalizedComponents::from_raw` thresholds are calibrated
|
||||
/// for this behaviour (`spectral_gap_norm` drops toward 0 as this value rises).
|
||||
pub fn compute_q_spectral_gap(&mut self) -> f32 {
|
||||
if self.q_readback_pinned.is_null() { return 1.0; }
|
||||
|
||||
if self.q_readback_pinned.is_null() { return 1.0; } // ok: defensive null-guard; 1.0 = "no spectral gap signal" (downstream calibration treats high values as collapse)
|
||||
// Pull the latest sample (total_actions floats) from the readback head.
|
||||
let n_cols = self.total_actions();
|
||||
if n_cols == 0 { return 1.0; }
|
||||
if n_cols == 0 { return 1.0; } // ok: degenerate config — no actions configured; return neutral 1.0 (matches null-guard semantics)
|
||||
let mut sample = Vec::with_capacity(n_cols);
|
||||
unsafe {
|
||||
let base = self.q_readback_pinned.add(7);
|
||||
@@ -14659,13 +14603,18 @@ impl GpuDqnTrainer {
|
||||
/// `d_adv_logits_buf`) directly. The entire backward
|
||||
/// chain stays in f32 to preserve gradient precision at large batch sizes.
|
||||
/// Accumulates weight gradients into `grad_buf`.
|
||||
pub(crate) fn launch_cublas_backward(&self) -> Result<(), MLError> {
|
||||
pub(crate) fn launch_cublas_backward(&mut self) -> Result<(), MLError> {
|
||||
self.launch_cublas_backward_to(self.ptrs.grad_buf)
|
||||
}
|
||||
|
||||
/// Backward pass writing weight gradients to an arbitrary output buffer.
|
||||
/// Used by the vaccine to write g_val to scratch without touching grad_buf.
|
||||
fn launch_cublas_backward_to(&self, grad_base: u64) -> Result<(), MLError> {
|
||||
///
|
||||
/// Plan 4 Task 2c.3c.4: takes `&mut self` because the GRN trunk-backward
|
||||
/// chain (`encoder_backward_chain`) borrows `&mut self.cublas_forward`
|
||||
/// (the GRN block's `backward_raw_phase1`/`backward_raw_phase2` need to
|
||||
/// scribble per-block partial-reduction scratch).
|
||||
fn launch_cublas_backward_to(&mut self, grad_base: u64) -> Result<(), MLError> {
|
||||
let bw = &self.cublas_backward;
|
||||
|
||||
let param_sizes = compute_param_sizes(&self.config);
|
||||
@@ -14782,6 +14731,34 @@ impl GpuDqnTrainer {
|
||||
)?;
|
||||
}
|
||||
|
||||
// ── GRN trunk backward (Plan 4 Task 2c.3c.4) ──
|
||||
// d_h_s2_ptr now holds the fully-accumulated gradient flowing into
|
||||
// the h_s2 GRN output. encoder_backward_chain unrolls both GRN
|
||||
// blocks symmetrically with `encoder_forward_only`, writing dW/dB
|
||||
// into goff[0..13) of `grad_base` and producing the trunk input
|
||||
// gradient at `s1_dx_output` (which `bn_d_concat_buf` aliases when
|
||||
// bottleneck is active).
|
||||
self.cublas_forward.encoder_backward_chain(
|
||||
&self.stream,
|
||||
bw,
|
||||
states_ptr,
|
||||
&w_ptrs,
|
||||
grad_base,
|
||||
¶m_sizes,
|
||||
d_h_s2_ptr,
|
||||
d_h_s1_ptr,
|
||||
s1_dx_output,
|
||||
self.bw_grn_h_s1_d_pre_ln.raw_ptr(),
|
||||
self.bw_grn_h_s1_d_linear_b_out.raw_ptr(),
|
||||
self.bw_grn_h_s1_d_elu_out.raw_ptr(),
|
||||
self.bw_grn_h_s1_d_linear_a_out.raw_ptr(),
|
||||
self.bw_grn_h_s2_d_pre_ln.raw_ptr(),
|
||||
self.bw_grn_h_s2_d_linear_b_out.raw_ptr(),
|
||||
self.bw_grn_h_s2_d_elu_out.raw_ptr(),
|
||||
self.bw_grn_h_s2_d_linear_a_out.raw_ptr(),
|
||||
h_s1_ptr,
|
||||
)?;
|
||||
|
||||
// ── #31 Bottleneck backward: chain rule through tanh + GEMM ──
|
||||
// After backward_full, d_bn_concat_buf has dL/d(bn_concat) [B, concat_dim] f32.
|
||||
// Chain rule: dL/d_bn = dL/d_concat[:, :bn_dim] * tanh'(bn_raw)
|
||||
|
||||
@@ -291,6 +291,8 @@ Plan 4 Task 2c.3c.2 (2026-04-24): backward infrastructure additions — two help
|
||||
|
||||
Plan 4 Task 2c.3c.3 (2026-04-24): GRN trunk backward scratch buffer allocation on `GpuDqnTrainer`. 8 new `CudaSlice<f32>` device buffers added (4 per GRN block × 2 blocks for h_s1 + h_s2): `bw_grn_h_s2_d_pre_ln` [B, SH2], `bw_grn_h_s2_d_linear_b_out` [B, 2*SH2], `bw_grn_h_s2_d_elu_out` [B, SH2], `bw_grn_h_s2_d_linear_a_out` [B, SH2], plus the four h_s1 mirrors at [B, SH1] (and 2*SH1 for the linear_b_out scratch — GLU's pre-activation has 2× hidden width because value-path + gate_pre are concatenated). The `d_pre_LN` buffer aliases both `d_glu_out` and `d_residual` (pure pass-through under `GrnBlock::backward_raw_phase1`'s pointer aliasing convention from 2c.3c.1); `d_linear_b_out` carries the phase1 → caller's Linear_b backward GEMM gradient; `d_elu_out` carries Linear_b backward's output into phase2; `d_linear_a_out` carries phase2's output to the caller's Linear_a backward GEMM. Allocated alongside the existing `alloc_backward_scratch` call in the trainer constructor (mirrors that helper's `stream.alloc_zeros::<f32>(size).map_err(...)?` pattern). Total footprint at SH1=SH2=256, B=512: 8 × 5 × 256 × 4B = 5.0 MB per fold (negligible against existing ~32 MB per-branch cuBLAS workspaces). Each field is `#[allow(dead_code)]` until 2c.3c.4's `apply_grn_trunk_backward_raw` helper consumes them — the comment `wired by Task 2c.3c.4` annotates each suppression. **Additive only — ZERO production callers in this commit**; the three backward panic gates (`backward_full`, `apply_iqn_trunk_gradient`, `apply_ensemble_diversity_backward`) remain in place until 2c.3c.4. Adam state for the new GRN tensors is auto-allocated since `m_buf`/`v_buf` are sized to `total_params + cutlass_tile_pad` which already includes the 13 GRN tensors after 2c.3a. cargo check clean at 11 warnings (baseline preserved); cargo build compiles 80 cubins (unchanged — no new kernel). +57 LOC. No new module / kernel / ISV slot / Orphan row.
|
||||
|
||||
Plan 4 Task 2c.3c.4 (2026-04-25): GRN backward chain wired through all 3 panic-gated trunk-backward sites — smoke validates. Three panic gates removed: `BatchedBackward::backward_full`, `GpuDqnTrainer::apply_iqn_trunk_gradient`, `GpuDqnTrainer::apply_ensemble_diversity_backward`. New helper `BatchedForward::encoder_backward_chain` orchestrates the full GRN trunk backward composition (h_s2 → h_s1) symmetrically with `encoder_forward_only`: `backward_raw_phase1` (LN_dx + LN_dgamma/dbeta no-atomicAdd reduction + GLU_bwd) → cuBLAS Linear_b dW + dX → `backward_raw_phase2` (ELU_bwd) → cuBLAS Linear_a dW + dX → for h_s1 only: Linear_residual dW (no_bias_lda) + dX accumulation when bottleneck active → for h_s2 only: `saxpy_inplace(alpha=1.0)` accumulates `d_pre_ln_h_s2` into `d_h_s1` (identity residual gradient). Used by 4 callers writing to 3 different gradient targets: main backward (writes to `grad_buf` for the 13 GRN trunk tensors at indices [0..13)), CQL backward (writes to `cql_grad_scratch` — same layout as grad_buf), and the two auxiliary paths (`apply_iqn_trunk_gradient` and `apply_ensemble_diversity_backward` write to `iqn_trunk_m` then SAXPY into `grad_buf` with their respective scale factors). `iqn_trunk_m` grown from legacy 4-tensor element count (`sh1*sd + sh1 + sh2*sh1 + sh2`) to `padded_byte_offset(¶m_sizes, 13) / sizeof(f32)` so the layout matches `grad_buf`'s padded byte offsets exactly for the trunk portion — element-wise SAXPY across that span lands every per-tensor gradient at the correct location without needing per-tensor offset computation. `apply_ensemble_diversity_backward` additionally fixes value-head weight indices: `w_v1` 4→13, `w_v2` 6→15 (legacy indices were stale post-2c.3a — caught by the panic gate before they could write garbage). `BatchedBackward::launch_dw_only_no_bias_lda` added because Linear_residual (h_s1, no bias) needs the padded states stride that plain `launch_dw_only_no_bias` doesn't accept; `launch_dw_only` would dereference NULL inside the bias-grad kernel. ReLU masks on `bw_d_h_s2` and `bw_d_h_s1` are gone — the GRN trunk ends in LayerNorm which has no truncation derivative; the value-head FC retains its ReLU mask (value head still uses ReLU for its hidden activation). `launch_cublas_backward_to` migrated from `&self` to `&mut self` because the GRN backward needs to scribble per-block partial-reduction scratch inside `grn_h_s{1,2}_online`; the borrow split lets the trainer hand `&mut BatchedForward` + `&BatchedBackward` to `encoder_backward_chain` simultaneously. Smoke (`cargo test … multi_fold_convergence --ignored`, 599.73s, 3 folds × 5 epochs): all 3 `dqn_fold{N}_best.safetensors` checkpoints written; per-fold best train Sharpe 7.52 / 60.94 / 10.40 at epochs 2 / 4 / 2; per-fold val_Sharpe peaks 0.51 / 0.64 / 0.82; fold 2 epoch 5 hit PF=5.29 with +988% return — gradients flow cleanly through the GRN composition (no NaN, no Inf, no policy collapse, no fold-boundary explosion). Audit row #11 (IQN target trunk orphan) was already retired by 2c.3b; this commit retires the 3 backward panic gates that 2c.3a introduced. cargo check clean at 11 warnings (baseline preserved). +582 / -320 LOC across `batched_backward.rs` (+154/-154 — panic-gate removal + `launch_dw_only_no_bias_lda`), `batched_forward.rs` (+284/-0 — `encoder_backward_chain`), and `gpu_dqn_trainer.rs` (+221/-243 — 3 trunk-backward call sites rewritten + `iqn_trunk_m` grown + 4 main/CQL/IQN/ensemble backward call sites updated). No new module / kernel / ISV slot. Closes 3 of the 6 panic gates introduced by 2c.3a (3 forward gates were retired by 2c.3b); panic-gate count: 6 → 0.
|
||||
|
||||
Plan 4 Task 2c.3a (2026-04-24): GRN trunk param-tensor reshuffle (build-clean, runtime-broken intentionally). `compute_param_sizes()` flat layout reshuffled from 86 → 95 tensors: the 4 legacy trunk Linear tensors (`w_s1`, `b_s1`, `w_s2`, `b_s2` at indices [0..4)) are replaced with **13 GRN tensors** at indices [0..13) — 7 for h_s1 GRN block (`w_a`, `b_a`, `w_b`, `b_b`, `w_residual`, `gamma`, `beta`; SD→SH1 with `Linear_residual` GEMM since dimensions differ) plus 6 for h_s2 GRN block (`w_a`, `b_a`, `w_b`, `b_b`, `gamma`, `beta`; SH1→SH2 with identity residual since SH1==SH2). Every tensor index ≥ 4 in the OLD layout shifts +9 in the NEW layout. `NUM_WEIGHT_TENSORS` 86→95, `FIRST_ISV_TENSOR` 68→77. `layout_fingerprint_seed()` updated to mirror the new tensor names + positions; new `LAYOUT_FINGERPRINT_CURRENT = 0xcf3a24b0a1f70057` (was `0xa504d3c2f275b8af`). All 93 `padded_byte_offset` call sites + ancillary index references migrated in lockstep per `feedback_no_partial_refactor.md`. `xavier_init_params_buf` extended with init paths for the 13 new tensors: Linear matrices (`w_a`, `w_b`, `w_residual`) get Xavier; LayerNorm γ initialised to 1.0; β + biases stay zero. Trunk-forward / trunk-backward / IQN-trunk / ensemble-diversity-backward callers gated by explicit `panic!("Plan 4 Task 2c.3a: trunk param layout migrated to GRN, …")` at function entry — `BatchedForward::encoder_forward_only`, `BatchedForward::forward_target_raw`, `BatchedForward::forward_online_f32`, `BatchedBackward::backward_full`, `GpuDqnTrainer::apply_iqn_trunk_gradient`, `GpuDqnTrainer::apply_ensemble_diversity_backward`. This makes accidental runtime execution fail loudly rather than producing silent garbage from running legacy `Linear→ReLU→Linear` GEMMs against GRN-shaped param tensors. Spectral-norm descriptor (13 matrices) keeps slots [0]/[1] mapped to `w_a_h_s1`/`w_a_h_s2` (shapes match legacy W_s1/W_s2 exactly — the GRN's first Linear has the same shape as the original Linear), so the existing spectral-norm constraint transfers cleanly to the GRN's first Linear; Linear_b / Linear_residual are NOT yet spectral-normed (Task 2c.3c follow-up). **Smoke intentionally NOT run** — build-clean is the validation; runtime would hit the panic at the first encoder forward (the desired behaviour, no need to verify explicitly). Task 2c.3b swaps in the GRN forward (gpu_grn::GrnBlock::forward) at all panic-gated forward callers in place; Task 2c.3c does the same for backward callers and runs the smoke test. cargo check clean at 11 warnings (baseline preserved); cargo build compiles all 58 cubins. No new module / kernel / ISV slot in this commit — pure structural reshuffle + assertion gates.
|
||||
|
||||
| Classification | Count |
|
||||
|
||||
Reference in New Issue
Block a user