From f94d857ebc8511428344d50e670bf539a1cc95aa Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Sat, 25 Apr 2026 15:09:24 +0200 Subject: [PATCH] =?UTF-8?q?feat(dqn-v2):=20Plan=204=20Task=202c.3c.4=20?= =?UTF-8?q?=E2=80=94=20wire=20GRN=20backward=20chain=20through=203=20panic?= =?UTF-8?q?-gated=20trunk-backward=20sites?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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) --- .../ml/src/cuda_pipeline/batched_backward.rs | 154 +++--- .../ml/src/cuda_pipeline/batched_forward.rs | 284 +++++++++++ .../ml/src/cuda_pipeline/gpu_dqn_trainer.rs | 463 +++++++++--------- docs/dqn-wire-up-audit.md | 2 + 4 files changed, 583 insertions(+), 320 deletions(-) diff --git a/crates/ml/src/cuda_pipeline/batched_backward.rs b/crates/ml/src/cuda_pipeline/batched_backward.rs index e7f4c234e..d3504dff7 100644 --- a/crates/ml/src/cuda_pipeline/batched_backward.rs +++ b/crates/ml/src/cuda_pipeline/batched_backward.rs @@ -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, + 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 diff --git a/crates/ml/src/cuda_pipeline/batched_forward.rs b/crates/ml/src/cuda_pipeline/batched_forward.rs index d147ae6a9..dba6482cd 100644 --- a/crates/ml/src/cuda_pipeline/batched_forward.rs +++ b/crates/ml/src/cuda_pipeline/batched_forward.rs @@ -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, + 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. /// diff --git a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs index b3fb66f8b..5064bb572 100644 --- a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs +++ b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs @@ -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::(); // iqn_trunk_m is now f32 + let f32_size = std::mem::size_of::(); - // 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::(); + 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::(); // iqn_trunk_m scratch is now f32 + let f32_size = std::mem::size_of::(); - // 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::(); let iqn_trunk_m = stream.alloc_zeros::(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) diff --git a/docs/dqn-wire-up-audit.md b/docs/dqn-wire-up-audit.md index 06b0d6972..5fb7293b7 100644 --- a/docs/dqn-wire-up-audit.md +++ b/docs/dqn-wire-up-audit.md @@ -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` 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::(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 |