From 7c19b59033d85ea2e2ea211573e9311d9eb6dd14 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Tue, 28 Apr 2026 19:08:51 +0200 Subject: [PATCH] =?UTF-8?q?fix(iqn):=20hard-sync=20target=20from=20online?= =?UTF-8?q?=20at=20fold=20boundary=20=E2=80=94=20kills=20geometric=20Q-dri?= =?UTF-8?q?ft=20in=20fold=201?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add `GpuIqnHead::sync_target_from_online` (DtoD copy of online_params → target_params) and call it from `FusedTrainingCtx::reset_for_fold` alongside the existing IQN Adam reset. Root cause: at fold boundary the main DQN does shrink-and-perturb + hard target sync (gpu_dqn_trainer.rs:12495 — the latter explicitly added to prevent fold-1 grad explosion). The IQN head was added later but only had Polyak EMA `target_ema_update` — no hard-sync existed. Polyak at τ=0.005/step closes ~0.5%/step, far too slow to close the fold-boundary online↔target gap before inflated Bellman TD-error compounds geometrically through the IQN backward pass. Bug signature this resolves: F0 ep5 Q=+0.79 (healthy) F1 ep1 Q=+0.82 (boundary fine; c51_alpha warmup blends IQN low) F1 ep2 Q=+2.23 (drift starts — c51_alpha ramp + IQN online↔target gap widens) F1 ep3 Q=+4.05 F1 ep4 Q=+10.66 (geometric ~2.3×/epoch) F1 ep5 NaN at step 5 (fp32 overflow past atom support) Per `feedback_no_partial_refactor.md`: when a shared contract (fold-boundary state migration) changes, every consumer must migrate together. The main DQN sync was added explicitly — the IQN, TLOB, IQL, and aux Adam states were added later but never extended to that contract. This commit closes the IQN gap (primary). Aux Adam + IQN readiness gaps follow in subsequent commits. Touched: - gpu_iqn_head.rs: new `sync_target_from_online` (mirrors gpu_dqn_trainer.rs:12495 exactly) - fused_training.rs: call the new sync inside the existing IQN reset block - docs/dqn-wire-up-audit.md: Invariant 7 entry covering Fix 1 of 3 cargo check -p ml --lib clean at 13 warnings (workspace baseline). No fingerprint change. DtoD memcpy only — no HtoD/HtoH/DtoH paths. Co-Authored-By: Claude Opus 4.7 (1M context) --- crates/ml/src/cuda_pipeline/gpu_iqn_head.rs | 22 +++++++++++++ crates/ml/src/trainers/dqn/fused_training.rs | 15 +++++++++ docs/dqn-wire-up-audit.md | 33 ++++++++++++++++++++ 3 files changed, 70 insertions(+) diff --git a/crates/ml/src/cuda_pipeline/gpu_iqn_head.rs b/crates/ml/src/cuda_pipeline/gpu_iqn_head.rs index 669baf6eb..895fad73a 100644 --- a/crates/ml/src/cuda_pipeline/gpu_iqn_head.rs +++ b/crates/ml/src/cuda_pipeline/gpu_iqn_head.rs @@ -1493,6 +1493,28 @@ impl GpuIqnHead { Ok(0.0) } + /// Hard-copy online IQN params into target params via DtoD. + /// + /// Mirrors `GpuDqnTrainer::sync_target_from_online`: at fold boundaries the + /// main DQN trunk is shrink-and-perturb'd and Adam reset, but the IQN + /// target buffer still holds end-of-prior-fold weights. Polyak EMA at + /// τ=0.005/step is far too slow to close that gap before the inflated + /// Bellman TD error compounds through the IQN backward pass — this drives + /// the geometric Q-drift observed in fold-1 (Q ≈ 0.8 → 2.2 → 4.0 → 10.7 + /// → fp32 overflow past atom support). Hard-syncing target ← online at + /// the same fold boundary as the main DQN closes the gap in a single + /// DtoD copy. No CPU staging — pure on-GPU. + pub fn sync_target_from_online(&mut self) -> Result<(), MLError> { + let n_bytes = self.total_params * std::mem::size_of::(); + let src = self.online_params.raw_ptr(); + let dst = self.target_params.raw_ptr(); + unsafe { + cudarc::driver::result::memcpy_dtod_async(dst, src, n_bytes, self.stream.cu_stream()) + .map_err(|e| MLError::ModelError(format!("IQN sync_target_from_online: {e}")))?; + } + Ok(()) + } + /// Update target IQN weights via Polyak EMA. /// /// `target[i] = (1 - tau) * target[i] + tau * online[i]` diff --git a/crates/ml/src/trainers/dqn/fused_training.rs b/crates/ml/src/trainers/dqn/fused_training.rs index f7e9e267f..fe29f7e47 100644 --- a/crates/ml/src/trainers/dqn/fused_training.rs +++ b/crates/ml/src/trainers/dqn/fused_training.rs @@ -896,6 +896,21 @@ impl FusedTrainingCtx { // producing oversized Adam steps in the first epochs whose effect // compounds through the IQN backward pass into runaway gradients. if let Some(iqn) = self.gpu_iqn.as_mut() { + // Hard-sync IQN target ← online at the same fold boundary as the + // main DQN's `sync_target_from_online` above. Without this, the + // IQN target buffer retains end-of-prior-fold weights while the + // online side is otherwise advancing through fold N+1's first + // gradient steps; Polyak EMA at τ=0.005 closes ~0.5%/step, far + // too slow before the inflated TD-error compounds geometrically + // (observed: Q ≈ 0.8 → 2.2 → 4.0 → 10.7 → fp32 overflow past + // atom support inside fold 1). Pairs with the IQN Adam reset + // immediately below — both are fold-boundary contract consumers + // that must migrate together (see `feedback_no_partial_refactor`). + if let Err(e) = iqn.sync_target_from_online() { + tracing::warn!("Fold-boundary IQN target sync failed (non-fatal): {e}"); + } else { + tracing::info!("Fold-boundary IQN target hard-synced from online"); + } if let Err(e) = iqn.reset_adam_state() { tracing::warn!("Fold-boundary IQN Adam reset failed (non-fatal): {e}"); } else { diff --git a/docs/dqn-wire-up-audit.md b/docs/dqn-wire-up-audit.md index c43fba508..6bdb4ebd6 100644 --- a/docs/dqn-wire-up-audit.md +++ b/docs/dqn-wire-up-audit.md @@ -2,6 +2,39 @@ **Status:** Populated during Plan 1 Task 6 (A.5 orphan audit). Updated on every commit per Invariant 7. +Fold-boundary reset gap — IQN target hard-sync (2026-04-28, Fix 1 of 3): +the main DQN's `sync_target_from_online` (DtoD `online_params → +target_params` at `gpu_dqn_trainer.rs:12495`) is wired into +`FusedTrainingCtx::reset_for_fold` (`fused_training.rs:887`) explicitly +to prevent fold-1 grad explosion. The IQN head, added later, only +exposed `target_ema_update` (Polyak τ=0.005/step) — no hard-sync. After +fold boundary the IQN online side advances through fold N+1's first +gradient steps while the IQN target buffer still holds end-of-prior- +fold weights, so the Bellman TD error gap widens by a factor that +Polyak takes hundreds of steps to close. Empirically that drives a +geometric ~2.3×/epoch Q-drift inside fold 1 +(`F0 ep5 Q=+0.79 → F1 ep1 +0.82 → ep2 +2.23 → ep3 +4.05 → ep4 +10.66 +→ ep5 NaN @ step 5` on fp32 overflow past atom support). The boundary +handoff at F1 ep1 looks fine because the c51_alpha warmup ramp blends +IQN low; once c51_alpha ramps up the IQN online↔target gap dominates. +**Fix**: add `GpuIqnHead::sync_target_from_online` (DtoD copy mirroring +the main DQN's helper exactly — same `memcpy_dtod_async` pattern, same +stream, same total_params) and invoke it from +`FusedTrainingCtx::reset_for_fold` immediately before the existing IQN +Adam reset. Per `feedback_no_partial_refactor.md`: the +fold-boundary state-migration contract has multiple consumers (main +DQN trunk, IQN head, aux Adam states, IQN readiness EMAs) and +extending it must migrate every consumer in lockstep. The IQN target +sync is the primary root cause; the aux Adam state resets and IQN +readiness reset land in the next two commits. + +Touched: `gpu_iqn_head.rs` (new `sync_target_from_online` method +between `reset_adam_state` and `target_ema_update`), `fused_training.rs` +(invoke the new sync inside the existing `if let Some(iqn) = …` block +in `reset_for_fold`, immediately before `iqn.reset_adam_state()`). +cargo check -p ml --lib clean at 13 warnings (workspace baseline). No +fingerprint change. DtoD only — no HtoD/HtoH/DtoH paths introduced. + MSE warmup loss target clamp to C51 atom support (2026-04-28): the fold-boundary PopArt carry-forward + iqr floor (`d710f9d50`) reduced fold-1 epoch-1 train loss readings 16× (318k → 19k) but ep-2/ep-3