diff --git a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs index 59eaad524..2e12e2a1f 100644 --- a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs +++ b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs @@ -1689,6 +1689,11 @@ pub struct GpuDqnTrainer { /// Formula: gamma_base + 0.005×(regime_stability − 0.5) − 0.05×(1 − health), /// clamped to [0.9, 0.995]. Populated by `apply_adaptive_gamma`. Default: 0.99. pub(crate) last_gamma_eff: f32, + + /// D1/N1: Ring buffer of best-health weight snapshots for temporal self-distillation. + 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, } impl GpuDqnTrainer { @@ -6976,6 +6981,9 @@ impl GpuDqnTrainer { let tau_final = 0.001_f32; let tau_anneal_steps: i32 = 100_000; + // D1/N1: clone Arc before it is moved into the struct literal. + let stream_for_snapshots = Arc::clone(&stream); + Ok(Self { config, stream, @@ -7478,6 +7486,11 @@ impl GpuDqnTrainer { last_c51_budget_eff: 0.55, last_ens_budget_eff: 0.05, last_gamma_eff: 0.99, + q_snapshots: crate::cuda_pipeline::q_snapshot::SnapshotRing::new( + stream_for_snapshots, + total_params, + ), + last_distill_active: false, }) } @@ -7702,6 +7715,130 @@ impl GpuDqnTrainer { gamma_clamped } + // ── D1/N1: Temporal self-distillation ───────────────────────────────── + + /// D1/N1: Whether distillation gradient was applied in the most recent epoch. + pub fn last_distill_active(&self) -> bool { + self.last_distill_active + } + + /// 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 { + use cudarc::driver::{DevicePtr, DevicePtrMut}; + use crate::cuda_pipeline::q_snapshot::SNAPSHOT_HEALTH_THRESHOLD; + + if health < SNAPSHOT_HEALTH_THRESHOLD { + return Ok(false); + } + + let ring = &mut self.q_snapshots; + if ring.snapshots.len() >= crate::cuda_pipeline::q_snapshot::MAX_SNAPSHOTS { + let worst_q_gap = ring.snapshots.iter().map(|s| s.q_gap).fold(f32::INFINITY, f32::min); + if q_gap <= worst_q_gap { + return Ok(false); + } + } + + let n = ring.param_count(); + let num_bytes = n * std::mem::size_of::(); + let mut new_weights = self.stream.alloc_zeros::(n) + .map_err(|e| MLError::ModelError(format!("snapshot alloc: {e}")))?; + { + let (src_ptr, _sg) = self.params_buf.device_ptr(&self.stream); + let (dst_ptr, _dg) = new_weights.device_ptr_mut(&self.stream); + #[allow(unsafe_code)] + unsafe { + cudarc::driver::result::memcpy_dtod_async( + dst_ptr, src_ptr, num_bytes, self.stream.cu_stream(), + ).map_err(|e| MLError::ModelError(format!("snapshot dtod: {e}")))?; + } + } + + let ring = &mut self.q_snapshots; + if ring.snapshots.len() >= crate::cuda_pipeline::q_snapshot::MAX_SNAPSHOTS { + let worst_idx = ring.snapshots.iter() + .enumerate() + .min_by(|(_, a), (_, b)| a.q_gap.partial_cmp(&b.q_gap).unwrap()) + .map(|(i, _)| i) + .unwrap(); + ring.snapshots.swap_remove(worst_idx); + } + + ring.snapshots.push(crate::cuda_pipeline::q_snapshot::QSnapshot { + weights: new_weights, + health, + q_gap, + epoch, + }); + Ok(true) + } + + /// D1/N1: Apply distillation gradient — pulls current weights toward the best + /// snapshot with strength `0.1 × (1 − health)`. Returns true if applied. + /// + /// Implemented as two SAXPY calls into grad_buf (ungraphed, stream-ordered): + /// grad += alpha * params (adds current) + /// grad += -alpha * best_snap (subtracts best) + /// Net effect: grad += alpha * (params - best_snap) + /// + /// Uses `saxpy_f32_aux` (dqn_saxpy_f32_kernel, aux_child handle) — same handle + /// used by `apply_cql_saxpy`, safe for ungraphed injection before graph_adam. + pub fn apply_distillation_gradient(&mut self) -> Result { + let (health, _) = self.read_isv_health_and_regime(); + let distill_weight = 0.1_f32 * (1.0 - health).max(0.0); + if distill_weight < 0.01 { + self.last_distill_active = false; + return Ok(false); + } + let best_ptr = match self.q_snapshots.best_raw_ptr() { + Some(p) => p, + None => { + self.last_distill_active = false; + return Ok(false); + } + }; + + let grad_ptr = self.ptrs.grad_buf; + let params_ptr = self.ptrs.params_ptr; + let n = self.total_params as i32; + let blocks = self.grad_norm_blocks as u32; + + unsafe { + // grad += distill_weight * params_buf + self.stream + .launch_builder(&self.saxpy_f32_aux) + .arg(&grad_ptr) + .arg(¶ms_ptr) + .arg(&distill_weight) + .arg(&n) + .launch(LaunchConfig { + grid_dim: (blocks, 1, 1), + block_dim: (256, 1, 1), + shared_mem_bytes: 0, + }) + .map_err(|e| MLError::ModelError(format!("distill saxpy (+params): {e}")))?; + + // grad += (-distill_weight) * best_snapshot + let neg_weight = -distill_weight; + self.stream + .launch_builder(&self.saxpy_f32_aux) + .arg(&grad_ptr) + .arg(&best_ptr) + .arg(&neg_weight) + .arg(&n) + .launch(LaunchConfig { + grid_dim: (blocks, 1, 1), + block_dim: (256, 1, 1), + shared_mem_bytes: 0, + }) + .map_err(|e| MLError::ModelError(format!("distill saxpy (-best): {e}")))?; + } + + self.last_distill_active = true; + Ok(true) + } + // ── A3: LearningHealth signal writers / readers ─────────────────────── /// Write a scalar into the ISV pinned buffer at the given index. diff --git a/crates/ml/src/cuda_pipeline/mod.rs b/crates/ml/src/cuda_pipeline/mod.rs index a8df655e9..99658f03a 100644 --- a/crates/ml/src/cuda_pipeline/mod.rs +++ b/crates/ml/src/cuda_pipeline/mod.rs @@ -38,6 +38,7 @@ pub mod gpu_iqn_head; pub mod gpu_attention; pub mod decision_transformer; pub mod learning_health; +pub mod q_snapshot; // gpu_replay_buffer moved to ml-dqn crate /// Maximum bytes allowed for a single GPU upload (2 GB safety limit). diff --git a/crates/ml/src/cuda_pipeline/q_snapshot.rs b/crates/ml/src/cuda_pipeline/q_snapshot.rs new file mode 100644 index 000000000..a63d88441 --- /dev/null +++ b/crates/ml/src/cuda_pipeline/q_snapshot.rs @@ -0,0 +1,123 @@ +//! N1: Q-network weight snapshots ring buffer for temporal self-distillation. +//! +//! Keeps up to `MAX_SNAPSHOTS` historical weight checkpoints taken at high-health +//! epochs. During collapse recovery, a gradient pull toward the best snapshot +//! is added to the master gradient. + +use std::sync::Arc; +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; + +/// A single weight snapshot with metadata. +pub struct QSnapshot { + pub weights: CudaSlice, + pub health: f32, + pub q_gap: f32, + pub epoch: u32, +} + +impl std::fmt::Debug for QSnapshot { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("QSnapshot") + .field("health", &self.health) + .field("q_gap", &self.q_gap) + .field("epoch", &self.epoch) + .finish() + } +} + +/// Ring buffer of best-health snapshots. +pub struct SnapshotRing { + stream: Arc, + pub(crate) snapshots: Vec, + param_count: usize, +} + +impl std::fmt::Debug for SnapshotRing { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("SnapshotRing") + .field("param_count", &self.param_count) + .field("snapshot_count", &self.snapshots.len()) + .finish() + } +} + +impl SnapshotRing { + pub fn new(stream: Arc, param_count: usize) -> Self { + Self { + stream, + snapshots: Vec::with_capacity(MAX_SNAPSHOTS), + param_count, + } + } + + pub fn is_empty(&self) -> bool { + self.snapshots.is_empty() + } + + pub fn param_count(&self) -> usize { + self.param_count + } + + /// Copy current weights into a new snapshot if health >= threshold and + /// q_gap beats the worst stored snapshot (when at capacity). + pub fn maybe_snapshot( + &mut self, + params_src: &CudaSlice, + health: f32, + q_gap: f32, + epoch: u32, + ) -> Result { + if health < SNAPSHOT_HEALTH_THRESHOLD { + return Ok(false); + } + + // At capacity, only replace if this q_gap beats the worst stored entry. + if self.snapshots.len() >= MAX_SNAPSHOTS { + let worst_q_gap = self.snapshots.iter() + .map(|s| s.q_gap) + .fold(f32::INFINITY, f32::min); + if q_gap <= worst_q_gap { + return Ok(false); + } + } + + let n = self.param_count; + let num_bytes = n * std::mem::size_of::(); + let mut new_weights = self.stream.alloc_zeros::(n) + .map_err(|e| MLError::ModelError(format!("snapshot alloc: {e}")))?; + { + let (src_ptr, _src_guard) = params_src.device_ptr(&self.stream); + let (dst_ptr, _dst_guard) = new_weights.device_ptr_mut(&self.stream); + #[allow(unsafe_code)] + unsafe { + cudarc::driver::result::memcpy_dtod_async( + dst_ptr, src_ptr, num_bytes, self.stream.cu_stream(), + ).map_err(|e| MLError::ModelError(format!("snapshot dtod: {e}")))?; + } + } + + if self.snapshots.len() >= MAX_SNAPSHOTS { + let worst_idx = self.snapshots.iter() + .enumerate() + .min_by(|(_, a), (_, b)| a.q_gap.partial_cmp(&b.q_gap).unwrap()) + .map(|(i, _)| i) + .unwrap(); + self.snapshots.swap_remove(worst_idx); + } + + self.snapshots.push(QSnapshot { weights: new_weights, health, q_gap, epoch }); + Ok(true) + } + + /// Best (highest q_gap) snapshot raw device pointer, if any. + pub fn best_raw_ptr(&self) -> Option { + self.snapshots.iter() + .max_by(|a, b| a.q_gap.partial_cmp(&b.q_gap).unwrap()) + .map(|s| s.weights.raw_ptr()) + } +} diff --git a/crates/ml/src/trainers/dqn/fused_training.rs b/crates/ml/src/trainers/dqn/fused_training.rs index 306613c9f..d5fc8a255 100644 --- a/crates/ml/src/trainers/dqn/fused_training.rs +++ b/crates/ml/src/trainers/dqn/fused_training.rs @@ -2247,6 +2247,25 @@ impl FusedTrainingCtx { self.trainer.compute_q_spectral_gap() } + /// D1/N1: Conditionally snapshot current weights at epoch boundary. + /// Snapshots if health >= 0.7 and q_gap beats the worst stored entry. + pub(crate) fn maybe_snapshot_qnet(&mut self, health: f32, q_gap: f32, epoch: u32) -> anyhow::Result { + self.trainer.maybe_snapshot_params(health, q_gap, epoch) + .map_err(|e| anyhow::anyhow!("snapshot: {e}")) + } + + /// D1/N1: Apply distillation gradient pull toward best snapshot. + /// No-op when health >= 0.4 or no snapshots exist yet. + pub(crate) fn apply_distillation(&mut self) -> anyhow::Result { + self.trainer.apply_distillation_gradient() + .map_err(|e| anyhow::anyhow!("distill: {e}")) + } + + /// D1/N1: Whether distillation was active in the most recent epoch boundary. + pub(crate) fn last_distill_active(&self) -> bool { + self.trainer.last_distill_active + } + 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 70ccce2c8..a797b3f02 100644 --- a/crates/ml/src/trainers/dqn/trainer/training_loop.rs +++ b/crates/ml/src/trainers/dqn/trainer/training_loop.rs @@ -1890,6 +1890,16 @@ impl DQNTrainer { self.last_gamma_eff = Some(fused.last_gamma_eff()); } + // D1/N1: snapshot high-health checkpoints and pull toward best during collapse. + 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 _ = fused.apply_distillation(); + } + self.last_distill_active = Some(fused.last_distill_active()); + } + // HEALTH_DIAG: components are [0, 1] normalized. effective = hyperparams after health-adaptation. novels = mechanism states. 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}]",