From d346e03d48cb0bb3fab1ece01b6cf0f9aab9062e Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Fri, 24 Apr 2026 17:53:30 +0200 Subject: [PATCH] =?UTF-8?q?feat(dqn-v2):=20C.6=20Task=2014=20=E2=80=94=20e?= =?UTF-8?q?psilon=20GPU=20kernel=20+=20CPU=20monitor=20(exploration)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit epsilon_update kernel computes effective epsilon from ISV[EPOCH_IDX=39, TOTAL_EPOCHS=40, LEARNING_HEALTH=12] with cosine schedule (eps_start -> eps_end) plus health-coupled boost (up to +0.1 at collapse). Single-thread cold-path kernel writes ISV[EPSILON_EFF_INDEX=41]. EpsilonMonitor is a read-only observer exposing eps_eff, epoch_idx, total_ep, health, progress, fire_rate in DiagSnapshot. Wires epsilon kernel launch at epoch boundary alongside tau kernel. Consumer migration: log_training_config ISV-adaptive epsilon block now reads ISV[EPSILON_EFF_INDEX] instead of computing `base_floor * (0.5 + volatility)`. Tests: 3 monitor unit tests pass. cargo check -p ml at 8-warning baseline. Plan 1 Task 14. Spec §4.C.6 (2026-04-24 revision). Co-Authored-By: Claude Opus 4.7 (1M context) --- crates/ml/build.rs | 1 + .../cuda_pipeline/epsilon_update_kernel.cu | 27 +++++ .../ml/src/cuda_pipeline/gpu_dqn_trainer.rs | 40 +++++++ crates/ml/src/trainers/dqn/fused_training.rs | 8 ++ .../trainers/dqn/monitors/epsilon_monitor.rs | 113 ++++++++++++++++++ crates/ml/src/trainers/dqn/monitors/mod.rs | 1 + .../src/trainers/dqn/trainer/training_loop.rs | 24 +++- docs/dqn-gpu-hot-path-audit.md | 1 + docs/dqn-wire-up-audit.md | 2 + 9 files changed, 211 insertions(+), 6 deletions(-) create mode 100644 crates/ml/src/cuda_pipeline/epsilon_update_kernel.cu create mode 100644 crates/ml/src/trainers/dqn/monitors/epsilon_monitor.rs diff --git a/crates/ml/build.rs b/crates/ml/build.rs index 6bdf265e6..bffc32839 100644 --- a/crates/ml/build.rs +++ b/crates/ml/build.rs @@ -85,6 +85,7 @@ fn main() { "branch_grad_balance_kernel.cu", "backtest_plan_kernel.cu", "tau_update_kernel.cu", + "epsilon_update_kernel.cu", ]; // ALL kernels get common header (BF16 types + wrappers) diff --git a/crates/ml/src/cuda_pipeline/epsilon_update_kernel.cu b/crates/ml/src/cuda_pipeline/epsilon_update_kernel.cu new file mode 100644 index 000000000..7c65f3256 --- /dev/null +++ b/crates/ml/src/cuda_pipeline/epsilon_update_kernel.cu @@ -0,0 +1,27 @@ +/* epsilon_update — effective exploration epsilon, GPU-driven (Plan 1 Task 14). + * Reads: ISV[EPOCH_IDX_INDEX=39], ISV[TOTAL_EPOCHS_INDEX=40], + * ISV[LEARNING_HEALTH_INDEX=12] + * Writes: ISV[EPSILON_EFF_INDEX=41] + * Cold path (per-epoch boundary). Single-thread kernel. + * + * Formula: cosine schedule (eps_start -> eps_end) over training, plus + * a volatility boost of up to +0.1 when health=0 (full collapse). + */ +extern "C" __global__ void epsilon_update( + const float* __restrict__ isv, + float* __restrict__ isv_out, + float eps_start, + float eps_end +) { + if (threadIdx.x != 0 || blockIdx.x != 0) return; + float epoch = isv[39]; + float total_ep = fmaxf(isv[40], 1.0f); + float health = fminf(fmaxf(isv[12], 0.0f), 1.0f); + float progress = fminf(fmaxf(epoch / total_ep, 0.0f), 1.0f); + float cos_f = 0.5f * (1.0f + cosf(3.14159265358979f * progress)); + float eps_cos = eps_end + (eps_start - eps_end) * cos_f; + float boost = (1.0f - health) * 0.1f; /* up to +0.1 when health=0 */ + float eps_eff = eps_cos + boost; + eps_eff = fminf(fmaxf(eps_eff, 0.0f), 1.0f); + isv_out[41] = eps_eff; +} diff --git a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs index 8161bb7c4..9976a3101 100644 --- a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs +++ b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs @@ -79,6 +79,8 @@ static GRAD_DECOMP_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/grad static BRANCH_GRAD_BALANCE_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/branch_grad_balance_kernel.cubin")); /// Plan 1 Task 13: GPU-driven Polyak-EMA tau coefficient. static TAU_UPDATE_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/tau_update_kernel.cubin")); +/// Plan 1 Task 14: GPU-driven effective exploration epsilon. +static EPSILON_UPDATE_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/epsilon_update_kernel.cubin")); /// Mamba2 temporal scan configuration. const MAMBA2_HISTORY_K: usize = 8; // Rolling history length @@ -1402,6 +1404,9 @@ pub struct GpuDqnTrainer { /// Plan 1 Task 13: GPU-driven Polyak-EMA tau. Cold-path (per-epoch). /// Reads ISV[39, 40, 12]; writes ISV[TAU_EFF_INDEX=42]. tau_update_kernel: CudaFunction, + /// Plan 1 Task 14: GPU-driven effective epsilon. Cold-path (per-epoch). + /// Reads ISV[39, 40, 12]; writes ISV[EPSILON_EFF_INDEX=41]. + epsilon_update_kernel: CudaFunction, /// Host-side cached per-component norms populated at epoch boundary /// from the pinned result slot. Index order matches `grad_mag_*_ratio` /// accessors: `[iqn, cql, cql_sx, c51, c51_bs, distill, rec, pred, ens]`. @@ -7042,6 +7047,14 @@ impl GpuDqnTrainer { .map_err(|e| MLError::ModelError(format!("tau_update load: {e}")))? }; + // Plan 1 Task 14: load epsilon_update kernel (cold-path, per-epoch). + let epsilon_update_kernel = { + let module = stream.context().load_cubin(EPSILON_UPDATE_CUBIN.to_vec()) + .map_err(|e| MLError::ModelError(format!("epsilon_update cubin load: {e}")))?; + module.load_function("epsilon_update") + .map_err(|e| MLError::ModelError(format!("epsilon_update load: {e}")))? + }; + // eval_v_range — pinned device-mapped (zero-copy write from CPU). let eval_v_range_pinned: *mut f32 = unsafe { let flags = cudarc::driver::sys::CU_MEMHOSTALLOC_DEVICEMAP; @@ -9099,6 +9112,7 @@ impl GpuDqnTrainer { branch_grad_scales_dev, branch_slice_max_len, tau_update_kernel, + epsilon_update_kernel, grad_component_norms_mag: [0.0_f32; 9], grad_component_norms_dir: [0.0_f32; 9], grad_component_norms_trunk: [0.0_f32; 9], @@ -10433,6 +10447,32 @@ impl GpuDqnTrainer { Ok(()) } + /// Plan 1 Task 14: GPU-driven epsilon update. + /// + /// Single-thread cold-path kernel: reads ISV[EPOCH_IDX=39, TOTAL_EPOCHS=40, + /// LEARNING_HEALTH=12]; writes ISV[EPSILON_EFF_INDEX=41] with cosine schedule + /// + health-coupled boost. Must be called once per epoch boundary BEFORE + /// any epsilon consumer reads ISV[EPSILON_EFF_INDEX]. + pub(crate) fn launch_epsilon_update(&self, eps_start: f32, eps_end: f32) -> Result<(), MLError> { + let isv_ptr = self.isv_signals_dev_ptr; + let isv_out = isv_ptr; // kernel writes to the same ISV bus (offset 41) + unsafe { + self.stream + .launch_builder(&self.epsilon_update_kernel) + .arg(&isv_ptr) + .arg(&isv_out) + .arg(&eps_start) + .arg(&eps_end) + .launch(LaunchConfig { + grid_dim: (1, 1, 1), + block_dim: (1, 1, 1), + shared_mem_bytes: 0, + }) + .map_err(|e| MLError::ModelError(format!("epsilon_update launch: {e}")))?; + } + Ok(()) + } + /// Read back the per-branch scales applied on the most recent /// `launch_branch_grad_balance` call — `[dir, mag, ord, urg]`. /// Values in `[0, 1]`; `1.0` means the branch passed through diff --git a/crates/ml/src/trainers/dqn/fused_training.rs b/crates/ml/src/trainers/dqn/fused_training.rs index 0658a51b7..360631a3c 100644 --- a/crates/ml/src/trainers/dqn/fused_training.rs +++ b/crates/ml/src/trainers/dqn/fused_training.rs @@ -3104,6 +3104,14 @@ impl FusedTrainingCtx { .map_err(|e| anyhow::anyhow!("launch_tau_update: {e}")) } + /// Plan 1 Task 14: delegate GPU epsilon update to trainer. + pub(crate) fn launch_epsilon_update(&self, eps_start: f32, eps_end: f32) + -> anyhow::Result<()> + { + self.trainer.launch_epsilon_update(eps_start, eps_end) + .map_err(|e| anyhow::anyhow!("launch_epsilon_update: {e}")) + } + /// G5: Set the EMA threshold for epistemic variance gating. pub(crate) fn set_var_ema(&self, val: f32) { self.trainer.set_var_ema(val); diff --git a/crates/ml/src/trainers/dqn/monitors/epsilon_monitor.rs b/crates/ml/src/trainers/dqn/monitors/epsilon_monitor.rs new file mode 100644 index 000000000..2ed394a21 --- /dev/null +++ b/crates/ml/src/trainers/dqn/monitors/epsilon_monitor.rs @@ -0,0 +1,113 @@ +//! CPU-side observer for the GPU-driven effective exploration epsilon. +//! +//! The `epsilon_update` CUDA kernel reads ISV[EPOCH_IDX_INDEX=39], +//! ISV[TOTAL_EPOCHS_INDEX=40], and ISV[LEARNING_HEALTH_INDEX=12] to compute +//! a cosine-annealed epsilon with a health-coupled boost, writing the result +//! to ISV[EPSILON_EFF_INDEX=41]. This monitor reads that slot for HEALTH_DIAG +//! emission and fire-rate tracking. +//! +//! The monitor's `read()` returns `isv[EPSILON_EFF_INDEX]`. `diagnose()` +//! exposes all four input slots plus the effective epsilon and fire-rate. + +use crate::cuda_pipeline::gpu_dqn_trainer::{ + EPOCH_IDX_INDEX, TOTAL_EPOCHS_INDEX, LEARNING_HEALTH_INDEX, EPSILON_EFF_INDEX, +}; +use crate::trainers::dqn::adaptive_monitor::{ + AdaptiveMonitor, DiagSnapshot, FireRateStats, IsvBus, +}; + +pub(crate) struct EpsilonMonitor { + last_eps: f32, + fire_rate: FireRateStats, +} + +impl EpsilonMonitor { + pub(crate) fn new() -> Self { + Self { last_eps: 0.0, fire_rate: FireRateStats::default() } + } +} + +impl AdaptiveMonitor for EpsilonMonitor { + fn read(&self, isv: &IsvBus<'_>) -> f32 { + isv.read(EPSILON_EFF_INDEX) + } + + fn diagnose(&self, isv: &IsvBus<'_>) -> DiagSnapshot { + let eps_eff = isv.read(EPSILON_EFF_INDEX); + let epoch = isv.read(EPOCH_IDX_INDEX); + let total_ep = isv.read(TOTAL_EPOCHS_INDEX); + let health = isv.read(LEARNING_HEALTH_INDEX); + let progress = if total_ep > 0.0 { (epoch / total_ep).clamp(0.0, 1.0) } else { 0.0 }; + DiagSnapshot::new() + .with("eps_eff", eps_eff) + .with("epoch_idx", epoch) + .with("total_ep", total_ep) + .with("health", health) + .with("progress", progress) + .with("fire_rate", self.fire_rate.fire_rate()) + } + + fn observe(&mut self, current: f32) { + let fired = (current - self.last_eps).abs() > 1e-6; + self.fire_rate.record_fire(fired); + self.last_eps = current; + } + + fn fire_rate(&self) -> &FireRateStats { + &self.fire_rate + } + + fn name(&self) -> &'static str { + "epsilon" + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn make_isv(epoch: f32, total: f32, health: f32, eps_eff: f32) -> Vec { + let mut v = vec![0.0f32; 49]; + v[EPOCH_IDX_INDEX] = epoch; + v[TOTAL_EPOCHS_INDEX] = total; + v[LEARNING_HEALTH_INDEX] = health; + v[EPSILON_EFF_INDEX] = eps_eff; + v + } + + #[test] + fn read_returns_epsilon_eff_slot() { + let mut buf = make_isv(5.0, 100.0, 0.8, 0.15); + let bus = IsvBus::new(&mut buf); + let mon = EpsilonMonitor::new(); + let v = mon.read(&bus); + assert!((v - 0.15).abs() < 1e-7, "expected eps_eff=0.15, got {v}"); + } + + #[test] + fn diagnose_exposes_all_fields() { + let mut buf = make_isv(20.0, 100.0, 0.7, 0.12); + let bus = IsvBus::new(&mut buf); + let mon = EpsilonMonitor::new(); + let snap = mon.diagnose(&bus); + assert!(snap.fields.contains_key("eps_eff")); + assert!(snap.fields.contains_key("epoch_idx")); + assert!(snap.fields.contains_key("total_ep")); + assert!(snap.fields.contains_key("health")); + assert!(snap.fields.contains_key("progress")); + assert!(snap.fields.contains_key("fire_rate")); + let progress = snap.fields["progress"]; + assert!((progress - 0.2).abs() < 1e-6, "expected progress=0.2, got {progress}"); + } + + #[test] + fn observe_tracks_fire_rate_on_change() { + let mut mon = EpsilonMonitor::new(); + mon.observe(0.20); + mon.observe(0.20); + mon.observe(0.15); + let rate = mon.fire_rate().fire_rate(); + assert!((rate - 2.0 / 3.0).abs() < 1e-6, + "expected 2/3 fire rate, got {rate}"); + } +} diff --git a/crates/ml/src/trainers/dqn/monitors/mod.rs b/crates/ml/src/trainers/dqn/monitors/mod.rs index f75e9e394..140333d87 100644 --- a/crates/ml/src/trainers/dqn/monitors/mod.rs +++ b/crates/ml/src/trainers/dqn/monitors/mod.rs @@ -5,5 +5,6 @@ //! read those slots for HEALTH_DIAG emission and `controller_activity` //! smoke-test fire-rate tracking. No CPU-side adaptive computation. +pub(crate) mod epsilon_monitor; pub(crate) mod grad_balancer_monitor; pub(crate) mod tau_monitor; diff --git a/crates/ml/src/trainers/dqn/trainer/training_loop.rs b/crates/ml/src/trainers/dqn/trainer/training_loop.rs index 9b4d5fb3b..377ce4c7b 100644 --- a/crates/ml/src/trainers/dqn/trainer/training_loop.rs +++ b/crates/ml/src/trainers/dqn/trainer/training_loop.rs @@ -21,7 +21,7 @@ use cudarc::driver::CudaSlice; use common::metrics::{questdb_sink, training_metrics}; use tracing::{debug, info, warn}; -use crate::cuda_pipeline::gpu_dqn_trainer::EPOCH_IDX_INDEX; +use crate::cuda_pipeline::gpu_dqn_trainer::{EPOCH_IDX_INDEX, EPSILON_EFF_INDEX}; use crate::dqn::logging::{log_epoch_start, log_epoch_end, log_training_progress}; use crate::dqn::target_update::convergence_half_life; use crate::evaluation::metrics::calculate_var_cvar; @@ -258,6 +258,15 @@ impl DQNTrainer { } } + // Plan 1 Task 14: launch GPU epsilon update kernel at epoch boundary. + if let Some(ref fused) = self.fused_ctx { + let eps_start = self.hyperparams.epsilon_start as f32; + let eps_end = self.hyperparams.epsilon_end as f32; + if let Err(e) = fused.launch_epsilon_update(eps_start, eps_end) { + tracing::warn!(epoch, "launch_epsilon_update failed (non-fatal): {e}"); + } + } + // Collect pending async validation from PREVIOUS epoch if let Some(val) = self.pending_val_loss.take() { // Store for this epoch's logging (one-epoch delay is acceptable) @@ -928,17 +937,20 @@ impl DQNTrainer { } { - // ISV-adaptive exploration: read regime signals, modulate epsilon + sigma. - // Volatile regime → explore more (higher epsilon, sigma). - // Stable regime → exploit (lower epsilon, sigma). + // Plan 1 Task 14: read GPU-computed epsilon from ISV[EPSILON_EFF_INDEX]. + // epsilon_update kernel ran at epoch boundary with cosine schedule + + // health-coupled boost; no CPU-side adaptive computation here. let (regime_stab, volatility) = if let Some(ref fused) = self.fused_ctx { fused.read_isv_regime() } else { (0.5, 0.5) }; - let base_floor = self.hyperparams.noisy_epsilon_floor.unwrap_or(0.05) as f32; - let adaptive_epsilon = base_floor * (0.5 + volatility); + let adaptive_epsilon = if let Some(ref fused) = self.fused_ctx { + fused.trainer().read_isv_signal_at(EPSILON_EFF_INDEX) + } else { + self.hyperparams.noisy_epsilon_floor.unwrap_or(0.05) as f32 + }; let mut agent = self.agent.write().await; agent.set_epsilon(adaptive_epsilon as f64); diff --git a/docs/dqn-gpu-hot-path-audit.md b/docs/dqn-gpu-hot-path-audit.md index 221855fde..a732739ec 100644 --- a/docs/dqn-gpu-hot-path-audit.md +++ b/docs/dqn-gpu-hot-path-audit.md @@ -68,6 +68,7 @@ | `training_loop.rs:1003` | `clone_htod_f32` → targets device buffer | COLD-PATH | Data loading: DBN targets uploaded before step loop | OK | | `training_loop.rs:1007` | `clone_htod_f32` → features device buffer | COLD-PATH | Data loading: feature matrix uploaded before step loop | OK | | `tau_update_kernel.cu` (Plan 1 Task 13) | single-thread kernel: reads ISV[39,40,12]; writes ISV[42] | COLD-PATH | Per-epoch boundary call from `training_loop.rs`; 1-thread, 1-block; zero hot-path overhead | OK | +| `epsilon_update_kernel.cu` (Plan 1 Task 14) | single-thread kernel: reads ISV[39,40,12]; writes ISV[41] | COLD-PATH | Per-epoch boundary call from `training_loop.rs`; 1-thread, 1-block; zero hot-path overhead | OK | ## Summary diff --git a/docs/dqn-wire-up-audit.md b/docs/dqn-wire-up-audit.md index abf62b2d6..ec0b60eb1 100644 --- a/docs/dqn-wire-up-audit.md +++ b/docs/dqn-wire-up-audit.md @@ -46,6 +46,8 @@ | `trainers/dqn/monitors/grad_balancer_monitor.rs` | Read-only observer for grad_balance_isv_update kernel output (ISV slots 31..35); consumers: HEALTH_DIAG + controller_activity smoke | Wired | Plan 1 Task 17 | — | | `trainers/dqn/monitors/tau_monitor.rs` | Read-only observer for tau_update kernel output (ISV slot 42); consumers: HEALTH_DIAG + controller_activity smoke | Wired | Plan 1 Task 13 | — | | `cuda_pipeline/tau_update_kernel.cu` | `GpuDqnTrainer::launch_tau_update` → `FusedTrainingCtx::launch_tau_update` → `training_loop.rs` epoch boundary | Wired | Plan 1 Task 13 — cold-path (per-epoch) | — | +| `trainers/dqn/monitors/epsilon_monitor.rs` | Read-only observer for epsilon_update kernel output (ISV slot 41); consumers: HEALTH_DIAG + controller_activity smoke | Wired | Plan 1 Task 14 | — | +| `cuda_pipeline/epsilon_update_kernel.cu` | `GpuDqnTrainer::launch_epsilon_update` → `FusedTrainingCtx::launch_epsilon_update` → `training_loop.rs` epoch boundary | Wired | Plan 1 Task 14 — cold-path (per-epoch) | — | ## CUDA Pipeline — Rust Wrappers