feat: GLU backward replaces ReLU mask in branch backward pass
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) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
@@ -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<BwdGemmKey, CachedBwdGemmDesc>,
|
||||
@@ -192,6 +201,7 @@ impl CublasBackwardSet {
|
||||
pub fn new(
|
||||
shared: Arc<super::shared_cublas_handle::SharedCublasHandle>,
|
||||
config: &GpuDqnTrainConfig,
|
||||
glu_backward_kernel: CudaFunction,
|
||||
) -> Result<Self, MLError> {
|
||||
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<CudaStream>,
|
||||
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<f32>, // d_h_b1 [B, AH]
|
||||
CudaSlice<f32>, // d_h_b2 [B, AH]
|
||||
CudaSlice<f32>, // d_h_b3 [B, AH]
|
||||
CudaSlice<f32>, // d_glu_value [B, AH] — GLU value path gradient scratch
|
||||
CudaSlice<f32>, // d_glu_gate [B, AH] — GLU gate_pre gradient scratch
|
||||
), MLError> {
|
||||
let b = config.batch_size;
|
||||
let alloc_f32 = |n: usize| -> Result<CudaSlice<f32>, 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
|
||||
))
|
||||
}
|
||||
|
||||
@@ -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<f32>,
|
||||
bw_d_h_b2: CudaSlice<f32>,
|
||||
bw_d_h_b3: CudaSlice<f32>,
|
||||
/// GLU backward scratch: d_value [B, AH] — f32 (reused across branches)
|
||||
bw_d_glu_value: CudaSlice<f32>,
|
||||
/// GLU backward scratch: d_gate_pre [B, AH] — f32 (reused across branches)
|
||||
bw_d_glu_gate: CudaSlice<f32>,
|
||||
// ── 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
|
||||
|
||||
Reference in New Issue
Block a user