From b4d4a8d046d6294b8a8aa8a4c846cec7615dcc2a Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Wed, 29 Apr 2026 18:31:36 +0200 Subject: [PATCH] =?UTF-8?q?feat(dqn):=20Plan=20C=20A.2=20=E2=80=94=20adapt?= =?UTF-8?q?ive=20Polyak=20tau=20coupled=20to=20q-drift=20rate?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Wires the q_drift_rate_ema kernel + ISV slot Q_DRIFT_RATE_INDEX=129 (landed in fc8dbb0a8) into the production training path: - StateResetRegistry FoldReset entry isv_q_drift_rate -> 0.0 (companion to A.1's prev_epoch_q_mean reset; both ensure no cross-fold leakage of drift state). - reset_named_state dispatch arm in training_loop.rs. - Per-epoch launch_q_drift_rate_ema call after the Q-drift kill check, gated on the same `prev_epoch_q_mean.abs() > 1e-6` cold-start guard. Both q_mean_curr and q_mean_prev are in scope BEFORE the `prev_epoch_q_mean = q_mean` update, so the producer sees the correct delta. - tau_update_kernel.cu multiplies the cosine-scheduled, health-coupled tau_eff by `1 / (1 + clip(ISV[129], 0, 4))` so tau ranges [tau_base/5, tau_base] — monotone dampening under drift; healthy runs (drift ≈ 0) unaffected. - reset_for_fold comment in trainer/mod.rs notes the registry entry handles the new ISV slot. Predicted impact: Plan C fold 0 ep2 with q_mean(t)=0.82, q_mean(t-1)=-0.018, ISV[16]+ISV[21]≈0.1 -> drift_rate ≈ 8.4 -> clipped to 4 -> tau_eff = tau_base × 0.2 (5× slower target sync, dampens Q-target optimism through the inflation spike). a52d99613 fold 0 with q_mean staying ~0.21 -> drift_rate ≈ 0 -> tau_eff = tau_base (unchanged). Per pearl_adaptive_moe_lambda.md — canonical "EMA-tracked diagnostic drives a controller" pattern. Per pearl_cold_path_no_exception_to_gpu_drives.md — cold-path scalar arithmetic stays on GPU. Per feedback_no_partial_refactor.md — consumer (tau_update + state_reset + producer launch) all migrate together in this commit. Layout fingerprint already shifted by fc8dbb0a8 (slot + ISV_TOTAL_DIM bump); no additional fingerprint shift in this commit. Co-Authored-By: Claude Opus 4.7 (1M context) --- .../ml/src/cuda_pipeline/tau_update_kernel.cu | 40 ++++++++++++++++- .../src/trainers/dqn/state_reset_registry.rs | 15 +++++++ crates/ml/src/trainers/dqn/trainer/mod.rs | 5 +++ .../src/trainers/dqn/trainer/training_loop.rs | 43 +++++++++++++++++++ docs/dqn-wire-up-audit.md | 2 + 5 files changed, 103 insertions(+), 2 deletions(-) diff --git a/crates/ml/src/cuda_pipeline/tau_update_kernel.cu b/crates/ml/src/cuda_pipeline/tau_update_kernel.cu index d345520fb..638187b78 100644 --- a/crates/ml/src/cuda_pipeline/tau_update_kernel.cu +++ b/crates/ml/src/cuda_pipeline/tau_update_kernel.cu @@ -1,11 +1,31 @@ /* tau_update — Polyak-EMA tau coefficient, GPU-driven (Plan 1 Task 13). * Reads: ISV[EPOCH_IDX_INDEX=39], ISV[TOTAL_EPOCHS_INDEX=40], - * ISV[LEARNING_HEALTH_INDEX=12] + * ISV[LEARNING_HEALTH_INDEX=12], ISV[Q_DRIFT_RATE_INDEX=129] * Writes: ISV[TAU_EFF_INDEX=42] * Cold path (per-epoch boundary). Single-thread kernel. * - * Formula: cosine schedule (tau_base -> tau_final) over training, with + * Base formula: cosine schedule (tau_base -> tau_final) over training, with * health-coupled floor that rises to 0.01 at full collapse. + * + * Plan C Phase 2 follow-up A.2 (2026-04-29) — adaptive Polyak-tau drift + * dampening: + * + * drift_factor = 1.0 / (1.0 + clip(ISV[Q_DRIFT_RATE_INDEX], 0, 4)) + * tau_eff = base_clamped × drift_factor + * + * The drift signal is produced by `q_drift_rate_ema_update` once per epoch + * boundary BEFORE this kernel fires. With drift saturated at 4, tau_eff + * drops to `tau_base / 5` (5× slower target sync); with drift = 0, + * `drift_factor = 1` and the schedule passes through unchanged. Monotone — + * tau only ever decreases under drift; healthy runs are unaffected. + * + * Per `feedback_no_partial_refactor.md`: when ISV[129] is freshly written by + * the producer kernel each epoch boundary, every consumer of the tau output + * (Polyak EMA target sync) sees the dampened value in the same step. Plan C + * fold 0 ep2's predicted impact: `q_mean(t)=0.82, q_mean(t-1)=-0.018, + * ISV[16]+ISV[21]≈0.1` → drift_rate ≈ 8.4 → clipped to 4 → drift_factor = + * 0.2 → tau_eff = tau_base × 0.2 (5× slower target sync, dampens Q-target + * optimism through the inflation spike). */ extern "C" __global__ void tau_update( const float* __restrict__ isv, @@ -23,5 +43,21 @@ extern "C" __global__ void tau_update( float floor_v = 0.01f * (1.0f - health); float tau_eff = fmaxf(tau_cos, floor_v); tau_eff = fminf(fmaxf(tau_eff, tau_final), tau_base); + + /* Plan C Phase 2 follow-up A.2: adaptive drift-rate dampening. + * Multiplicative factor in [1/5, 1] applied AFTER the [tau_final, + * tau_base] clamp so dampening can drive tau_eff below tau_final on + * runaway drift (this is the explicit intent — tau_final is the + * design floor for healthy runs, but a runaway requires further + * suppression to stop the bootstrap target chasing the online Q). + * The producer kernel `q_drift_rate_ema_update` clips ISV[129] to + * [0, 4]; the `fminf(..., 4.0f)` here is a defensive belt-and-braces + * bound (Invariant 1 carve-out) in case a future producer-kernel bug + * lets a larger value through. */ + float drift_raw = isv[129]; + float drift_clip = fminf(fmaxf(drift_raw, 0.0f), 4.0f); + float drift_factor = 1.0f / (1.0f + drift_clip); + tau_eff = tau_eff * drift_factor; + isv_out[42] = tau_eff; } diff --git a/crates/ml/src/trainers/dqn/state_reset_registry.rs b/crates/ml/src/trainers/dqn/state_reset_registry.rs index 09f10c928..7cf6e2a73 100644 --- a/crates/ml/src/trainers/dqn/state_reset_registry.rs +++ b/crates/ml/src/trainers/dqn/state_reset_registry.rs @@ -409,6 +409,21 @@ impl StateResetRegistry { category: ResetCategory::FoldReset, description: "ISV[MOE_LAMBDA_EFF_INDEX=128] — adaptive load-balance λ from `moe_lambda_eff_update` controller; reset to config.moe_lambda_floor at fold boundary (gate uniform-init ⇒ deficit=0 ⇒ λ_floor); per-step kernel raises λ_eff up to (floor + max_extra) when gate-entropy EMA decays below entropy_target_frac × ln(K)", }, + // Plan C Phase 2 follow-up A.2 (2026-04-29): adaptive Polyak-tau + // drift signal (slot 129). Reset to 0.0 at fold boundary so the + // new fold starts with no inherited drift dampening — the first + // epoch's tau passes through unchanged (drift_factor = 1) and + // the producer kernel `q_drift_rate_ema_update` only re-fires + // once `DQNTrainer::prev_epoch_q_mean` has graduated above 1e-6 + // (gated in `training_loop.rs`). Companion to A.1's + // prev_epoch_q_mean reset in `reset_for_fold` — both ensure + // cross-fold drift state cannot leak into the new fold's + // tau-dampening computation. + RegistryEntry { + name: "isv_q_drift_rate", + category: ResetCategory::FoldReset, + description: "ISV[Q_DRIFT_RATE_INDEX=129] — adaptive Polyak-tau drift signal ∈ [0, 4]; GPU q_drift_rate_ema_update kernel fills (Plan C A.2). Cold-start 0.0 (no-op dampening factor) reapplied at fold boundary; consumer tau_update_kernel.cu multiplies tau_eff by 1/(1+ISV[129]) so dampening is monotone — tau only ever decreases under drift, healthy runs unaffected", + }, ]; Self { entries } } diff --git a/crates/ml/src/trainers/dqn/trainer/mod.rs b/crates/ml/src/trainers/dqn/trainer/mod.rs index cb561a10a..e6ce45b63 100644 --- a/crates/ml/src/trainers/dqn/trainer/mod.rs +++ b/crates/ml/src/trainers/dqn/trainer/mod.rs @@ -1717,6 +1717,11 @@ impl DQNTrainer { // spuriously when fold N legitimately starts from a different // q-scale (different policy, different replay distribution). // Pairs with q_value_history clear above (the source of q_mean). + // Plan C Phase 2 follow-up A.2 (2026-04-29) companion: the + // `isv_q_drift_rate` registry entry below resets ISV[Q_DRIFT_RATE_ + // INDEX=129] to 0.0 — both must clear together so the new fold's + // first epoch sees no inherited drift dampening (the launch_q_drift + // _rate_ema gate `prev_epoch_q_mean.abs() > 1e-6` skips cold-start). self.prev_epoch_q_mean = 0.0; self.adaptive_tau = f64::from(self.hyperparams.tau); diff --git a/crates/ml/src/trainers/dqn/trainer/training_loop.rs b/crates/ml/src/trainers/dqn/trainer/training_loop.rs index aad7eeb6b..ab7d98efb 100644 --- a/crates/ml/src/trainers/dqn/trainer/training_loop.rs +++ b/crates/ml/src/trainers/dqn/trainer/training_loop.rs @@ -4104,6 +4104,33 @@ impl DQNTrainer { )); } } + + // Plan C Phase 2 follow-up A.2 (2026-04-29): adaptive Polyak-tau + // drift signal producer. Computes + // drift_rate = |q_mean(t) - q_mean(t-1)| / max(|q_mean(t-1)|, + // ISV[Q_ABS_REF] + ISV[Q_DIR_ABS_REF], 1e-6) + // clipped to [0, 4] and writes ISV[Q_DRIFT_RATE_INDEX=129]. The + // next epoch's `launch_tau_update` (called at the top of the next + // iteration of this loop) reads the slot and multiplies the + // cosine-scheduled tau_base by `1/(1+ISV[129])` so the effective + // Polyak rate dampens monotonically as q-drift rises. + // + // Gated on `prev_epoch_q_mean.abs() > 1e-6` — same criterion the + // kill-check above uses to skip the first epoch (no meaningful + // delta). Cold-start ISV[129]=0.0 (constructor + FoldReset), so a + // skipped launch keeps the dampening factor at 1.0 (no-op) until a + // real prev value exists. Per `feedback_no_partial_refactor.md`: + // every consumer of the new contract migrates in this commit. + if self.prev_epoch_q_mean.abs() > 1e-6 { + if let Some(ref fused) = self.fused_ctx { + let q_curr_f = q_mean as f32; + let q_prev_f = self.prev_epoch_q_mean as f32; + if let Err(e) = fused.trainer().launch_q_drift_rate_ema(q_curr_f, q_prev_f) { + tracing::warn!(epoch, "launch_q_drift_rate_ema failed (non-fatal): {e}"); + } + } + } + self.prev_epoch_q_mean = q_mean; // Action diversity — per-branch from GPU monitoring (not re-derived through OrderRouter) @@ -5192,6 +5219,22 @@ impl DQNTrainer { ); } } + "isv_q_drift_rate" => { + // Plan C Phase 2 follow-up A.2 (2026-04-29): adaptive + // Polyak-tau drift signal. Reset to 0.0 at fold boundary so + // the new fold's first `tau_update` consumer sees a no-op + // dampening factor of 1.0 (`drift_factor = 1/(1+0) = 1`) + // until `q_drift_rate_ema_update` fires with a meaningful + // delta. Companion to A.1's `prev_epoch_q_mean = 0.0` + // reset in `DQNTrainer::reset_for_fold` — both ensure no + // cross-fold leakage of drift state. + if let Some(ref fused) = self.fused_ctx { + fused.trainer().write_isv_signal_at( + crate::cuda_pipeline::gpu_dqn_trainer::Q_DRIFT_RATE_INDEX, + 0.0_f32, + ); + } + } "isv_gamma_dir_eff" | "isv_gamma_mag_eff" | "isv_gamma_ord_eff" | "isv_gamma_urg_eff" => { // D.2 per-branch gamma slots. Reset to 0.0 at fold boundary; // per_branch_gamma_update GPU kernel re-populates on next epoch-boundary launch. diff --git a/docs/dqn-wire-up-audit.md b/docs/dqn-wire-up-audit.md index a5356460f..bf16e1874 100644 --- a/docs/dqn-wire-up-audit.md +++ b/docs/dqn-wire-up-audit.md @@ -2190,3 +2190,5 @@ Plan C Phase 2 follow-up A.2 (2026-04-29): adaptive Polyak-tau coupled to per-ep Predicted impact: Plan C fold 0 ep2 with `q_mean(t)=0.82, q_mean(t-1)=-0.018, ISV[16]+ISV[21]≈0.1` → drift_rate ≈ |0.84|/max(0.018, 0.1, ε) = 8.4 → clipped to 4 → tau_eff = tau_base × 0.2 (5× slower target sync, dampens Q-target optimism through inflation spike). a52d99613 fold 0 with q_mean staying ~0.21 → drift_rate ≈ 0 → tau_eff = tau_base (unchanged, healthy run unaffected). Per `pearl_adaptive_moe_lambda.md` — canonical "EMA-tracked diagnostic drives a controller" pattern: kernel + ISV slot + bootstrap (0.0 cold-start) + reset (0.0 fold boundary) + observability. Per `pearl_cold_path_no_exception_to_gpu_drives.md` — even cold-path scalar arithmetic stays on GPU when its reduction inputs already live on GPU (ISV[16,21] both produced by `q_stats_kernel.cu`). Per `feedback_adaptive_not_tuned.md` — tau decay is ISV-driven, not a constant. Per `feedback_isv_for_adaptive_bounds.md` — the `4.0` upper clip and `1e-6` denom floor are structural bounds (drift saturation matching the kill-criterion's 3.0× ratio threshold; numerical-stability floor) carried in ISV slot semantics. + +Plan C A.2 wire-up (2026-04-29, second commit): `tau_update_kernel.cu`, `state_reset_registry.rs`, `training_loop.rs::reset_named_state` + per-epoch launch site, and `trainer/mod.rs::reset_for_fold` comment all updated in lockstep with the kernel+slot landed by the previous commit. Producer launch lives between the existing Q-drift kill criterion and the `prev_epoch_q_mean = q_mean` update — both prev and curr q_means are in scope, and the `prev_epoch_q_mean.abs() > 1e-6` cold-start gate matches the kill check's gate. Tau consumer applies `1/(1+clip(ISV[129],0,4))` AFTER the existing `[tau_final, tau_base]` clamp so dampening can drive tau below tau_final on runaway drift (intent: tau_final is the design floor for healthy runs, but a runaway requires further suppression). Defensive `fminf(...,4.0f)` in the consumer guards against producer-side bugs even though the producer already clips (Invariant 1 carve-out, belt-and-braces).