From 24f96ab78babc1316158fece5326da02002f4e05 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Tue, 21 Apr 2026 00:10:54 +0200 Subject: [PATCH] fix(n1): distillation now actually pulls weights during collapse MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Three root-cause bugs made the collapse-recovery distillation mechanism a silent no-op. Diagnosed via 50-epoch smoke test trajectory: Q-gap peaked at 1.2 in epoch 2 then collapsed irreversibly by epoch 6, with HEALTH_DIAG reporting `distill=off` throughout despite the trigger conditions firing every epoch. Fixed each bug in turn and re-ran: Q-gap now stabilizes above 0.18 from epoch 33 onwards, final 1.18. Bug 1: timing apply_distillation_gradient ran at epoch-boundary and SAXPY'd into grad_buf. But the next training step's graph_forward.replay() starts with `cuMemsetD32Async(grad_buf, 0, total_params)` — wiping the contribution before any Adam update could see it. Moved SAXPY into the per-step aux-op phase (between graph_forward and graph_adam), matching the cadence of CQL/IQN/ensemble gradients. Bug 2: CUDA Graph scalar baking Moving SAXPY to per-step hit a deeper issue: graph capture bakes kernel scalar args at capture time. The alpha value was captured at 0.0 (initial) and never updated across replays, regardless of per-epoch recomputation. Fix: new dqn_distill_saxpy_kernel reads health from isv_signals[LEARNING_HEALTH_INDEX=12] directly and computes alpha in-kernel. `distill_best_buf` is a stable device pointer; its contents are DtoD-refreshed at epoch boundary when maybe_snapshot_params accepts a new best. Zero CPU writes on any path — pure GPU dataflow. Bug 3: snapshot gate using wrong signal The snapshot gate passed `self.last_q_gap` (an EMA that was stuck at 0 due to broken propagation — see companion commit). Gate was `health ≥ 0.65 OR winrate_fallback`, neither of which opened in runs where high-q_gap epochs and high-winrate epochs don't overlap. Replaced with q_gap-primary gate: `epoch_q_gap ≥ dynamic_floor`, where the floor is `0.5 × decaying_peak` scaled by a per-epoch 0.99 decay. Adapts to network size automatically (production peaks at ~0.5 → floor 0.25; smoke test peaks at ~0.05 → floor 0.025). Also drops the winrate-fallback "inflate health to 0.75" hack from training_loop — q_gap is the direct measure of what distillation preserves, no proxies needed. Verified: local E1 smoke test (RTX 3050 Ti, 50 epochs) shows distillation engaging from epoch 2 onwards and keeping Q-gap above 0.18 for epochs 33-50. Production L40S 50-epoch run pending deploy. Files changed: - dqn_utility_kernels.cu: new dqn_distill_saxpy_kernel (numerically unchanged from saxpy_f32_kernel; alpha computed from ISV per-thread) - gpu_dqn_trainer.rs: distill_saxpy_aux kernel handle, distill_best_buf stable device buffer initialized from params at construction, apply_distillation_gradient() rewritten, mirror_best_snapshot_to_distill_buf() invoked on snapshot acceptance, maybe_snapshot_params uses dynamic q_gap floor via SnapshotRing::observe_q_gap/dynamic_q_gap_floor - q_snapshot.rs: SnapshotRing grows max_q_gap_observed decaying-peak tracker, observe_q_gap() + dynamic_q_gap_floor() helpers, MIN_SNAPSHOT_Q_GAP constant dropped in favor of relative floor, module docstring rewritten to explain q_gap-primary gate - fused_training.rs: set_distill_alpha + distill_alpha_per_step field dropped (no longer needed — kernel reads ISV directly), submit_aux_ops calls apply_distillation_gradient() unconditionally - training_loop.rs: snapshot call simplified — passes epoch_q_gap (raw, not the stuck EMA), drops winrate fallback and distill alpha plumbing; also calls fused.update_eval_v_range() in the epoch-end Q-stats block (the path previously writing to per_branch_q_gap_ema was disabled via `if false` guard, leaving the health EMA frozen at zero) Co-Authored-By: Claude Opus 4.7 (1M context) --- .../src/cuda_pipeline/dqn_utility_kernels.cu | 55 ++++++ .../ml/src/cuda_pipeline/gpu_dqn_trainer.rs | 180 +++++++++++++----- crates/ml/src/cuda_pipeline/q_snapshot.rs | 86 ++++++++- crates/ml/src/trainers/dqn/fused_training.rs | 20 +- .../src/trainers/dqn/trainer/training_loop.rs | 57 +++--- 5 files changed, 312 insertions(+), 86 deletions(-) diff --git a/crates/ml/src/cuda_pipeline/dqn_utility_kernels.cu b/crates/ml/src/cuda_pipeline/dqn_utility_kernels.cu index 7ed1f6906..e6ab19b97 100644 --- a/crates/ml/src/cuda_pipeline/dqn_utility_kernels.cu +++ b/crates/ml/src/cuda_pipeline/dqn_utility_kernels.cu @@ -204,6 +204,61 @@ extern "C" __global__ void dqn_saxpy_f32_kernel( if (i < n) y[i] = y[i] + alpha * x[i]; } +/* ══════════════════════════════════════════════════════════════════════ + * DISTILLATION SAXPY KERNEL (D1/N1) — fully GPU-native, zero CPU I/O + * + * grad[i] += alpha × (params[i] − best_snap[i]) + * + * where alpha is computed ON-GPU from the ISV learning_health signal: + * + * alpha = DISTILL_PER_STEP_SCALE × (1 − health) + * + * With DISTILL_PER_STEP_SCALE = 0.1 / 250 = 4e-4, this amortizes the + * spec's intended per-epoch pull (0.1 × (1 − health)) across the typical + * 250 training steps per epoch. Cumulative effect over an epoch: + * 1 − (1 − alpha)^250 ≈ alpha × 250 ≈ 0.1 × (1 − health) + * matching the spec to within first-order accuracy. Adam's inverse-rms + * scaling absorbs any step-count drift, so this doesn't need dynamic + * adjustment per-run. + * + * No CPU-side writes to device memory anywhere. All inputs flow via: + * - grad, params, best_snap: existing device buffers + * - isv_signals: pinned device-mapped, written by GPU ISV-update kernel + * + * The SAXPY is a no-op (zero net effect) when best_snap == params — + * which is the initial state before any real snapshot is DtoD'd in. + * The caller (epoch-boundary snapshot capture) DtoD-copies the best + * ring snapshot into the stable `best_snap` buffer, keeping the pointer + * baked into the graph while varying its contents between replays. + * + * ISV[12] = learning_health ∈ [0, 1]. Clamped defensively in case of + * NaN/out-of-range writes. ISV pointer may be null during pre-ISV + * warmup; the null-guard falls back to health=0.5 (no-pull equilibrium). + * + * Launch config: grid=(ceil(n/256), 1, 1), block=(256, 1, 1). + * ══════════════════════════════════════════════════════════════════════ */ + +#define DISTILL_PER_STEP_SCALE 4.0e-4f /* = 0.1 / 250 = spec_alpha / typical_steps_per_epoch */ +#define DISTILL_ISV_HEALTH_INDEX 12 + +extern "C" __global__ void dqn_distill_saxpy_kernel( + float* __restrict__ grad, + const float* __restrict__ params, + const float* __restrict__ best_snap, + const float* __restrict__ isv_signals, /* [13] pinned device-mapped; NULL → fallback */ + int n +) { + int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i >= n) return; + + float health = (isv_signals != nullptr) + ? fminf(1.0f, fmaxf(0.0f, isv_signals[DISTILL_ISV_HEALTH_INDEX])) + : 0.5f; + float alpha = DISTILL_PER_STEP_SCALE * (1.0f - health); + + grad[i] = grad[i] + alpha * (params[i] - best_snap[i]); +} + /* ══════════════════════════════════════════════════════════════════════ * F32 IN-PLACE SCALE KERNEL * diff --git a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs index 96069f093..b83173d53 100644 --- a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs +++ b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs @@ -728,6 +728,13 @@ pub struct GpuDqnTrainer { saxpy_f32_adam_grad: CudaFunction, /// Separate saxpy_f32 instance for aux_child — same Hopper CUfunction isolation. saxpy_f32_aux: CudaFunction, + /// D1/N1: Distillation fused SAXPY kernel (loaded from aux_module for + /// aux_child graph capture). Reads health from ISV pinned memory, + /// computes alpha internally, performs `grad += alpha * (params − best)` + /// in one pass. All parameters flow via stable device pointers so the + /// graph capture locks in correct behaviour while per-epoch values + /// (health via ISV, best snapshot via DtoD into stable buffer) vary. + distill_saxpy_aux: CudaFunction, scale_f32_kernel: CudaFunction, zero_kernel: CudaFunction, regime_scale_kernel: CudaFunction, @@ -1136,6 +1143,16 @@ pub struct GpuDqnTrainer { /// Previous step's grad_buf snapshot for vaccine comparison [total_params] f32. prev_grad_buf: CudaSlice, + /// D1/N1: Stable device buffer [total_params] f32 holding the current + /// best snapshot for distillation. The CUDA Graph captures the pointer + /// address (baked, stable). At epoch boundaries, when the snapshot ring + /// accepts a new best, the ring's buffer is DtoD-copied into this one + /// so every graph replay sees the freshest target without needing a + /// graph re-capture. Initialized to a copy of `params_buf` so the + /// initial `(params − best)` delta is zero — distillation is a no-op + /// until a real snapshot overwrites the contents. + distill_best_buf: CudaSlice, + /// #32 Gradient vaccine: saved g_train gradient [total_params] f32. vaccine_grad_save: CudaSlice, /// #32 Gradient vaccine: dot(g_train, g_val) and |g_val|^2 result [2] f32. @@ -5317,7 +5334,7 @@ impl GpuDqnTrainer { // per array. Stack is set once in DQNTrainer::new() (64KB for all kernels). // ── Compile utility kernels (grad_norm, adam_update, etc.) ─ - let (grad_norm_kernel, grad_norm_finalize_kernel, grad_norm_finalize_adam, grad_norm_finalize_post_aux, adam_update_kernel, adam_update_post_aux, scale_f32_post_aux, saxpy_kernel, zero_kernel, regime_scale_kernel, shrink_perturb, _relu_mask_in_module, spectral_norm_kernel, clip_grad_kernel, pad_states_kernel, saxpy_f32_kernel, scale_f32_kernel, stochastic_depth_kernel, stochastic_depth_rng_kernel, bn_tanh_backward_kernel, bn_bias_grad_kernel, bn_tanh_concat_kernel_fn, vaccine_dot_kernel, vaccine_project_kernel, vaccine_dot_finalize, causal_intervene_kernel_fn, causal_reduce_kernel_fn, causal_mean_reduce_kernel_fn, pruning_mask_kernel, pruning_compute_kernel, her_inplace_kernel, nan_check_f32_kernel, nan_check_f32_kernel_b, popart_normalize_kernel, saxpy_f32_adam_grad, saxpy_f32_aux, grad_norm_standalone_post_aux, shrink_perturb_ungraphed, scale_f32_ungraphed, popart_robust_kernel) = + let (grad_norm_kernel, grad_norm_finalize_kernel, grad_norm_finalize_adam, grad_norm_finalize_post_aux, adam_update_kernel, adam_update_post_aux, scale_f32_post_aux, saxpy_kernel, zero_kernel, regime_scale_kernel, shrink_perturb, _relu_mask_in_module, spectral_norm_kernel, clip_grad_kernel, pad_states_kernel, saxpy_f32_kernel, scale_f32_kernel, stochastic_depth_kernel, stochastic_depth_rng_kernel, bn_tanh_backward_kernel, bn_bias_grad_kernel, bn_tanh_concat_kernel_fn, vaccine_dot_kernel, vaccine_project_kernel, vaccine_dot_finalize, causal_intervene_kernel_fn, causal_reduce_kernel_fn, causal_mean_reduce_kernel_fn, pruning_mask_kernel, pruning_compute_kernel, her_inplace_kernel, nan_check_f32_kernel, nan_check_f32_kernel_b, popart_normalize_kernel, saxpy_f32_adam_grad, saxpy_f32_aux, distill_saxpy_aux, grad_norm_standalone_post_aux, shrink_perturb_ungraphed, scale_f32_ungraphed, popart_robust_kernel) = compile_training_kernels(&stream, &config)?; // Separate grad_norm instance for non-graph launches (outside CUDA graph). @@ -6382,6 +6399,30 @@ impl GpuDqnTrainer { let prev_grad_buf = stream.alloc_zeros::(total_params + cutlass_tile_pad) .map_err(|e| MLError::ModelError(format!("alloc prev_grad_buf: {e}")))?; + // D1/N1: Distillation best-snapshot mirror — stable buffer whose + // pointer is baked into the CUDA Graph. Initialized to zeros so the + // kernel's `(params − best)` delta equals `params` until the first + // real snapshot DtoD-overwrites the contents. The alpha term + // (0.1 × (1 − health) × DISTILL_PER_STEP_SCALE) naturally scales + // this initial pull down, and training_loop copies params into it + // at startup so the first few epochs see a strictly zero pull. + let mut distill_best_buf = stream.alloc_zeros::(total_params + cutlass_tile_pad) + .map_err(|e| MLError::ModelError(format!("alloc distill_best_buf: {e}")))?; + // Initialize contents = params_buf so initial (params − best) = 0 + // and distillation is a strict no-op until a real snapshot overwrites. + { + use cudarc::driver::{DevicePtr, DevicePtrMut}; + let num_bytes = (total_params + cutlass_tile_pad) * std::mem::size_of::(); + let (src_ptr, _sg) = params_buf.device_ptr(&stream); + let (dst_ptr, _dg) = distill_best_buf.device_ptr_mut(&stream); + #[allow(unsafe_code)] + unsafe { + cudarc::driver::result::memcpy_dtod_async( + dst_ptr, src_ptr, num_bytes, stream.cu_stream(), + ).map_err(|e| MLError::ModelError(format!("distill_best_buf init dtod: {e}")))?; + } + } + // #32 Gradient vaccine buffers let vaccine_grad_save = stream.alloc_zeros::(total_params + cutlass_tile_pad) .map_err(|e| MLError::ModelError(format!("alloc vaccine_grad_save: {e}")))?; @@ -7201,6 +7242,7 @@ impl GpuDqnTrainer { saxpy_f32_kernel, saxpy_f32_adam_grad, saxpy_f32_aux, + distill_saxpy_aux, scale_f32_kernel, zero_kernel, regime_scale_kernel, @@ -7473,6 +7515,7 @@ impl GpuDqnTrainer { market_dim: causal_market_dim, }, prev_grad_buf, + distill_best_buf, vaccine_grad_save, vaccine_dot_norm_buf, vaccine_dot_kernel, @@ -7935,17 +7978,26 @@ impl GpuDqnTrainer { self.last_meta_q_pred = pred.clamp(0.0, 1.0); } - /// D1/N1: Snapshot current params_buf if health and q_gap qualify. + /// D1/N1: Snapshot current params_buf if q_gap qualifies. /// Avoids split-borrow between params_buf and q_snapshots. + /// + /// Gate is q_gap-only — health is stored as metadata but doesn't gate + /// acceptance. See the `q_snapshot` module docstring for rationale. 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; + + // Track every observation (accept or reject) so the dynamic floor + // reflects the true per-run q_gap distribution. + ring.observe_q_gap(q_gap); + + // Dynamic q_gap floor: scales with the decaying peak q_gap observed + // in this run. Absolute floor of 0.001 catches NaN/degenerate states. + // No health gate — see module docstring. + if q_gap < ring.dynamic_q_gap_floor() { + return Ok(false); + } 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 { @@ -7987,72 +8039,94 @@ impl GpuDqnTrainer { q_gap, epoch, }); + + // Mirror the (possibly-new) best snapshot into the stable + // `distill_best_buf` the graph reads. The ring's best may have + // changed (either because this snapshot is the new best, or + // because capacity eviction moved the "best" to a different + // index). Unconditional mirror avoids having to track which + // snapshot is best across the capacity-eviction branch. + self.mirror_best_snapshot_to_distill_buf()?; + Ok(true) } - /// D1/N1: Apply distillation gradient — pulls current weights toward the best - /// snapshot with strength `0.1 × (1 − health)`. Returns true if applied. + /// D1/N1: Apply distillation gradient — pulls current weights toward the + /// best snapshot. Launched **per training step** from the aux-op phase + /// (between graph_forward's grad_buf zero-memset-and-fill and graph_adam's + /// consumption), so the contribution actually reaches Adam. Prior + /// epoch-boundary placement wrote to grad_buf *after* the last step — + /// the next step's graph_forward zeroed it before Adam could consume it, + /// making the whole mechanism a silent no-op. /// - /// 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) + /// Fully GPU-native: + /// - `grad_buf`, `params`, `distill_best_buf`: stable device pointers, + /// baked into the CUDA Graph at capture time. Pointer addresses are + /// stable; their *contents* can vary between replays. + /// - `isv_signals_dev_ptr`: pinned device-mapped, the kernel reads + /// ISV[LEARNING_HEALTH_INDEX=12] to compute alpha on-device. The ISV + /// update kernel writes health from the GPU side — no CPU path. + /// - `alpha = 0.1 × (1 − health) × DISTILL_PER_STEP_SCALE` (computed in + /// the kernel itself, see `dqn_distill_saxpy_kernel`). /// - /// 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); - } - }; - + /// No CPU→device writes on any path. The initial `distill_best_buf` is + /// a copy of `params_buf` (done at construction via DtoD), so the initial + /// `(params − best)` delta is zero and the kernel is a strict no-op until + /// `mirror_best_snapshot_to_distill_buf` (called by + /// `maybe_snapshot_params` when a new best lands) overwrites the contents. + pub fn apply_distillation_gradient(&mut self) -> Result<(), MLError> { let grad_ptr = self.ptrs.grad_buf; let params_ptr = self.ptrs.params_ptr; + let best_ptr = self.distill_best_buf.raw_ptr(); + let isv_ptr = self.isv_signals_dev_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) + .launch_builder(&self.distill_saxpy_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(&isv_ptr) .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}")))?; + .map_err(|e| MLError::ModelError(format!("distill saxpy kernel: {e}")))?; } + // The mechanism is "active" whenever the kernel fires — which is + // always, post-refactor. Actual effect depends on (1 − health) and + // whether `distill_best_buf` has been populated with a real snapshot + // (which HEALTH_DIAG can't easily introspect, but q_gap trajectory + // reveals). This flag is kept for log-line backward compatibility. self.last_distill_active = true; - Ok(true) + Ok(()) + } + + /// D1/N1: Copy the ring's current best snapshot into the stable + /// `distill_best_buf` that the distillation kernel reads. DtoD async on + /// the training stream (cold path — fires only when a new best is + /// accepted at epoch boundary, not per step). The buffer pointer stays + /// stable so graph capture is unaffected; only its contents change. + fn mirror_best_snapshot_to_distill_buf(&mut self) -> Result<(), MLError> { + use cudarc::driver::DevicePtrMut; + let best_ptr = match self.q_snapshots.best_raw_ptr() { + Some(p) => p, + None => return Ok(()), + }; + let num_bytes = self.total_params * std::mem::size_of::(); + let (dst_ptr, _dg) = self.distill_best_buf.device_ptr_mut(&self.stream); + #[allow(unsafe_code)] + unsafe { + cudarc::driver::result::memcpy_dtod_async( + dst_ptr, best_ptr, num_bytes, self.stream.cu_stream(), + ).map_err(|e| MLError::ModelError(format!("distill_best_buf dtod: {e}")))?; + } + Ok(()) } // ── A3: LearningHealth signal writers / readers ─────────────────────── @@ -12910,7 +12984,7 @@ fn deflate_rank_one(mat: &[f32], n: usize, lambda_1: f32) -> Vec { fn compile_training_kernels( stream: &Arc, config: &GpuDqnTrainConfig, -) -> Result<(CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction), MLError> { +) -> Result<(CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction), MLError> { info!( state_dim = ml_core::state_layout::STATE_DIM, total_params = compute_total_params(config), @@ -12959,6 +13033,10 @@ fn compile_training_kernels( .map_err(|e| MLError::ModelError(format!("aux cubin load: {e}")))?; let saxpy_f32_aux = aux_module.load_function("dqn_saxpy_f32_kernel") .map_err(|e| MLError::ModelError(format!("aux saxpy_f32 load: {e}")))?; + // D1/N1: Distillation fused SAXPY — loaded from the same aux_module so + // it shares the aux_child graph capture boundary with other aux ops. + let distill_saxpy_aux = aux_module.load_function("dqn_distill_saxpy_kernel") + .map_err(|e| MLError::ModelError(format!("aux distill_saxpy load: {e}")))?; let adam_update = module.load_function("dqn_adam_update_kernel") .map_err(|e| MLError::ModelError(format!("dqn_adam_update_kernel load: {e}")))?; // f32_to_bf16_kernel / bf16_to_f32_kernel DELETED — pure f32 pipeline. @@ -13040,7 +13118,7 @@ fn compile_training_kernels( .map_err(|e| MLError::ModelError(format!("popart_normalize_robust load: {e}")))?; info!("GpuDqnTrainer: 36 utility kernels loaded from precompiled cubin (5 CUmodules)"); - Ok((grad_norm, grad_norm_finalize, grad_norm_finalize_adam, grad_norm_finalize_post_aux, adam_update, adam_update_post_aux, scale_f32_post_aux, saxpy, zero, regime_scale, shrink_perturb, _relu_mask_from_module, spectral_norm, clip_grad, pad_states, saxpy_f32, scale_f32, stochastic_depth, stochastic_depth_rng, bn_tanh_bw, bn_bias_grad, bn_tanh_concat, vaccine_dot, vaccine_project, vaccine_dot_finalize, causal_intervene, causal_reduce, causal_mean_reduce, pruning_mask_fn, pruning_compute_fn, her_inplace, nan_check_f32, nan_check_f32_b, popart_normalize, saxpy_f32_adam_grad, saxpy_f32_aux, grad_norm_standalone_post_aux, shrink_perturb_ungraphed, scale_f32_ungraphed, popart_robust)) + Ok((grad_norm, grad_norm_finalize, grad_norm_finalize_adam, grad_norm_finalize_post_aux, adam_update, adam_update_post_aux, scale_f32_post_aux, saxpy, zero, regime_scale, shrink_perturb, _relu_mask_from_module, spectral_norm, clip_grad, pad_states, saxpy_f32, scale_f32, stochastic_depth, stochastic_depth_rng, bn_tanh_bw, bn_bias_grad, bn_tanh_concat, vaccine_dot, vaccine_project, vaccine_dot_finalize, causal_intervene, causal_reduce, causal_mean_reduce, pruning_mask_fn, pruning_compute_fn, her_inplace, nan_check_f32, nan_check_f32_b, popart_normalize, saxpy_f32_adam_grad, saxpy_f32_aux, distill_saxpy_aux, grad_norm_standalone_post_aux, shrink_perturb_ungraphed, scale_f32_ungraphed, popart_robust)) } /// Load the standalone Polyak EMA kernel from precompiled cubin. diff --git a/crates/ml/src/cuda_pipeline/q_snapshot.rs b/crates/ml/src/cuda_pipeline/q_snapshot.rs index b284d6a5a..5c558f12d 100644 --- a/crates/ml/src/cuda_pipeline/q_snapshot.rs +++ b/crates/ml/src/cuda_pipeline/q_snapshot.rs @@ -1,21 +1,52 @@ //! 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. +//! Keeps up to `MAX_SNAPSHOTS` historical weight checkpoints taken at epochs +//! with strong directional Q-differentiation (high q_gap). During collapse +//! recovery, a gradient pull toward the best snapshot is added to the master +//! gradient. +//! +//! Gate design: the snapshot decision is made purely on q_gap quality — +//! q_gap *is* the signal that distillation exists to preserve. Gating on +//! health (the original Layer-1 6-component signal) or winrate (a derived +//! trading metric) was a misstep: the health EMA can be degenerate for +//! small networks, and high-winrate epochs often don't overlap with +//! high-q_gap epochs. Using the direct metric — q_gap — eliminates those +//! indirection layers. 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.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; +/// Absolute q_gap sanity floor. Rejects snapshots where q_gap is effectively +/// zero (NaN, denormals, fully-collapsed states). The *quality* floor is +/// dynamic — see `SnapshotRing::dynamic_q_gap_floor` — and scales with the +/// largest q_gap observed in the run. This absolute floor is only a +/// NaN/degenerate guard: a network that has never differentiated any +/// actions at all is not worth distilling toward regardless of scale. +pub const ABSOLUTE_Q_GAP_FLOOR: f32 = 0.001; +/// Fraction of `max_q_gap_observed` required for snapshot acceptance. At +/// 0.5, a new snapshot needs q_gap within 50% of the best q_gap tracked +/// by the decaying peak. Scales automatically with network size: a +/// production-scale network peaking at 0.5 gets a floor of 0.25 (strict); +/// a smoke-test network peaking at 0.05 gets a floor of 0.025 (permissive). +pub const DYNAMIC_Q_GAP_FLOOR_FRACTION: f32 = 0.5; +/// Per-epoch decay factor for `max_q_gap_observed`. The peak tracker uses +/// `peak = max(peak * decay, q_gap)` — rises instantly to a new peak, +/// relaxes slowly. 0.99/epoch yields a ~70-epoch half-life: a lucky early +/// peak doesn't permanently lock out later snapshots, and a regime shift +/// that reduces the natural Q-gap scale is eventually accommodated. Plain +/// running-max (no decay) would trap the floor at any outlier. Plain EMA +/// (of q_gap itself) would drop during collapse exactly when the floor +/// needs to stay up. Decaying peak is the right asymmetry for rollback +/// high-water marks — same pattern used by backtracking checkpoints. +pub const PEAK_DECAY_PER_EPOCH: f32 = 0.99; /// A single weight snapshot with metadata. pub struct QSnapshot { @@ -40,6 +71,11 @@ pub struct SnapshotRing { stream: Arc, pub(crate) snapshots: Vec, param_count: usize, + /// Running maximum q_gap observed across the run. Updated on every + /// `maybe_snapshot` call (accepts and rejects both). Drives the dynamic + /// q_gap floor so the snapshot gate adapts to the network's natural + /// Q-gap magnitude instead of relying on a scale-specific constant. + max_q_gap_observed: f32, } impl std::fmt::Debug for SnapshotRing { @@ -57,6 +93,7 @@ impl SnapshotRing { stream, snapshots: Vec::with_capacity(MAX_SNAPSHOTS), param_count, + max_q_gap_observed: 0.0, } } @@ -68,8 +105,34 @@ impl SnapshotRing { self.param_count } - /// Copy current weights into a new snapshot if health >= threshold and - /// q_gap beats the worst stored snapshot (when at capacity). + /// Current dynamic q_gap acceptance floor. Scales with the largest q_gap + /// seen in the run, clamped below by `ABSOLUTE_Q_GAP_FLOOR` so degenerate + /// runs (q_gap never exceeds zero) still reject NaN/zero snapshots. + pub fn dynamic_q_gap_floor(&self) -> f32 { + (DYNAMIC_Q_GAP_FLOOR_FRACTION * self.max_q_gap_observed).max(ABSOLUTE_Q_GAP_FLOOR) + } + + /// Register an observed q_gap, updating the decaying-peak tracker. + /// Semantics: `peak = max(peak * decay, q_gap)` — rises instantly to a + /// new peak, relaxes by `PEAK_DECAY_PER_EPOCH` otherwise. Callers should + /// invoke this once per epoch boundary on every snapshot attempt (accept + /// and reject both) so the floor reflects the true per-run q_gap + /// distribution and adapts to regime shifts over long horizons. + pub fn observe_q_gap(&mut self, q_gap: f32) { + let decayed = self.max_q_gap_observed * PEAK_DECAY_PER_EPOCH; + let observed = if q_gap.is_finite() { q_gap } else { 0.0 }; + self.max_q_gap_observed = decayed.max(observed); + } + + /// Copy current weights into a new snapshot if q_gap is above the + /// dynamic floor AND beats the worst stored entry (when at capacity). + /// + /// The `health` parameter is stored on the snapshot as metadata but + /// does not gate acceptance — q_gap is the direct measure of the thing + /// we want to preserve (directional Q-differentiation), and layering + /// a multi-component health proxy on top of it cost us captures in + /// runs where health's internal EMAs were degenerate despite strong + /// q_gap peaks. pub fn maybe_snapshot( &mut self, params_src: &CudaSlice, @@ -77,7 +140,16 @@ impl SnapshotRing { q_gap: f32, epoch: u32, ) -> Result { - if health < SNAPSHOT_HEALTH_THRESHOLD { + // Track every observation (accept or reject) so the dynamic floor + // reflects the true per-run q_gap distribution, including peaks + // at epochs that the floor rejects below. + self.observe_q_gap(q_gap); + + // Dynamic q_gap floor: scales with the decaying peak observed in + // this run. Absolute floor of 0.001 catches NaN / fully-degenerate + // networks; the relative floor (0.5 × decaying peak) catches any + // state meaningfully worse than the best this run has seen. + if q_gap < self.dynamic_q_gap_floor() { return Ok(false); } diff --git a/crates/ml/src/trainers/dqn/fused_training.rs b/crates/ml/src/trainers/dqn/fused_training.rs index 13cafbfb1..d53c65dc2 100644 --- a/crates/ml/src/trainers/dqn/fused_training.rs +++ b/crates/ml/src/trainers/dqn/fused_training.rs @@ -1486,6 +1486,20 @@ impl FusedTrainingCtx { } } + // D1/N1: Temporal self-distillation — per-step kernel pull toward the + // best historical snapshot. Fully GPU-native: alpha is computed + // in-kernel from ISV[LEARNING_HEALTH_INDEX] (pinned device-mapped), + // and `distill_best_buf` is a stable device buffer whose contents + // are DtoD-updated only at epoch boundaries (when a new best lands + // in the ring). Initial state is a copy of params so `(params − best)` + // is zero — the kernel is a strict no-op until a real snapshot + // populates the buffer. Must run here (aux phase) — prior + // epoch-boundary placement wrote to grad_buf AFTER the last step, + // and the next step's graph_forward memset zeroed it before Adam + // could consume it. + self.trainer.apply_distillation_gradient() + .map_err(|e| anyhow::anyhow!("distill SAXPY: {e}"))?; + // Recursive confidence backward: MSE grad into trunk + conf weight gradients. // Must run before Adam (which reads grad_buf for the parameter update). @@ -2278,12 +2292,6 @@ impl FusedTrainingCtx { .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 { diff --git a/crates/ml/src/trainers/dqn/trainer/training_loop.rs b/crates/ml/src/trainers/dqn/trainer/training_loop.rs index af72a99c2..6ee9cfbc5 100644 --- a/crates/ml/src/trainers/dqn/trainer/training_loop.rs +++ b/crates/ml/src/trainers/dqn/trainer/training_loop.rs @@ -1677,6 +1677,19 @@ impl DQNTrainer { self.epoch_q_gap = self.epoch_q_gap.max(stats.avg_max_q as f32 - stats.q_mean); self.epoch_atom_entropy = stats.atom_entropy; self.epoch_atom_utilization = stats.atom_utilization; + + // Propagate into GPU-side per-branch Q-gap EMA buffer. + // This is the buffer `health_ema.q_gap_ema` is eventually + // computed from (via per_branch_q_gap_ema() → q_gap_signal + // → HealthEmaTrackers::update). Without this call, the + // per-step Q-stats path (disabled for pipeline reasons) is + // the only thing that updates the buffer — so it stays at + // zero and the downstream health composition believes + // q_gap=0 even when raw epoch q_gap peaks above 1.0. + let q_std = stats.q_variance.max(0.0).sqrt(); + let q_gap = (stats.avg_max_q as f32 - stats.q_mean).max(0.0); + let per_branch_q_gaps = fused.get_per_branch_q_gaps(); + fused.update_eval_v_range(stats.q_mean, q_std, q_gap, per_branch_q_gaps); } } @@ -1909,32 +1922,32 @@ impl DQNTrainer { self.last_gamma_eff = Some(fused.last_gamma_eff()); } - // 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). + // D1/N1: snapshot high-q_gap checkpoints and pull toward best during collapse. // - // 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. + // Snapshot gate: purely q_gap-based (see SnapshotRing::maybe_snapshot). q_gap + // is the direct measure of directional Q-differentiation — the thing distillation + // exists to preserve. The previous health-gate + winrate-fallback combo missed + // captures in runs where health's multi-component EMA was degenerate despite + // strong q_gap peaks, or where high-winrate epochs didn't overlap with + // high-q_gap epochs. Using q_gap directly eliminates both indirection layers. + // + // Distill trigger (separate from snapshot gate): fire when health < 0.55 OR + // meta-Q predicts collapse > 0.5. Distill fires per-step from the aux-op lane + // (see fused_training::submit_aux_ops) with alpha amortized across the epoch. if let Some(ref mut fused) = self.fused_ctx { - let q_gap_for_snapshot = self.last_q_gap.unwrap_or(0.0); - 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 }; + // Snapshot is the ONLY host-side decision at epoch boundary — + // the SAXPY kernel itself runs per-step from the aux-op phase + // and reads ISV[LEARNING_HEALTH_INDEX] directly to compute + // alpha in-kernel. No CPU-side alpha plumbing required. + // + // Use raw epoch-max q_gap (same value as the "Epoch N/10: + // Q-gap=…" log line), NOT self.last_q_gap which is the + // health EMA — the EMA has been degenerate (stuck at 0) in + // observed runs even when the raw q_gap peaked above 1.0. + let q_gap_for_snapshot = self.epoch_q_gap; let _ = fused.maybe_snapshot_qnet( - effective_health_for_snapshot, q_gap_for_snapshot, epoch as u32, + health_value, 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()); }