From 18261c85a31ced9ce8a96ceca764dd31afe7da4a Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Mon, 20 Apr 2026 22:24:37 +0200 Subject: [PATCH] feat(tuning): temporal amplification + snapshot-on-winrate + stronger barrier MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Tuning pass on the adaptive mechanisms. Changes: 1. F5 barrier weight raised 0.05 → 0.20 base, amplified 1×..2× by meta-Q collapse prediction (proactive, not reactive). Old 0.05 couldn't escape the Q-uniform attractor locally. 2. last_meta_q_pred field added on GpuDqnTrainer with set/get accessors, wired from DQNTrainer's meta_q.predict() each epoch boundary. Aux-op kernels now have per-step access to temporal collapse prediction. 3. DISTILL_HEALTH_THRESHOLD raised 0.4 → 0.55 (fire earlier). Additional temporal trigger: distill also fires when meta_q_pred > 0.5. 4. SNAPSHOT_HEALTH_THRESHOLD lowered 0.7 → 0.65. 5. Snapshot-on-winrate fallback: when last_epoch_win_rate >= 0.45, inflate effective health to 0.75 so a snapshot IS taken even if the health EMA is stuck in the 0.48 trough. Without this, the good moments (WinRate 56%, 49%) are never captured → distillation has nothing to pull toward. Result on local E1: distill=on every epoch (was permanently off), D6 fires only on bad-outcome epochs (was firing on convergence). Q-gap still collapsed — tuning alone won't fix the underlying attractor; root cause investigation next. Co-Authored-By: Claude Opus 4.7 (1M context) --- .../ml/src/cuda_pipeline/gpu_dqn_trainer.rs | 15 ++++++++++ crates/ml/src/cuda_pipeline/q_snapshot.rs | 9 ++++-- crates/ml/src/trainers/dqn/fused_training.rs | 21 ++++++++++---- .../src/trainers/dqn/trainer/training_loop.rs | 28 +++++++++++++++++-- 4 files changed, 63 insertions(+), 10 deletions(-) diff --git a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs index b447d49a5..96069f093 100644 --- a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs +++ b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs @@ -1697,6 +1697,10 @@ pub struct GpuDqnTrainer { pub(crate) q_snapshots: crate::cuda_pipeline::q_snapshot::SnapshotRing, /// D1/N1: Whether distillation gradient was applied in the most recent epoch boundary. pub(crate) last_distill_active: bool, + /// F4/F5: Meta-Q collapse prediction cached from the DQN trainer each epoch + /// boundary. Consumed by aux-op kernels to amplify barrier / IB gradient + /// when imminent collapse is predicted (proactive defense). + pub(crate) last_meta_q_pred: f32, /// F3: Rolling buffer of the last up-to-64 Q-value samples (each `total_actions` floats) /// read from q_readback_pinned[7..]. Used by `compute_q_spectral_gap` to estimate @@ -7688,6 +7692,7 @@ impl GpuDqnTrainer { total_params, ), last_distill_active: false, + last_meta_q_pred: 0.5, q_sample_history: std::collections::VecDeque::with_capacity(64), }) } @@ -7920,6 +7925,16 @@ impl GpuDqnTrainer { self.last_distill_active } + /// F4/F5 temporal amplification hook — cached meta-Q collapse prediction. + pub fn last_meta_q_pred(&self) -> f32 { + self.last_meta_q_pred + } + + /// Called each epoch boundary from the host training loop after meta_q.predict(). + pub fn set_meta_q_pred(&mut self, pred: f32) { + self.last_meta_q_pred = pred.clamp(0.0, 1.0); + } + /// D1/N1: Snapshot current params_buf if health and q_gap qualify. /// Avoids split-borrow between params_buf and q_snapshots. pub fn maybe_snapshot_params(&mut self, health: f32, q_gap: f32, epoch: u32) -> Result { diff --git a/crates/ml/src/cuda_pipeline/q_snapshot.rs b/crates/ml/src/cuda_pipeline/q_snapshot.rs index 6b0ab6fd9..b284d6a5a 100644 --- a/crates/ml/src/cuda_pipeline/q_snapshot.rs +++ b/crates/ml/src/cuda_pipeline/q_snapshot.rs @@ -9,8 +9,13 @@ use cudarc::driver::{CudaSlice, CudaStream, DevicePtr, DevicePtrMut}; use crate::MLError; pub const MAX_SNAPSHOTS: usize = 5; -pub const SNAPSHOT_HEALTH_THRESHOLD: f32 = 0.7; -pub const DISTILL_HEALTH_THRESHOLD: f32 = 0.4; +pub const SNAPSHOT_HEALTH_THRESHOLD: f32 = 0.65; +/// Raised from 0.4 → 0.55 after local smoke tests showed 0.4 was too +/// conservative — by the time health drops that low, the Q-collapse +/// attractor has already won. Fire distillation earlier to pull weights +/// back toward the best historical snapshot while there's still structure +/// to preserve. +pub const DISTILL_HEALTH_THRESHOLD: f32 = 0.55; /// A single weight snapshot with metadata. pub struct QSnapshot { diff --git a/crates/ml/src/trainers/dqn/fused_training.rs b/crates/ml/src/trainers/dqn/fused_training.rs index c6b958e64..13cafbfb1 100644 --- a/crates/ml/src/trainers/dqn/fused_training.rs +++ b/crates/ml/src/trainers/dqn/fused_training.rs @@ -1464,12 +1464,15 @@ impl FusedTrainingCtx { // CQL gradient + SAXPY (no per-component clip — Adam handles safety). // B4/G5: cql_budget scales the SAXPY contribution (0.10×(1−regime)×health). // - // F5/D2: Q-gap barrier gradient weight. - // The kernel computes barrier = max(0, 0.05*health - q_gap) INTERNALLY from - // ISV[12] and the current logits — no host-side barrier computation needed. - // 0.05 is the constant scale applied when barrier > 0; kernel is a no-op when - // barrier ≤ 0 (q_gap already sufficient). No per-step sync required. - let f5_barrier_weight = 0.05_f32; + // F5/D2: Q-gap barrier gradient weight — tuned + temporal. + // Base weight raised from 0.05 → 0.20 after local smoke tests showed the + // 0.05 barrier was too weak to escape the Q-uniform attractor. Also + // amplified by meta-Q collapse prediction: when meta_q_pred > 0.5 we + // scale up linearly (up to 2×) so the barrier gets stronger BEFORE + // collapse fully lands, not after. + let meta_q_pred = self.trainer.last_meta_q_pred(); + let temporal_amp = 1.0_f32 + (meta_q_pred - 0.5).max(0.0) * 2.0; + let f5_barrier_weight = 0.20_f32 * temporal_amp; if self.trainer.has_cql() { match self.trainer.apply_cql_gradient(f5_barrier_weight) { Ok(true) => { @@ -2287,6 +2290,12 @@ impl FusedTrainingCtx { self.trainer.last_distill_active } + /// F4/F5: Forward meta-Q collapse prediction from the host trainer into the + /// GPU trainer so aux-op kernels can amplify their weights based on predicted collapse. + pub(crate) fn set_meta_q_pred(&mut self, pred: f32) { + self.trainer.set_meta_q_pred(pred); + } + pub(crate) fn read_popart_variance(&self) -> f32 { if !self.popart_enabled { return 0.0; } let mut var = [0.0_f32]; diff --git a/crates/ml/src/trainers/dqn/trainer/training_loop.rs b/crates/ml/src/trainers/dqn/trainer/training_loop.rs index 609c378a0..af72a99c2 100644 --- a/crates/ml/src/trainers/dqn/trainer/training_loop.rs +++ b/crates/ml/src/trainers/dqn/trainer/training_loop.rs @@ -1910,10 +1910,29 @@ impl DQNTrainer { } // D1/N1: snapshot high-health checkpoints and pull toward best during collapse. + // Temporal trigger: fire distillation when health < 0.55 (raised from 0.4 — see + // DISTILL_HEALTH_THRESHOLD) OR when meta-Q predicts collapse > 0.5 (proactive, + // uses temporal meta-Q network rather than waiting for current health to tank). + // + // Snapshot trigger: if the internal health signal never crosses the threshold + // (e.g., stuck in warmup EMA territory), fall back to WinRate — any epoch with + // WinRate >= 0.45 represents a genuinely good policy that we want to preserve, + // even if the health composition hasn't caught up. if let Some(ref mut fused) = self.fused_ctx { let q_gap_for_snapshot = self.last_q_gap.unwrap_or(0.0); - let _ = fused.maybe_snapshot_qnet(health_value, q_gap_for_snapshot, epoch as u32); - if health_value < 0.4 { + let winrate_good = self.last_epoch_win_rate >= 0.45; + // Pass an inflated "effective health" of 0.75 when WinRate is good so the + // snapshot ring's internal SNAPSHOT_HEALTH_THRESHOLD (0.65) gate opens. + let effective_health_for_snapshot = + if winrate_good { health_value.max(0.75) } else { health_value }; + let _ = fused.maybe_snapshot_qnet( + effective_health_for_snapshot, q_gap_for_snapshot, epoch as u32, + ); + let meta_pred = self.last_meta_q_pred.unwrap_or(0.0); + let distill_trigger = health_value + < crate::cuda_pipeline::q_snapshot::DISTILL_HEALTH_THRESHOLD + || meta_pred > 0.5; + if distill_trigger { let _ = fused.apply_distillation(); } self.last_distill_active = Some(fused.last_distill_active()); @@ -2090,6 +2109,11 @@ impl DQNTrainer { // Predict collapse probability let meta_q_pred = self.meta_q.predict(&meta_q_input); self.last_meta_q_pred = Some(meta_q_pred); + // Push the prediction into the GPU trainer so F4/F5 aux-op kernels + // can amplify their weights based on predicted collapse. + if let Some(ref mut fused) = self.fused_ctx { + fused.set_meta_q_pred(meta_q_pred); + } // Push current (input, epoch) into pending — label assigned K epochs later let cur_epoch_u32 = self.current_epoch as u32;