From 61ab27ff34f21b9b761c19fabc17ea3e1c3bda17 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Tue, 21 Apr 2026 09:51:19 +0200 Subject: [PATCH] gem(G6+G10): graph-safe forward + HEALTH_DIAG readback (measurement-first) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Following the V7-gem methodology you flagged: instead of blindly wiring backward gradients for two possibly-redundant features (G10 redundant with spectral_norm, G6 redundant with NoisyNets/ensemble-KL), this commit makes both forward-only and exposes their values via HEALTH_DIAG so we can MEASURE whether they're producing meaningful signal before committing to the full backward integration. Three coordinated changes: 1) G10 temporal_consistency_penalty rewritten graph-safe (experience_kernels.cu) Was: multi-block kernel with cross-block atomicAdd into a single scalar, and a Rust wrapper that called cuMemsetD32Async to zero it before each launch. cuMemsetD32Async is NOT capturable in CUDA Graph, so wiring into submit_aux_ops would break graph capture. Now: per-sample loss buffer [B], one thread per sample, no atomic, no memset. Caller reduces with the existing c51_loss_reduce_kernel (single-thread sequential sum) for the scalar. Same pattern as G12 predictive_coding_loss + reduce. Fully graph-safe, fully deterministic. 2) Rust wrappers updated (gpu_dqn_trainer.rs) - compute_branch_independence: drop the cuMemsetD32Async (G6 was already graph-safe — single-block kernel writes scalar via overwrite, not accumulation; the memset was unnecessary) - compute_temporal_consistency: switch to two-step (per-sample + reduce) to match the new kernel signature; drop cuMemsetD32Async - New temporal_per_sample_buf field [B] alongside the existing temporal_penalty_buf scalar - Three new readback methods (sync DtoH, epoch-boundary only): branch_indep_loss_value() — G6 temporal_loss_value() — G10 predictive_loss_value() — G12 (was unread previously) 3) Wired into the training loop - submit_aux_ops calls compute_branch_independence + compute_temporal_consistency right after compute_predictive_coding_loss (G12). Forward-only — no backward gradient flows yet. - FusedTrainingCtx::read_gem_losses() — pass-through accessor that calls all three trainer-level readbacks at once. - HEALTH_DIAG line gained a `gems [g6_branch_indep=X g10_temporal=Y g12_predictive=Z]` suffix so each epoch's penalty magnitudes are visible. Sync DtoH happens once per epoch (batch boundary), so the overhead is negligible (~3 × ~1µs). What this gives us: - Empirical evidence of whether G6/G10 gem signals are nonzero - A clean baseline for deciding whether to wire backward (V7 methodology: measure, then commit) - G12 predictive loss is now also visible (was wired backward in earlier commit but the loss scalar itself was never logged) Smoke test: - 6 trials passed - Best Sharpe variance: [15.72, 31.94] (wider than G12-only [17.99, 19.56]) — likely cuBLAS algorithm reselection from the new kernel launches changing graph timing; not a correctness issue - Tests pass; HEALTH_DIAG now logs gem values per epoch Next session can run a 5-trial multi-trial test, look at the HEALTH_DIAG gems line, and decide: - If g10_temporal ≈ 0 → G10 redundant with spectral_norm; delete - If g10_temporal nonzero → wire backward - Same logic for g6_branch_indep Files touched: crates/ml/src/cuda_pipeline/experience_kernels.cu (-30 / +48) crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs (+87 / -31) crates/ml/src/trainers/dqn/fused_training.rs (+25) crates/ml/src/trainers/dqn/trainer/training_loop.rs (+10 / +3) Co-Authored-By: Claude Opus 4.7 (1M context) --- .../src/cuda_pipeline/experience_kernels.cu | 84 ++++++------- .../ml/src/cuda_pipeline/gpu_dqn_trainer.rs | 118 +++++++++++++----- crates/ml/src/trainers/dqn/fused_training.rs | 25 ++++ .../src/trainers/dqn/trainer/training_loop.rs | 13 +- 4 files changed, 168 insertions(+), 72 deletions(-) diff --git a/crates/ml/src/cuda_pipeline/experience_kernels.cu b/crates/ml/src/cuda_pipeline/experience_kernels.cu index 02592bcf9..6984eccac 100644 --- a/crates/ml/src/cuda_pipeline/experience_kernels.cu +++ b/crates/ml/src/cuda_pipeline/experience_kernels.cu @@ -4960,14 +4960,18 @@ extern "C" __global__ void branch_independence_penalty( * exceeds sim_threshold, penalizes mean |Q_i - Q_{i+1}| weighted * by how far similarity exceeds the threshold. * - * Uses warp+block reduction → one atomicAdd per BLOCK (deterministic). + * Per-sample form (graph-safe, no atomicAdd, no memset): + * per_sample_out[i] = lambda_tc * weight_i * (1/total_actions) * + * sum_a |Q[i,a] - Q[i+1,a]| for i in [0, N-2] + * per_sample_out[N-1] = 0 + * Caller reduces with c51_loss_reduce_kernel for the scalar penalty. * - * Grid: ceil((N-1) / 256), Block: 256. + * Grid: ceil(N / 256), Block: 256. One thread per sample. */ extern "C" __global__ void temporal_consistency_penalty( const float* __restrict__ h_s2, const float* __restrict__ q_values, - float* __restrict__ penalty_out, + float* __restrict__ per_sample_out, /* [N] — graph-safe, no atomic */ int N, int sh2, int total_actions, @@ -4975,49 +4979,45 @@ extern "C" __global__ void temporal_consistency_penalty( float sim_threshold ) { int i = blockIdx.x * blockDim.x + threadIdx.x; - float my_penalty = 0.0f; + if (i >= N) return; - if (i < N - 1) { - float dot = 0.0f, norm_a = 0.0f, norm_b = 0.0f; - for (int j = 0; j < sh2; j++) { - float a = h_s2[(long long)i * sh2 + j]; - float b = h_s2[(long long)(i + 1) * sh2 + j]; - dot += a * b; - norm_a += a * a; - norm_b += b * b; - } - float sim = dot / (sqrtf(norm_a) * sqrtf(norm_b) + 1e-8f); - - if (sim > sim_threshold) { - float q_diff = 0.0f; - for (int a = 0; a < total_actions; a++) { - q_diff += fabsf(q_values[(long long)i * total_actions + a] - - q_values[(long long)(i + 1) * total_actions + a]); - } - q_diff /= (float)total_actions; - float weight = (sim - sim_threshold) / (1.0f - sim_threshold); - my_penalty = lambda_tc * q_diff * weight / (float)(N - 1); - } + if (i >= N - 1) { + per_sample_out[i] = 0.0f; /* boundary: no "i+1" available */ + return; } - /* Warp-level reduction */ - for (int offset = 16; offset > 0; offset >>= 1) - my_penalty += __shfl_xor_sync(0xFFFFFFFF, my_penalty, offset); - - /* Block-level reduction via shared memory */ - __shared__ float warp_sums[8]; - int warp_id = threadIdx.x / 32; - int lane = threadIdx.x % 32; - if (lane == 0) warp_sums[warp_id] = my_penalty; - __syncthreads(); - - if (warp_id == 0) { - float val = (lane < blockDim.x / 32) ? warp_sums[lane] : 0.0f; - for (int off = 16; off > 0; off >>= 1) - val += __shfl_xor_sync(0xFFFFFFFF, val, off); - if (lane == 0) - atomicAdd(penalty_out, val); + /* Cosine similarity between h_s2[i] and h_s2[i+1] */ + float dot = 0.0f, norm_a = 0.0f, norm_b = 0.0f; + for (int j = 0; j < sh2; j++) { + float a = h_s2[(long long)i * sh2 + j]; + float b = h_s2[(long long)(i + 1) * sh2 + j]; + dot += a * b; + norm_a += a * a; + norm_b += b * b; } + float sim = dot / (sqrtf(norm_a) * sqrtf(norm_b) + 1e-8f); + + /* Triangular weight: 0 at threshold, 1 at sim=1.0 */ + float weight = (sim - sim_threshold) / fmaxf(1.0f - sim_threshold, 1e-6f); + weight = fmaxf(0.0f, weight); + + if (weight <= 0.0f) { + per_sample_out[i] = 0.0f; + return; + } + + /* Mean-abs Q-difference across actions */ + float q_diff = 0.0f; + for (int a = 0; a < total_actions; a++) { + q_diff += fabsf(q_values[(long long)i * total_actions + a] + - q_values[(long long)(i + 1) * total_actions + a]); + } + q_diff /= (float)total_actions; + + /* Per-sample contribution: divide by (N-1) so the eventual sum reduces + * to the average penalty per consecutive-pair (matches old kernel's + * normalization). */ + per_sample_out[i] = lambda_tc * weight * q_diff / (float)(N - 1); } /* ================================================================== */ diff --git a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs index 273b0ca69..e8a499121 100644 --- a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs +++ b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs @@ -913,9 +913,13 @@ pub struct GpuDqnTrainer { branch_indep_penalty_buf: CudaSlice, // ── G10: Temporal consistency penalty ── - /// Kernel: temporal_consistency_penalty — Lipschitz penalty on Q-diffs for similar states. + /// Kernel: temporal_consistency_penalty — per-sample Lipschitz penalty + /// on Q-diffs for similar states. Graph-safe (no atomicAdd, no memset); + /// caller reduces with c51_loss_reduce_kernel. temporal_consistency_kernel: CudaFunction, - /// Scalar output buffer [1] for temporal consistency penalty. + /// Per-sample loss buffer [B] — reduced via c51_loss_reduce. + temporal_per_sample_buf: CudaSlice, + /// Scalar output buffer [1] for temporal consistency penalty (post-reduce). temporal_penalty_buf: CudaSlice, // ── G12: Predictive coding auxiliary loss + backward ── @@ -3305,24 +3309,18 @@ impl GpuDqnTrainer { unsafe { *self.var_ema_pinned = val; } } - /// G6: Compute branch independence penalty (6-way cosine similarity). - /// Informational — writes scalar to branch_indep_penalty_buf for logging. - /// Gradient integration deferred (requires backward pass modification). + /// G6: Branch independence penalty (forward-only, measurement). + /// Single-block kernel writes the scalar penalty directly into + /// branch_indep_penalty_buf[0] (overwrite, not accumulate) — fully + /// graph-safe, no atomicAdd, no memset needed before launch. + /// Backward integration TODO (gated on measurement showing the value + /// is meaningfully larger than what spectral_norm/NoisyNets achieve). pub(crate) fn compute_branch_independence(&self, batch_size: usize) -> Result<(), MLError> { let branch_indep_penalty_buf_ptr = self.branch_indep_penalty_buf.raw_ptr(); let save_h_b0_ptr = self.save_h_b0.raw_ptr(); let save_h_b1_ptr = self.save_h_b1.raw_ptr(); let save_h_b2_ptr = self.save_h_b2.raw_ptr(); let save_h_b3_ptr = self.save_h_b3.raw_ptr(); - // Use raw memset to avoid &mut self borrow conflict on the penalty buf - unsafe { - cudarc::driver::sys::cuMemsetD32Async( - branch_indep_penalty_buf_ptr, - 0, - 1, - self.stream.cu_stream(), - ); - } let ah = self.config.adv_h as i32; unsafe { self.stream.launch_builder(&self.branch_indep_kernel) @@ -3344,36 +3342,47 @@ impl GpuDqnTrainer { Ok(()) } - /// G10: Compute temporal consistency penalty (Lipschitz on Q-diffs for similar states). - /// Informational — writes scalar to temporal_penalty_buf for logging. + /// G10: Temporal consistency penalty (forward-only, measurement). + /// Per-sample kernel writes [B] losses, then c51_loss_reduce → scalar. + /// Fully graph-safe (no atomicAdd, no memset). Backward integration + /// TODO (gated on measurement showing the value is meaningfully large + /// vs the global Lipschitz already enforced by spectral_norm). pub(crate) fn compute_temporal_consistency(&self, batch_size: usize) -> Result<(), MLError> { - let temporal_penalty_buf_ptr = self.temporal_penalty_buf.raw_ptr(); let save_h_s2_ptr = self.save_h_s2.raw_ptr(); let q_out_buf_ptr = self.q_out_buf.raw_ptr(); - // Use raw memset to avoid &mut self borrow conflict on the penalty buf - unsafe { - cudarc::driver::sys::cuMemsetD32Async( - temporal_penalty_buf_ptr, - 0, - 1, - self.stream.cu_stream(), - ); - } + let per_sample_buf_ptr = self.temporal_per_sample_buf.raw_ptr(); + let temporal_penalty_buf_ptr = self.temporal_penalty_buf.raw_ptr(); let ta = self.total_actions() as i32; let sh2 = self.config.shared_h2 as i32; + + // Step 1: per-sample penalty → temporal_per_sample_buf [B] unsafe { self.stream.launch_builder(&self.temporal_consistency_kernel) .arg(&save_h_s2_ptr) .arg(&q_out_buf_ptr) - .arg(&temporal_penalty_buf_ptr) + .arg(&per_sample_buf_ptr) .arg(&(batch_size as i32)) .arg(&sh2) .arg(&ta) .arg(&0.005_f32) // lambda_tc .arg(&0.95_f32) // sim_threshold - .launch(LaunchConfig::for_num_elems((batch_size - 1) as u32)) + .launch(LaunchConfig::for_num_elems(batch_size as u32)) .map_err(|e| MLError::ModelError(format!("temporal_consistency: {e}")))?; } + // Step 2: deterministic reduce → temporal_penalty_buf[0] for monitoring + unsafe { + self.stream + .launch_builder(&self.c51_loss_reduce_kernel) + .arg(&per_sample_buf_ptr) + .arg(&temporal_penalty_buf_ptr) + .arg(&(batch_size as i32)) + .launch(LaunchConfig { + grid_dim: (1, 1, 1), + block_dim: (1, 1, 1), + shared_mem_bytes: 0, + }) + .map_err(|e| MLError::ModelError(format!("temporal_consistency reduce: {e}")))?; + } Ok(()) } @@ -5840,7 +5849,11 @@ impl GpuDqnTrainer { let branch_indep_penalty_buf = stream.alloc_zeros::(1) .map_err(|e| MLError::ModelError(format!("alloc branch_indep_penalty_buf: {e}")))?; - // ── G10: Temporal consistency penalty buffer [1] ── + // ── G10: Temporal consistency penalty buffers ── + // Per-sample loss [B] (graph-safe, no atomicAdd, no memset) → + // reduced via c51_loss_reduce → scalar [1] for monitoring readback. + let temporal_per_sample_buf = stream.alloc_zeros::(config.batch_size) + .map_err(|e| MLError::ModelError(format!("alloc temporal_per_sample_buf: {e}")))?; let temporal_penalty_buf = stream.alloc_zeros::(1) .map_err(|e| MLError::ModelError(format!("alloc temporal_penalty_buf: {e}")))?; @@ -7372,6 +7385,7 @@ impl GpuDqnTrainer { branch_indep_kernel, branch_indep_penalty_buf, temporal_consistency_kernel, + temporal_per_sample_buf, temporal_penalty_buf, predictive_coding_kernel, predictive_coding_backward_kernel, @@ -8271,6 +8285,52 @@ impl GpuDqnTrainer { /// Read eval_q_std_ema. pub fn eval_q_std_ema(&self) -> f32 { self.eval_q_std_ema } + /// G6: Read branch_independence penalty scalar (sync DtoH; epoch-boundary only). + /// Returns 0.0 if buffer is uninitialized or sync fails — caller treats + /// as "no signal". Use HEALTH_DIAG to log. + pub fn branch_indep_loss_value(&self) -> f32 { + // SAFETY: sync DtoH at epoch boundary; not in graph capture region. + unsafe { cudarc::driver::sys::cuStreamSynchronize(self.stream.cu_stream()); } + let mut v = 0.0_f32; + unsafe { + cudarc::driver::sys::cuMemcpyDtoH_v2( + std::ptr::addr_of_mut!(v).cast(), + self.branch_indep_penalty_buf.raw_ptr(), + std::mem::size_of::(), + ); + } + v + } + + /// G10: Read temporal_consistency penalty scalar (sync DtoH; epoch-boundary only). + pub fn temporal_loss_value(&self) -> f32 { + unsafe { cudarc::driver::sys::cuStreamSynchronize(self.stream.cu_stream()); } + let mut v = 0.0_f32; + unsafe { + cudarc::driver::sys::cuMemcpyDtoH_v2( + std::ptr::addr_of_mut!(v).cast(), + self.temporal_penalty_buf.raw_ptr(), + std::mem::size_of::(), + ); + } + v + } + + /// G12: Read predictive_coding loss scalar (sync DtoH; epoch-boundary only). + /// Same pattern as G6/G10. Used for HEALTH_DIAG monitoring. + pub fn predictive_loss_value(&self) -> f32 { + unsafe { cudarc::driver::sys::cuStreamSynchronize(self.stream.cu_stream()); } + let mut v = 0.0_f32; + unsafe { + cudarc::driver::sys::cuMemcpyDtoH_v2( + std::ptr::addr_of_mut!(v).cast(), + self.predictive_loss_buf.raw_ptr(), + std::mem::size_of::(), + ); + } + v + } + /// Read per_branch_q_gap_ema — synchronous DtoH for trajectory backtracking. pub fn per_branch_q_gap_ema(&self) -> [f32; 4] { unsafe { cudarc::driver::sys::cuStreamSynchronize(self.stream.cu_stream()); } diff --git a/crates/ml/src/trainers/dqn/fused_training.rs b/crates/ml/src/trainers/dqn/fused_training.rs index d8d7799d2..98e5c48c2 100644 --- a/crates/ml/src/trainers/dqn/fused_training.rs +++ b/crates/ml/src/trainers/dqn/fused_training.rs @@ -868,6 +868,19 @@ impl FusedTrainingCtx { } } + /// G6/G10/G12 gem loss readbacks (sync DtoH; epoch-boundary only). + /// Returns (branch_indep_loss, temporal_consistency_loss, predictive_coding_loss). + /// Used by HEALTH_DIAG to surface gem-penalty magnitudes for the V7 + /// measurement-first methodology — values informs whether the existing + /// mechanisms (spectral_norm, NoisyNets) already cover them. + pub(crate) fn read_gem_losses(&self) -> (f32, f32, f32) { + ( + self.trainer.branch_indep_loss_value(), + self.trainer.temporal_loss_value(), + self.trainer.predictive_loss_value(), + ) + } + /// Flush the last in-flight async readback from the training step. /// Called once at epoch end to retrieve the final step's scalars. pub(crate) fn flush_readback(&mut self) -> Result { @@ -1466,6 +1479,18 @@ impl FusedTrainingCtx { self.trainer.compute_predictive_coding_loss(self.batch_size) .map_err(|e| anyhow::anyhow!("Predictive coding (G12): {e}"))?; + // G6 + G10 forward-only (measurement). Both kernels are now graph-safe + // (G6: single-block direct overwrite; G10: per-sample → c51_loss_reduce). + // No backward yet — wired only to log the penalty magnitudes via + // HEALTH_DIAG. Decision on backward integration depends on whether + // the measured penalty is meaningfully larger than what existing + // mechanisms (spectral_norm for G10, NoisyNets/ensemble-KL for G6) + // already enforce. V7-gem methodology: measure before commit. + self.trainer.compute_branch_independence(self.batch_size) + .map_err(|e| anyhow::anyhow!("Branch independence (G6 measure): {e}"))?; + self.trainer.compute_temporal_consistency(self.batch_size) + .map_err(|e| anyhow::anyhow!("Temporal consistency (G10 measure): {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() { diff --git a/crates/ml/src/trainers/dqn/trainer/training_loop.rs b/crates/ml/src/trainers/dqn/trainer/training_loop.rs index 914b0b12b..8e9c4f277 100644 --- a/crates/ml/src/trainers/dqn/trainer/training_loop.rs +++ b/crates/ml/src/trainers/dqn/trainer/training_loop.rs @@ -2148,8 +2148,16 @@ impl DQNTrainer { // diag = pure monitoring signals (not in reward): sharpe EMA (formerly w_dsr // reward input) and action entropy normalized to [0,1] (formerly // position_entropy_weight reward). See Phase 1 inventory better-form taxonomy. + // gems = G6/G10/G12 gem-penalty values (sync DtoH from GPU; once per epoch). + // V7 methodology: log first to measure whether the gem is doing real work + // before deciding to wire its backward gradient. + let (g6_loss, g10_loss, g12_loss) = if let Some(ref fused) = self.fused_ctx { + fused.read_gem_losses() + } else { + (0.0, 0.0, 0.0) + }; tracing::info!( - "HEALTH_DIAG[{}]: health={:.2} components [q_gap={:.2} q_var={:.2} atoms={:.2} grad_stable={:.2} ens_agree={:.2} grad_cos={:.2} spectral={:.2}] effective [cql_alpha={:.4} iqn_budget={:.2} cql_budget={:.2} c51_budget={:.2} tau={:.5} sarsa_tau={:.2} gamma={:.3} cf_ratio={:.2}] novels [distill={} barrier={:.3} plasticity={} ib={:.3} ensemble_collapse={:.2} contrarian={} meta_q_pred={:.2}] diag [sharpe_ema={:.3} action_entropy={:.2}]", + "HEALTH_DIAG[{}]: health={:.2} components [q_gap={:.2} q_var={:.2} atoms={:.2} grad_stable={:.2} ens_agree={:.2} grad_cos={:.2} spectral={:.2}] effective [cql_alpha={:.4} iqn_budget={:.2} cql_budget={:.2} c51_budget={:.2} tau={:.5} sarsa_tau={:.2} gamma={:.3} cf_ratio={:.2}] novels [distill={} barrier={:.3} plasticity={} ib={:.3} ensemble_collapse={:.2} contrarian={} meta_q_pred={:.2}] diag [sharpe_ema={:.3} action_entropy={:.2}] gems [g6_branch_indep={:.6} g10_temporal={:.6} g12_predictive={:.6}]", epoch, health_value, self.learning_health.components.q_gap_norm, @@ -2176,6 +2184,9 @@ impl DQNTrainer { self.last_meta_q_pred.unwrap_or(0.0), self.training_sharpe_ema, self.last_action_entropy.unwrap_or(1.0), + g6_loss, + g10_loss, + g12_loss, ); // C1/P1: propagate health to GPU replay buffer for diversity-weighted priorities.