gitadmin/foxhunt
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

feat(dqn): Plan C A.2 — adaptive Polyak tau coupled to q-drift rate

Browse Source 
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) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-04-29 18:31:36 +02:00
parent fc8dbb0a85
commit b4d4a8d046
5 changed files with 103 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

40
crates/ml/src/cuda_pipeline/tau_update_kernel.cu
View File

@@ -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;
}

15
crates/ml/src/trainers/dqn/state_reset_registry.rs
View File

@@ -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 }
}

5
crates/ml/src/trainers/dqn/trainer/mod.rs
View File

@@ -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);

43
crates/ml/src/trainers/dqn/trainer/training_loop.rs
View File

@@ -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.

2
docs/dqn-wire-up-audit.md
View File

@@ -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).