From 524d81a45c514e53b3aca45c1c3340a6edac4ece Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Wed, 15 Apr 2026 00:14:48 +0200 Subject: [PATCH] feat: GLU backward replaces ReLU mask in branch backward pass MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit d_h_bd → glu_backward → d_value + d_gate_pre. Two weight gradient GEMMs per branch (value + gate). Upstream gradient sums both paths. VSN backward DEFERRED (W_vsn init to identity, zero gradient for now). Gradient flows to gate weights at indices 34-41. Co-Authored-By: Claude Opus 4.6 (1M context) --- .../ml/src/cuda_pipeline/batched_backward.rs | 236 +++++++++++++----- .../ml/src/cuda_pipeline/gpu_dqn_trainer.rs | 55 +++- 2 files changed, 221 insertions(+), 70 deletions(-) diff --git a/crates/ml/src/cuda_pipeline/batched_backward.rs b/crates/ml/src/cuda_pipeline/batched_backward.rs index 61d9453c2..5010058e2 100644 --- a/crates/ml/src/cuda_pipeline/batched_backward.rs +++ b/crates/ml/src/cuda_pipeline/batched_backward.rs @@ -87,7 +87,7 @@ use cudarc::cublaslt::result as cublaslt_result; use cudarc::driver::{CudaFunction, CudaSlice, CudaStream, DevicePtr, LaunchConfig, PushKernelArg}; use crate::MLError; -use super::gpu_dqn_trainer::GpuDqnTrainConfig; +use super::gpu_dqn_trainer::{GpuDqnTrainConfig, compute_param_sizes, padded_byte_offset}; // ── Cached backward GEMM descriptor ───────────────────────────────────────── @@ -155,6 +155,11 @@ pub struct CublasBackwardSet { /// `bias_grad_reduce_f32_kernel(dy, db, out_dim, batch_size)` — reduce f32 dY over batch. bias_grad_kernel: CudaFunction, + /// `glu_gate_backward(d_output, gate_pre, value, d_gate_pre, d_value, n)` — + /// GLU backward pass: splits d_h_bd into d_value and d_gate_pre components. + /// Replaces ReLU mask in the branch backward pass. + glu_backward_kernel: CudaFunction, + // ── Network dimensions (baked at construction) ── batch_size: usize, state_dim: usize, @@ -172,6 +177,10 @@ pub struct CublasBackwardSet { branch_2_size: usize, branch_3_size: usize, + /// Pre-computed parameter sizes for all 42 weight tensors. + /// Used to compute `padded_byte_offset` for gradient buffer offsets. + param_sizes: [usize; super::gpu_dqn_trainer::NUM_WEIGHT_TENSORS], + // ── Cached GEMM descriptors (created once at init, reused per-call) ── /// Map from (transa, transb, m, n, k, lda, ldb, ldc) → pre-created descriptors + algo. gemm_cache: HashMap, @@ -192,6 +201,7 @@ impl CublasBackwardSet { pub fn new( shared: Arc, config: &GpuDqnTrainConfig, + glu_backward_kernel: CudaFunction, ) -> Result { let stream = &shared.stream; let lt_raw_handle = shared.lt_handle.0; @@ -200,6 +210,8 @@ impl CublasBackwardSet { // ── Compile helper kernels ────────────────────────────────── let (relu_mask_kernel, bias_grad_kernel) = compile_backward_kernels(stream)?; + let param_sizes = compute_param_sizes(config); + // ── Pre-create backward GEMM descriptors for all unique shapes ── let batch = config.batch_size; let sh1 = config.shared_h1; @@ -274,6 +286,7 @@ impl CublasBackwardSet { handle: shared, relu_mask_kernel, bias_grad_kernel, + glu_backward_kernel, batch_size: config.batch_size, state_dim: config.state_dim, state_dim_padded: sd_pad, @@ -287,6 +300,7 @@ impl CublasBackwardSet { branch_1_size: config.branch_1_size, branch_2_size: config.branch_2_size, branch_3_size: config.branch_3_size, + param_sizes, gemm_cache, }) } @@ -657,6 +671,49 @@ impl CublasBackwardSet { Ok(()) } + /// GLU backward: split upstream gradient `d_output` into `d_value` and + /// `d_gate_pre` using saved forward activations. + /// + /// ```text + /// d_value[i] = d_output[i] * sigmoid(gate_pre[i]) + /// d_gate_pre[i] = d_output[i] * value[i] * sigmoid(gate_pre[i]) * (1 - sigmoid(gate_pre[i])) + /// ``` + /// + /// Replaces `relu_mask` for branch FC layers that use GLU gating. + /// Grid: `ceil(n / 256)`, Block: 256. + pub fn launch_glu_backward( + &self, + stream: &Arc, + d_output: u64, // [B, AH] upstream gradient from branch output backward + gate_pre: u64, // [B, AH] saved pre-sigmoid gate activation + value: u64, // [B, AH] saved value path output + d_gate_pre: u64, // [B, AH] output: gradient for gate path + d_value: u64, // [B, AH] output: gradient for value path + n: usize, + ) -> Result<(), MLError> { + let n_i32 = n as i32; + let blocks = ((n + 255) / 256) as u32; + + unsafe { + stream + .launch_builder(&self.glu_backward_kernel) + .arg(&d_output) + .arg(&gate_pre) + .arg(&value) + .arg(&d_gate_pre) + .arg(&d_value) + .arg(&n_i32) + .launch(LaunchConfig { + grid_dim: (blocks, 1, 1), + block_dim: (256, 1, 1), + shared_mem_bytes: 0, + }) + .map_err(|e| MLError::ModelError(format!("glu_backward_kernel: {e}")))?; + } + + Ok(()) + } + // ══════════════════════════════════════════════════════════════════════════ // Full backward pass // ══════════════════════════════════════════════════════════════════════════ @@ -684,17 +741,21 @@ impl CublasBackwardSet { /// ## Arguments /// /// - `d_value_logits` : `[B, NA]` gradient from C51 loss for the value head - /// - `d_adv_logits` : three `[B, branch_k * NA]` gradients for each branch + /// - `d_adv_logits` : four `[B, branch_k * NA]` gradients for each branch /// - `states` : `[B, SD]` raw input states (for last shared layer dW) /// - `save_h_s1` : `[B, SH1]` saved shared layer 1 ReLU output /// - `save_h_s2` : `[B, SH2]` saved shared layer 2 ReLU output /// - `save_h_v` : `[B, VH]` saved value FC ReLU output - /// - `save_h_b` : three `[B, AH]` saved branch FC ReLU outputs - /// - `w_ptrs` : 20-element array of raw F32 device pointers into + /// - `save_h_b` : four `[B, AH]` saved GLU output (post-gating activation) + /// - `w_ptrs` : 42-element array of raw F32 device pointers into /// `params_buf` (online weights, read-only) /// - `grad_buf` : flat gradient accumulator `[TOTAL_PARAMS]` — must be /// zeroed before this call /// - `scratch_dh` : scratch buffers for layer-to-layer gradient passing + /// - `glu_gate_pre` : four `[B, AH]` saved pre-sigmoid gate activations + /// - `glu_value` : four `[B, AH]` saved value path outputs (pre-gating) + /// - `scratch_d_glu_value`: `[B, AH]` scratch for GLU value gradient + /// - `scratch_d_glu_gate` : `[B, AH]` scratch for GLU gate_pre gradient #[allow(clippy::too_many_arguments)] pub fn backward_full( &self, @@ -707,13 +768,13 @@ impl CublasBackwardSet { save_h_s1: u64, // [B, SH1] save_h_s2: u64, // [B, SH2] save_h_v: u64, // [B, VH] - save_h_b: &[u64; 4], // [B, AH] each + save_h_b: &[u64; 4], // [B, AH] each — GLU output // Online weight pointers (read-only for W^T computation) w_ptrs: &[u64; super::gpu_dqn_trainer::NUM_WEIGHT_TENSORS], // Flat gradient accumulator (must be zeroed by caller) grad_buf_base: u64, // Scratch buffers for inter-layer gradients (f32) - scratch_d_h_s2: u64, // [B, SH2] — f32, accumulated from value + all 3 branches + scratch_d_h_s2: u64, // [B, SH2] — f32, accumulated from value + all 4 branches scratch_d_h_s1: u64, // [B, SH1] — f32 scratch_d_h_v: u64, // [B, VH] — f32 scratch_d_h_b: &[u64; 4], // [B, AH] each — f32 @@ -724,6 +785,12 @@ impl CublasBackwardSet { // Magnitude branch conditioning: saved forward concat and dX output mag_concat_ptr: u64, // [B, SH2+3] saved forward concat (for dW) d_mag_concat_ptr: u64, // [B, SH2+3] dX output for branch 1 (caller accumulates) + // GLU saved forward activations for backward pass + glu_gate_pre: &[u64; 4], // [B, AH] each — saved pre-sigmoid gate activations + glu_value: &[u64; 4], // [B, AH] each — saved value path outputs + // GLU backward scratch buffers (reused across branches) + scratch_d_glu_value: u64, // [B, AH] — d_value output from GLU backward + scratch_d_glu_gate: u64, // [B, AH] — d_gate_pre output from GLU backward ) -> Result<(), MLError> { let b = self.batch_size; let na = self.num_atoms; @@ -804,23 +871,42 @@ impl CublasBackwardSet { // [12]=w_b1fc, [13]=b_b1fc, [14]=w_b1out, [15]=b_b1out, // [16]=w_b2fc, [17]=b_b2fc, [18]=w_b2out, [19]=b_b2out, // [20]=w_b3fc, [21]=b_b3fc, [22]=w_b3out, [23]=b_b3out + // [34]=w_gate_0, [35]=b_gate_0, [36]=w_gate_1, [37]=b_gate_1, + // [38]=w_gate_2, [39]=b_gate_2, [40]=w_gate_3, [41]=b_gate_3 let w_bfc_idx = [8_usize, 12, 16, 20]; let w_bout_idx = [10_usize, 14, 18, 22]; + // GLU gate weight indices: w_gate_d at 34+d*2, b_gate_d at 34+d*2+1 + let w_gate_idx = [34_usize, 36, 38, 40]; + + // GLU gate gradient offsets (use padded_byte_offset for correctness) + let goff_w_gate = [ + padded_byte_offset(&self.param_sizes, 34), + padded_byte_offset(&self.param_sizes, 36), + padded_byte_offset(&self.param_sizes, 38), + padded_byte_offset(&self.param_sizes, 40), + ]; + let goff_b_gate = [ + padded_byte_offset(&self.param_sizes, 35), + padded_byte_offset(&self.param_sizes, 37), + padded_byte_offset(&self.param_sizes, 39), + padded_byte_offset(&self.param_sizes, 41), + ]; // ══════════════════════════════════════════════════════════════════ - // BRANCHES (all 4, in parallel with shared h_s2 accumulation) + // BRANCHES (all 4, with GLU backward replacing ReLU mask) // ══════════════════════════════════════════════════════════════════ for d in 0..4 { let n_d = branch_n[d]; let w_out = w_ptrs[w_bout_idx[d]]; // W_bdk_out[n_d, AH] - let w_fc = w_ptrs[w_bfc_idx[d]]; // W_bdk_fc[AH, SH2] + let w_fc = w_ptrs[w_bfc_idx[d]]; // W_bdk_fc[AH, SH2] (value path weight) + let w_gate = w_ptrs[w_gate_idx[d]]; // W_gate_d[AH, SH2] (gate path weight) - // ── Branch output layer (no ReLU — logit layer) ────────── + // ── Step 1: Branch output layer (no activation — logit layer) ── // dY = d_adv_logits[d] [B, n_d] - // X = save_h_b[d] [B, AH] + // X = save_h_b[d] [B, AH] (GLU output) // W = W_bdk_out [n_d, AH] - // dW += dY^T @ X, db += sum(dY), dX_b = dY @ W^T + // dW += dY^T @ X, db += sum(dY), dX_b = dY @ W^T → scratch_d_h_b[d] self.backward_fc_layer( stream, d_adv_logits[d], // dY [B, n_d] @@ -828,53 +914,28 @@ impl CublasBackwardSet { w_out, // W [n_d, AH] grad_buf_base + goff_w_bout[d], // dW grad_buf_base + goff_b_bout[d], // db - scratch_d_h_b[d], // dX [B, AH] + scratch_d_h_b[d], // dX [B, AH] → d_h_bd (gradient into GLU output) n_d, // out_dim self.adv_h, // in_dim b, )?; - // No ReLU mask needed — branch output is a logit layer. - // ── Branch FC layer (ReLU activation) ──────────────────── - // dY = scratch_d_h_b[d] [B, AH] - // X = save_h_s2 [B, SH2] - // W = W_bdk_fc [AH, SH2] - // dW += dY^T @ X, db += sum(dY), dX_s2 written to scratch_d_h_s2 - // - // NOTE: We accumulate contributions from all 3 branches into - // scratch_d_h_s2. To do this safely, we use beta=1.0 in the - // backward_fc_layer call. However, backward_fc_layer always uses - // beta=0.0 for dX. We therefore use a temporary per-branch dX - // and manually add it into scratch_d_h_s2. - // - // For simplicity, we write into scratch_d_h_b[d] (reuse — output - // layer has already consumed it) and then accumulate. - // - // Actually, scratch_d_h_b[d] is already holding the branch FC - // upstream gradient dX that we need to produce. We write it - // there, apply ReLU, then accumulate into scratch_d_h_s2. - // - // So we need a separate branch_dh_s2 scratch. For now, use - // scratch_d_h_b[d] as the per-branch dh_s2 temporary: after the - // branch out backward we no longer need d_h_b[d] for anything, - // so we can repurpose it. But its size is [B, AH], not [B, SH2]. - // We cannot reuse it directly. - // - // Solution: run the branch FC backward writing dX into - // scratch_d_h_s2 with beta=0 first, then for branches 1 and 2 - // we need to add rather than overwrite. We handle this by - // setting dx=0 for the fc layer and doing the upstream gradient - // via a separate GEMM call with beta=1.0 for branches d>0. - - // Apply ReLU mask to branch output dX: d_h_b[d] *= (save_h_b[d] > 0) - self.relu_mask( + // ── Step 2: GLU backward (replaces ReLU mask) ────────────── + // d_h_bd → d_value + d_gate_pre using saved gate_pre and value + // d_value[i] = d_h_bd[i] * sigmoid(gate_pre[i]) + // d_gate_pre[i] = d_h_bd[i] * value[i] * sigmoid(gate_pre[i]) * (1 - sigmoid(gate_pre[i])) + self.launch_glu_backward( stream, - scratch_d_h_b[d], // f32 dx to gate - save_h_b[d], // f32 saved post-ReLU activation + scratch_d_h_b[d], // d_output = d_h_bd [B, AH] + glu_gate_pre[d], // saved gate_pre [B, AH] + glu_value[d], // saved value [B, AH] + scratch_d_glu_gate, // d_gate_pre output [B, AH] + scratch_d_glu_value, // d_value output [B, AH] b * self.adv_h, )?; - // dW for branch FC: dW += dY^T @ X + // ── Step 3: Value path weight gradients ──────────────────── + // dW_bdf += d_value^T @ vsn_input, db_bdf += sum(d_value) // Branch 1 (magnitude) uses wider concat input [B, SH2+3] for dW. let (fc_input, fc_in_dim) = if d == 1 && mag_concat_ptr != 0 { (mag_concat_ptr, self.shared_h2 + 3) @@ -883,8 +944,8 @@ impl CublasBackwardSet { }; self.launch_dw_only( stream, - scratch_d_h_b[d], // dY [B, AH] — f32 - fc_input, // X [B, fc_in_dim] + scratch_d_glu_value, // dY = d_value [B, AH] — f32 + fc_input, // X [B, fc_in_dim] (vsn_input) grad_buf_base + goff_w_bfc[d], // dW [AH, fc_in_dim] grad_buf_base + goff_b_bfc[d], // db [AH] self.adv_h, // out_dim @@ -892,35 +953,72 @@ impl CublasBackwardSet { b, )?; - // dX for branch FC: d_h_s2 += dY @ W_bdk_fc - // Branch 1 (magnitude): dX writes to d_mag_concat [B, SH2+3]. - // The caller accumulates first SH2 columns into d_h_s2 via strided_accumulate. + // ── Step 4: Gate path weight gradients ───────────────────── + // dW_gate += d_gate_pre^T @ vsn_input, db_gate += sum(d_gate_pre) + self.launch_dw_only( + stream, + scratch_d_glu_gate, // dY = d_gate_pre [B, AH] — f32 + fc_input, // X [B, fc_in_dim] (same vsn_input) + grad_buf_base + goff_w_gate[d], // dW_gate [AH, fc_in_dim] + grad_buf_base + goff_b_gate[d], // db_gate [AH] + self.adv_h, // out_dim + fc_in_dim, // in_dim + b, + )?; + + // ── Step 5: Upstream gradient to VSN input ───────────────── + // d_vsn_input = W_bdf^T @ d_value + W_gate^T @ d_gate_pre + // Sum of both paths (value + gate) using beta accumulation. if d == 1 && mag_concat_ptr != 0 { - // Magnitude: dX = dY @ W_b1fc → [B, SH2+3], write to d_mag_concat + // Magnitude: dX writes to d_mag_concat [B, SH2+3]. + // Value path: beta=0 (overwrite) self.launch_dx_only( stream, - scratch_d_h_b[d], // dY [B, AH] — f32 - w_fc, // W [AH, SH2+3] - d_mag_concat_ptr, // dX [B, SH2+3] — f32 - self.adv_h, // out_dim - self.shared_h2 + 3, // in_dim + scratch_d_glu_value, // dY = d_value [B, AH] + w_fc, // W_bdf [AH, SH2+3] + d_mag_concat_ptr, // dX [B, SH2+3] + self.adv_h, + self.shared_h2 + 3, b, - 0.0_f32, // overwrite (fresh buffer) + 0.0_f32, // overwrite + )?; + // Gate path: beta=1 (accumulate on top of value path) + self.launch_dx_only( + stream, + scratch_d_glu_gate, // dY = d_gate_pre [B, AH] + w_gate, // W_gate [AH, SH2+3] + d_mag_concat_ptr, // dX [B, SH2+3] (accumulate) + self.adv_h, + self.shared_h2 + 3, + b, + 1.0_f32, // accumulate )?; // Caller will call accumulate_d_h_s2_from_concat(d_mag_concat, d_h_s2, B, 1.0) // after backward_full returns. } else { // Other branches: accumulate into d_h_s2 directly. let beta_s2 = if d == 0 { 0.0_f32 } else { 1.0_f32 }; + // Value path self.launch_dx_only( stream, - scratch_d_h_b[d], // dY [B, AH] — f32 - w_fc, // W [AH, SH2] - scratch_d_h_s2, // dX [B, SH2] — f32 - self.adv_h, // out_dim - self.shared_h2, // in_dim + scratch_d_glu_value, // dY = d_value [B, AH] + w_fc, // W_bdf [AH, SH2] + scratch_d_h_s2, // dX [B, SH2] + self.adv_h, + self.shared_h2, b, - beta_s2, + beta_s2, // d==0: overwrite, d>0: accumulate + )?; + // Gate path: always beta=1 (accumulate on top of value path) + self.launch_dx_only( + stream, + scratch_d_glu_gate, // dY = d_gate_pre [B, AH] + w_gate, // W_gate [AH, SH2] + scratch_d_h_s2, // dX [B, SH2] (accumulate) + self.adv_h, + self.shared_h2, + b, + 1.0_f32, // always accumulate )?; } } @@ -1305,6 +1403,8 @@ pub fn alloc_backward_scratch( CudaSlice, // d_h_b1 [B, AH] CudaSlice, // d_h_b2 [B, AH] CudaSlice, // d_h_b3 [B, AH] + CudaSlice, // d_glu_value [B, AH] — GLU value path gradient scratch + CudaSlice, // d_glu_gate [B, AH] — GLU gate_pre gradient scratch ), MLError> { let b = config.batch_size; let alloc_f32 = |n: usize| -> Result, MLError> { @@ -1320,5 +1420,7 @@ pub fn alloc_backward_scratch( alloc_f32(b * config.adv_h)?, alloc_f32(b * config.adv_h)?, alloc_f32(b * config.adv_h)?, + alloc_f32(b * config.adv_h)?, // GLU d_value scratch + alloc_f32(b * config.adv_h)?, // GLU d_gate_pre scratch )) } diff --git a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs index ef50eba16..494f5ab86 100644 --- a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs +++ b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs @@ -485,6 +485,8 @@ struct CachedPtrs { bw_d_h_b1: u64, bw_d_h_b2: u64, bw_d_h_b3: u64, + bw_d_glu_value: u64, + bw_d_glu_gate: u64, iqn_trunk_m: u64, iqn_trunk_grad_norm: u64, td_errors_buf: u64, @@ -968,6 +970,10 @@ pub struct GpuDqnTrainer { bw_d_h_b1: CudaSlice, bw_d_h_b2: CudaSlice, bw_d_h_b3: CudaSlice, + /// GLU backward scratch: d_value [B, AH] — f32 (reused across branches) + bw_d_glu_value: CudaSlice, + /// GLU backward scratch: d_gate_pre [B, AH] — f32 (reused across branches) + bw_d_glu_gate: CudaSlice, // ── Expected Q-value kernel (ad-hoc validation, not captured in CUDA Graph) ─ /// Converts C51 value+advantage logits → expected Q-values (validation path). /// Input: on_v_logits_buf [B, NA], on_b_logits_buf [B, (B0+B1+B2+B3)*NA] @@ -1997,6 +2003,20 @@ impl GpuDqnTrainer { self.stream.memset_zeros(&mut self.cql_grad_scratch) .map_err(|e| MLError::ModelError(format!("zero cql_grad_scratch: {e}")))?; + // GLU saved activation pointers + let glu_gate_pre_ptrs = [ + self.glu_gate_pre_buf[0].raw_ptr(), + self.glu_gate_pre_buf[1].raw_ptr(), + self.glu_gate_pre_buf[2].raw_ptr(), + self.glu_gate_pre_buf[3].raw_ptr(), + ]; + let glu_value_ptrs = [ + self.glu_value_buf[0].raw_ptr(), + self.glu_value_buf[1].raw_ptr(), + self.glu_value_buf[2].raw_ptr(), + self.glu_value_buf[3].raw_ptr(), + ]; + // Run full backward pass with CQL logit gradients into ISOLATED scratch buffer. // Produces CQL parameter gradients WITHOUT mixing with C51's grad_buf. self.cublas_backward.backward_full( @@ -2013,6 +2033,10 @@ impl GpuDqnTrainer { 0, // CQL backward: no bottleneck dX needed (separate gradient budget) self.ptrs.mag_concat_buf, self.ptrs.d_mag_concat_buf, + &glu_gate_pre_ptrs, + &glu_value_ptrs, + self.ptrs.bw_d_glu_value, + self.ptrs.bw_d_glu_gate, ).map_err(|e| MLError::ModelError(format!("CQL backward_full: {e}")))?; // CQL: accumulate magnitude branch dX into scratch_d_h_s2 @@ -2963,13 +2987,16 @@ impl GpuDqnTrainer { info!("GpuDqnTrainer: cuBLAS batched forward (DDQN) initialized (VSN+GLU wired)"); // ── Initialize cuBLAS backward context (required) ────────── - let cublas_backward = CublasBackwardSet::new(Arc::clone(&shared_cublas), &config)?; - info!("GpuDqnTrainer: cuBLAS batched backward initialized"); + let cublas_backward = CublasBackwardSet::new( + Arc::clone(&shared_cublas), &config, glu_backward_kernel.clone(), + )?; + info!("GpuDqnTrainer: cuBLAS batched backward initialized (GLU backward wired)"); // ── Backward scratch buffers ──────────────────────────────── // Pre-allocate inter-layer gradient buffers for the cuBLAS backward // pass. These are separate from the activation saves used in forward. - let (bw_d_h_s2, bw_d_h_s1, bw_d_h_v, bw_d_h_b0, bw_d_h_b1, bw_d_h_b2, bw_d_h_b3) = + let (bw_d_h_s2, bw_d_h_s1, bw_d_h_v, bw_d_h_b0, bw_d_h_b1, bw_d_h_b2, bw_d_h_b3, + bw_d_glu_value, bw_d_glu_gate) = alloc_backward_scratch(&stream, &config) .map_err(|e| MLError::ModelError(format!("backward scratch alloc: {e}")))?; @@ -3061,6 +3088,8 @@ impl GpuDqnTrainer { bw_d_h_b1: bw_d_h_b1.raw_ptr(), bw_d_h_b2: bw_d_h_b2.raw_ptr(), bw_d_h_b3: bw_d_h_b3.raw_ptr(), + bw_d_glu_value: bw_d_glu_value.raw_ptr(), + bw_d_glu_gate: bw_d_glu_gate.raw_ptr(), iqn_trunk_m: iqn_trunk_m.raw_ptr(), iqn_trunk_grad_norm: iqn_trunk_grad_norm.raw_ptr(), td_errors_buf: td_errors_buf.raw_ptr(), @@ -3338,6 +3367,8 @@ impl GpuDqnTrainer { bw_d_h_b1, bw_d_h_b2, bw_d_h_b3, + bw_d_glu_value, + bw_d_glu_gate, expected_q_kernel, q_stats_kernel, q_stats_buf, @@ -5873,6 +5904,20 @@ impl GpuDqnTrainer { 0u64 // no dX needed — states not trainable }; + // GLU saved activation pointers + let glu_gate_pre_ptrs = [ + self.glu_gate_pre_buf[0].raw_ptr(), + self.glu_gate_pre_buf[1].raw_ptr(), + self.glu_gate_pre_buf[2].raw_ptr(), + self.glu_gate_pre_buf[3].raw_ptr(), + ]; + let glu_value_ptrs = [ + self.glu_value_buf[0].raw_ptr(), + self.glu_value_buf[1].raw_ptr(), + self.glu_value_buf[2].raw_ptr(), + self.glu_value_buf[3].raw_ptr(), + ]; + bw.backward_full( &self.stream, d_value_logits_ptr, @@ -5891,6 +5936,10 @@ impl GpuDqnTrainer { s1_dx_output, self.ptrs.mag_concat_buf, self.ptrs.d_mag_concat_buf, + &glu_gate_pre_ptrs, + &glu_value_ptrs, + self.ptrs.bw_d_glu_value, + self.ptrs.bw_d_glu_gate, )?; // Accumulate magnitude branch dX (first SH2 columns) into d_h_s2