diff --git a/crates/ml/src/cuda_pipeline/experience_kernels.cu b/crates/ml/src/cuda_pipeline/experience_kernels.cu index abd9b2dc9..02592bcf9 100644 --- a/crates/ml/src/cuda_pipeline/experience_kernels.cu +++ b/crates/ml/src/cuda_pipeline/experience_kernels.cu @@ -5065,6 +5065,70 @@ extern "C" __global__ void predictive_coding_loss( } } +/** + * predictive_coding_backward — gradient of the predictive_coding_loss above. + * + * Loss shape: + * L = sum_{i=0..N-2} lambda_pred * (1/SH2) * sum_j (h[i,j] - h[i+1,j])^2 + * + * Gradient w.r.t. h[i,j]: + * dL/dh[i,j] = (2*lambda_pred/SH2) * sum over loss terms that involve h[i,j] + * + * Each h[i,j] appears in TWO loss terms (except boundaries): + * - as "current" in term i: contribution +diff(i,j) = (h[i,j] - h[i+1,j]) + * - as "next" in term i-1: contribution -diff(i-1,j) = -(h[i-1,j] - h[i,j]) + * = (h[i,j] - h[i-1,j]) + * + * Combined for interior i in [1, N-2]: + * dL/dh[i,j] = (2*lambda/SH2) * ( (h[i,j] - h[i+1,j]) + (h[i,j] - h[i-1,j]) ) + * = (2*lambda/SH2) * ( 2*h[i,j] - h[i-1,j] - h[i+1,j] ) + * + * Boundary i=0 (only "current" term): dL/dh[0,j] = (2*lambda/SH2) * (h[0,j] - h[1,j]) + * Boundary i=N-1 (only "next" term): dL/dh[N-1,j] = (2*lambda/SH2) * (h[N-1,j] - h[N-2,j]) + * + * Grid: ceil(N*SH2 / 256), Block: 256. One thread per (i, j) cell. + * Output: ADDS gradient into bw_d_h_s2[i, j] (does not overwrite — caller's + * trunk backward path expects accumulation semantics). + * + * Determinism: each thread writes to a unique (i, j) slot — plain += is safe, + * no atomicAdd needed. + */ +extern "C" __global__ void predictive_coding_backward( + const float* __restrict__ h_s2, /* [N, SH2] same buffer as forward */ + float* __restrict__ d_h_s2, /* [N, SH2] accumulated trunk gradient (in/out) */ + int N, + int sh2, + float lambda_pred /* must match forward */ +) { + int idx = blockIdx.x * blockDim.x + threadIdx.x; + int total = N * sh2; + if (idx >= total) return; + + int i = idx / sh2; + int j = idx % sh2; + + /* Pre-factor: 2 * lambda / SH2 */ + float k = (2.0f * lambda_pred) / (float)sh2; + long long off = (long long)i * sh2 + j; + + float v_self = h_s2[off]; + float grad = 0.0f; + + /* "current" term i: diff = (h[i] - h[i+1]) */ + if (i < N - 1) { + float v_next = h_s2[(long long)(i + 1) * sh2 + j]; + grad += k * (v_self - v_next); + } + /* "next" term i-1: -diff(i-1) = (h[i] - h[i-1]) */ + if (i > 0) { + float v_prev = h_s2[(long long)(i - 1) * sh2 + j]; + grad += k * (v_self - v_prev); + } + + /* Accumulate (caller's existing trunk backward expects += semantics) */ + d_h_s2[off] += grad; +} + /* ================================================================== */ /* Kernel: homeostatic_regularizer (G16) */ /* ================================================================== */ diff --git a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs index f7eccc859..273b0ca69 100644 --- a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs +++ b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs @@ -918,7 +918,10 @@ pub struct GpuDqnTrainer { /// Scalar output buffer [1] for temporal consistency penalty. temporal_penalty_buf: CudaSlice, - // ── G12: Predictive coding auxiliary loss ── + // ── G12: Predictive coding auxiliary loss + backward ── + /// Backward kernel: writes per-(sample, feature) gradient of the + /// predictive_coding_loss into bw_d_h_s2 via plain += (no atomicAdd). + predictive_coding_backward_kernel: CudaFunction, /// Kernel: predictive_coding_loss — self-supervised temporal smoothness on enriched trunk. predictive_coding_kernel: CudaFunction, /// Per-sample loss buffer [B] for predictive coding — reduced by c51_loss_reduce. @@ -3374,28 +3377,40 @@ impl GpuDqnTrainer { Ok(()) } - /// G12: Predictive coding auxiliary loss — temporal smoothness on enriched trunk. - /// MSE between consecutive h_s2 samples provides a noise-free self-supervised - /// signal for the trunk. Uses save_h_s2 (enriched after mamba2_step). + /// G12: Predictive coding auxiliary loss + backward — temporal smoothness + /// on the enriched trunk. MSE between consecutive h_s2 samples gives a + /// noise-free self-supervised signal that flows back through the trunk. /// - /// Graph-safe: writes per-sample loss to predictive_per_sample_buf, then - /// reduces via c51_loss_reduce kernel. No cuMemsetD32Async or atomicAdd. + /// Three steps, all graph-safe (no memset, no atomicAdd): + /// 1. Forward: per-sample MSE → predictive_per_sample_buf [B] + /// 2. Reduce → predictive_loss_buf [1] (for monitoring readback) + /// 3. Backward: gradient on h_s2 → ACCUMULATED into bw_d_h_s2 + /// + /// The trunk backward (W_s2 → W_s1) reads bw_d_h_s2, so steps 1+3 effectively + /// add a regularization term to the trunk gradient. lambda_pred = 0.1 keeps + /// the predictive contribution small relative to the C51/IQN gradient. + /// + /// Caller MUST invoke this AFTER the main heads have written their share + /// of bw_d_h_s2 (so we accumulate, not overwrite) and BEFORE the trunk + /// backward GEMMs (so the trunk picks up the combined gradient). pub(crate) fn compute_predictive_coding_loss(&self, batch_size: usize) -> Result<(), MLError> { let sh2 = self.config.shared_h2 as i32; let save_h_s2_ptr = self.save_h_s2.raw_ptr(); let predictive_per_sample_buf_ptr = self.predictive_per_sample_buf.raw_ptr(); - // Step 1: per-sample MSE loss (graph-captured kernel, no memset needed) + let lambda_pred = 0.1_f32; + + // Step 1: per-sample MSE loss unsafe { self.stream.launch_builder(&self.predictive_coding_kernel) .arg(&save_h_s2_ptr) .arg(&predictive_per_sample_buf_ptr) .arg(&(batch_size as i32)) .arg(&sh2) - .arg(&0.1_f32) // lambda_pred + .arg(&lambda_pred) .launch(LaunchConfig::for_num_elems(batch_size as u32)) .map_err(|e| MLError::ModelError(format!("predictive_coding: {e}")))?; } - // Step 2: deterministic reduce → predictive_loss_buf[0] + // Step 2: deterministic reduce → predictive_loss_buf[0] (monitoring) let loss_ptr = self.predictive_loss_buf.raw_ptr(); unsafe { self.stream @@ -3410,6 +3425,20 @@ impl GpuDqnTrainer { }) .map_err(|e| MLError::ModelError(format!("predictive_coding reduce: {e}")))?; } + // Step 3: backward → ACCUMULATE gradient into bw_d_h_s2 (one thread + // per (sample, feature) cell, no atomic — unique writes). + let bw_d_h_s2_ptr = self.bw_d_h_s2.raw_ptr(); + let total = (batch_size * (sh2 as usize)) as u32; + unsafe { + self.stream.launch_builder(&self.predictive_coding_backward_kernel) + .arg(&save_h_s2_ptr) + .arg(&bw_d_h_s2_ptr) + .arg(&(batch_size as i32)) + .arg(&sh2) + .arg(&lambda_pred) + .launch(LaunchConfig::for_num_elems(total)) + .map_err(|e| MLError::ModelError(format!("predictive_coding_backward: {e}")))?; + } Ok(()) } @@ -5751,6 +5780,8 @@ impl GpuDqnTrainer { .map_err(|e| MLError::ModelError(format!("branch_independence_penalty load: {e}")))?; let temporal_consistency_kernel = exp_module_for_mag.load_function("temporal_consistency_penalty") .map_err(|e| MLError::ModelError(format!("temporal_consistency_penalty load: {e}")))?; + let predictive_coding_backward_kernel = exp_module_for_mag.load_function("predictive_coding_backward") + .map_err(|e| MLError::ModelError(format!("predictive_coding_backward load: {e}")))?; let predictive_coding_kernel = exp_module_for_mag.load_function("predictive_coding_loss") .map_err(|e| MLError::ModelError(format!("predictive_coding_loss load: {e}")))?; let risk_forward_kernel = exp_module_for_mag.load_function("risk_budget_forward") @@ -7343,6 +7374,7 @@ impl GpuDqnTrainer { temporal_consistency_kernel, temporal_penalty_buf, predictive_coding_kernel, + predictive_coding_backward_kernel, predictive_per_sample_buf, predictive_loss_buf, q_anchor_kernel, diff --git a/crates/ml/src/trainers/dqn/fused_training.rs b/crates/ml/src/trainers/dqn/fused_training.rs index 1a1debba1..d8d7799d2 100644 --- a/crates/ml/src/trainers/dqn/fused_training.rs +++ b/crates/ml/src/trainers/dqn/fused_training.rs @@ -1457,6 +1457,15 @@ impl FusedTrainingCtx { self.trainer.launch_recursive_confidence_backward(self.batch_size) .map_err(|e| anyhow::anyhow!("Recursive confidence backward: {e}"))?; + // G12: Predictive coding auxiliary loss + backward. + // Self-supervised temporal smoothness on the enriched trunk h_s2 — + // adds 2*lambda*(h[i] - h[i+1])/SH2 regularization gradient into + // bw_d_h_s2. Runs HERE so the gradient lands in bw_d_h_s2 before + // the trunk backward GEMMs consume it (and before recursive + // confidence's own trunk gradient — both accumulate via plain +=). + self.trainer.compute_predictive_coding_loss(self.batch_size) + .map_err(|e| anyhow::anyhow!("Predictive coding (G12): {e}"))?; + // Ensemble diversity: heads 1..K-1 forward + KL gradient + trunk backward. // Runs inside aux_child graph — gradients land in grad_buf before Adam. if !self.ensemble_extra_heads.is_empty() {