From e72885e8b6fa63bbb319226f8760c49b93c2fdbd Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Tue, 21 Apr 2026 09:23:28 +0200 Subject: [PATCH] =?UTF-8?q?gem(G12):=20finish=20predictive=5Fcoding=20?= =?UTF-8?q?=E2=80=94=20write=20backward=20+=20wire=20into=20training=20loo?= =?UTF-8?q?p?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Found by V7-gem audit: predictive_coding_loss kernel existed (forward only, per-sample MSE on consecutive trunk activations), and compute_predictive_coding_loss Rust wrapper existed, but no backward and no caller. Pure scaffolding. Self-supervised temporal smoothness on the enriched trunk h_s2 IS in the "better form" by the V7 methodology — it operates at the gradient level on the network representation, not as a reward shaping term. Worth wiring up. Three changes: 1) New CUDA kernel `predictive_coding_backward` (experience_kernels.cu) Loss: L = sum_{i=0..N-2} (lambda/SH2) * sum_j (h[i,j] - h[i+1,j])^2 Grad: dL/dh[i,j] = (2*lambda/SH2) * sum-of-affected-loss-terms Each h[i,j] appears in TWO loss terms (interior i): "current" of term i and "next" of term i-1. Boundaries (i=0, i=N-1) have one. Closed-form gradient written directly with no atomicAdd — one thread per (sample, feature) cell, each writes to a unique slot, plain += accumulates into bw_d_h_s2. Bit-deterministic. 2) Rust glue (gpu_dqn_trainer.rs) - Load `predictive_coding_backward` kernel via existing exp_module_for_mag - New field on DQNTrainer (predictive_coding_backward_kernel) - Extend `compute_predictive_coding_loss` to also launch backward as step 3 (after forward + reduce). Now the function name accurately describes what it does — both compute and accumulate gradient. 3) Integration (fused_training.rs::submit_aux_ops) Inserted the call right after `launch_recursive_confidence_backward`, before regime_scale_td_errors. Both spots accumulate into bw_d_h_s2 via plain +=, so ordering is irrelevant for correctness — what matters is that this runs INSIDE the aux_child CUDA-graph capture window AND BEFORE the trunk W_s2 → W_s1 backward GEMMs read bw_d_h_s2. Why not also G6 (branch_independence) and G10 (temporal_consistency)? Per V7-gem methodology — check for redundancy first: - G10 wants Lipschitz on Q for similar states. Spectral normalization (already wired on all 12 weight tensors) achieves *global* Lipschitz. G10 adds *local-pair* Lipschitz on top. Possibly redundant — needs a measurement before wiring blindly. - G6 wants the 4 advantage branches to stay diverse. NoisyNets already adds different parameter noise per layer → naturally diverse heads. Ensemble KL gradient pushes ensemble heads apart (different mechanism but similar intent). Possibly redundant for the 4-branch case. G12 is unambiguously useful — trunk smoothness is a genuine gem with no existing equivalent in the codebase. G6/G10 land separately if measurement shows they add real value beyond spectral_norm + NoisyNets + ensemble KL. Files touched: crates/ml/src/cuda_pipeline/experience_kernels.cu (+45) crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs (+33 / -3) crates/ml/src/trainers/dqn/fused_training.rs (+9) Verified: cargo check passes. Smoke test running to verify training is stable with G12 active (lambda_pred=0.1 — small enough that any regression is from a real bug, not dominance over the C51/IQN gradient). Co-Authored-By: Claude Opus 4.7 (1M context) --- .../src/cuda_pipeline/experience_kernels.cu | 64 +++++++++++++++++++ .../ml/src/cuda_pipeline/gpu_dqn_trainer.rs | 50 ++++++++++++--- crates/ml/src/trainers/dqn/fused_training.rs | 9 +++ 3 files changed, 114 insertions(+), 9 deletions(-) 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() {