From e1efa8994f7bab4a7db4c473e2a629ae09d8bce9 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Tue, 28 Apr 2026 19:15:08 +0200 Subject: [PATCH] =?UTF-8?q?fix(adam):=20reset=20all=207=20auxiliary=20Adam?= =?UTF-8?q?=20states=20at=20fold=20boundary=20=E2=80=94=20close=20partial-?= =?UTF-8?q?refactor=20gap?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Extend the fold-boundary Adam reset contract to every auxiliary optimizer added after the main DQN's `reset_adam_state` was wired: - `GpuDqnTrainer::reset_adam_state` now also memsets q_attn / sel / denoise / mamba2 / ofi_embed (m, v) buffer pairs and zeroes the corresponding adam_step counters. - New `pub(crate) fn reset_adam_state` on `GpuTlob` (memsets adam_m/v, zeroes adam_step + writes 0 to the device-mapped pinned host counter so the next adam_step kernel reads t=1 after `increment_adam_step`). - New `pub fn reset_adam_state` on `GpuIqlTrainer` (same pattern, m_buf/v_buf/adam_step/t_pinned). - `FusedTrainingCtx::reset_for_fold` invokes the three new helpers immediately after the existing IQN Adam reset block. Both `gpu_iql` and `gpu_iql_low` are reset. Per `feedback_no_partial_refactor.md` — the fold-boundary state-migration contract has multiple consumers (main DQN trunk, IQN head, aux Adam states across q_attn / sel / denoise / mamba2 / ofi_embed / tlob / iql / iql_low). Adding `sync_target_from_online` to the main DQN years ago without extending it to subsequently- added optimizers left a "partial migration strictly worse than the original" state. This commit + Fix 1 + Fix 3 close all three remaining gaps surfaced by the fold-boundary audit. All new resets follow the existing main-DQN pattern: DtoD memset- zero only, pinned host counter written through the device-mapped page (no HtoD/HtoH/DtoH paths). cargo check -p ml --lib clean at 13 warnings (workspace baseline). cargo test -p ml --lib --no-run clean. Co-Authored-By: Claude Opus 4.7 (1M context) --- .../ml/src/cuda_pipeline/gpu_dqn_trainer.rs | 40 +++++++++++++++++++ .../ml/src/cuda_pipeline/gpu_iql_trainer.rs | 23 +++++++++++ crates/ml/src/cuda_pipeline/gpu_tlob.rs | 21 ++++++++++ crates/ml/src/trainers/dqn/fused_training.rs | 24 +++++++++++ docs/dqn-wire-up-audit.md | 40 +++++++++++++++++++ 5 files changed, 148 insertions(+) diff --git a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs index 6e25d9acf..1e4d907bb 100644 --- a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs +++ b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs @@ -3442,6 +3442,46 @@ impl GpuDqnTrainer { .map_err(|e| MLError::ModelError(format!("reset iqn t: {e}")))?; self.iqn_trunk_adam_step = 0; + // Auxiliary Adam states — reset alongside the main Adam state so + // every per-fold optimizer enters fold N+1 with fresh momentum. + // Cross-Branch Q-Attention, selectivity gate, diffusion denoiser, + // Mamba2 update, and OFI embed all carry their own (m, v, step) + // tuple separate from the main `m_buf`/`v_buf`. Without this + // reset the auxiliary optimizers enter fold N+1 with fold N's + // momentum, producing oversized Adam steps in the first epochs + // whose effect compounds through the corresponding backward path. + // Closes the partial-refactor gap diagnosed in the fold-boundary + // audit (companion to the IQN target sync + IQN Adam reset). + self.stream.memset_zeros(&mut self.q_attn_adam_m) + .map_err(|e| MLError::ModelError(format!("reset q_attn_adam_m: {e}")))?; + self.stream.memset_zeros(&mut self.q_attn_adam_v) + .map_err(|e| MLError::ModelError(format!("reset q_attn_adam_v: {e}")))?; + self.q_attn_adam_step = 0; + + self.stream.memset_zeros(&mut self.sel_adam_m) + .map_err(|e| MLError::ModelError(format!("reset sel_adam_m: {e}")))?; + self.stream.memset_zeros(&mut self.sel_adam_v) + .map_err(|e| MLError::ModelError(format!("reset sel_adam_v: {e}")))?; + self.sel_adam_step = 0; + + self.stream.memset_zeros(&mut self.denoise_adam_m) + .map_err(|e| MLError::ModelError(format!("reset denoise_adam_m: {e}")))?; + self.stream.memset_zeros(&mut self.denoise_adam_v) + .map_err(|e| MLError::ModelError(format!("reset denoise_adam_v: {e}")))?; + self.denoise_adam_step = 0; + + self.stream.memset_zeros(&mut self.mamba2_adam_m) + .map_err(|e| MLError::ModelError(format!("reset mamba2_adam_m: {e}")))?; + self.stream.memset_zeros(&mut self.mamba2_adam_v) + .map_err(|e| MLError::ModelError(format!("reset mamba2_adam_v: {e}")))?; + self.mamba2_adam_step = 0; + + self.stream.memset_zeros(&mut self.ofi_embed_adam_m) + .map_err(|e| MLError::ModelError(format!("reset ofi_embed_adam_m: {e}")))?; + self.stream.memset_zeros(&mut self.ofi_embed_adam_v) + .map_err(|e| MLError::ModelError(format!("reset ofi_embed_adam_v: {e}")))?; + self.ofi_embed_adam_step = 0; + // Reset PopArt running statistics — prevents fold 1's reward distribution // from contaminating fold 2's normalization. self.stream.memset_zeros(&mut self.popart_mean) diff --git a/crates/ml/src/cuda_pipeline/gpu_iql_trainer.rs b/crates/ml/src/cuda_pipeline/gpu_iql_trainer.rs index 11cb4770b..1a3197d6b 100644 --- a/crates/ml/src/cuda_pipeline/gpu_iql_trainer.rs +++ b/crates/ml/src/cuda_pipeline/gpu_iql_trainer.rs @@ -1114,6 +1114,29 @@ impl GpuIqlTrainer { self.total_loss_buf.raw_ptr() } + /// Reset Adam optimizer state (m, v, step) at fold boundary. + /// + /// IQL carries its own Adam state separate from the main DQN. Without + /// this reset the IQL value-head optimizer enters fold N+1 with fold + /// N's momentum, producing oversized Adam steps in the first epochs + /// whose effect compounds through the expectile-loss backward pass — + /// degrading the per-sample support / advantage-weight stream that + /// the main C51/IQN updates depend on. Closes the partial-refactor + /// gap diagnosed in the fold-boundary audit (companion to the IQN + /// target sync + main-DQN aux Adam resets). Pinned host counter is + /// written through the device-mapped page so the next `adam_step` + /// kernel reads `t = 1` after the next `increment_adam_step`, + /// matching the in-place pattern used by the main DQN trainer. + pub fn reset_adam_state(&mut self) -> Result<(), MLError> { + self.stream.memset_zeros(&mut self.m_buf) + .map_err(|e| MLError::ModelError(format!("reset iql m_buf: {e}")))?; + self.stream.memset_zeros(&mut self.v_buf) + .map_err(|e| MLError::ModelError(format!("reset iql v_buf: {e}")))?; + self.adam_step = 0; + unsafe { *self.t_pinned = 0; } + Ok(()) + } + /// Current Adam step count. pub fn adam_step(&self) -> i32 { self.adam_step diff --git a/crates/ml/src/cuda_pipeline/gpu_tlob.rs b/crates/ml/src/cuda_pipeline/gpu_tlob.rs index ab31951e4..f408bbf22 100644 --- a/crates/ml/src/cuda_pipeline/gpu_tlob.rs +++ b/crates/ml/src/cuda_pipeline/gpu_tlob.rs @@ -712,6 +712,27 @@ impl GpuTlob { Ok(()) } + /// Reset Adam optimizer state (m, v, step) at fold boundary. + /// + /// TLOB carries its own Adam state separate from the main DQN. Without + /// this reset the TLOB optimizer enters fold N+1 with fold N's + /// momentum, producing oversized Adam steps in the first epochs whose + /// effect compounds through the QKV/output projections into the trunk + /// gradient. Closes the partial-refactor gap diagnosed in the + /// fold-boundary audit (companion to the IQN target sync + main-DQN + /// aux Adam resets). Pinned host counter is written through the + /// device-mapped page so the next `adam_step` launch reads `t = 1`, + /// matching the in-place pattern used by the main DQN trainer. + pub(crate) fn reset_adam_state(&mut self) -> Result<(), MLError> { + self.stream.memset_zeros(&mut self.adam_m) + .map_err(|e| MLError::ModelError(format!("reset tlob adam_m: {e}")))?; + self.stream.memset_zeros(&mut self.adam_v) + .map_err(|e| MLError::ModelError(format!("reset tlob adam_v: {e}")))?; + self.adam_step = 0; + unsafe { *self.t_pinned = 0; } + Ok(()) + } + /// Copy trained weight parameters from `src` into this instance via DtoD async copy. /// /// Used to synchronise the val-path TLOB (owned by `GpuBacktestEvaluator`) with diff --git a/crates/ml/src/trainers/dqn/fused_training.rs b/crates/ml/src/trainers/dqn/fused_training.rs index fe29f7e47..2671a7bc7 100644 --- a/crates/ml/src/trainers/dqn/fused_training.rs +++ b/crates/ml/src/trainers/dqn/fused_training.rs @@ -918,6 +918,30 @@ impl FusedTrainingCtx { } } + // TLOB and IQL each carry their own Adam state separate from the + // main DQN's `m_buf` / `v_buf` (and from the IQN head's). The + // fold-boundary contract that resets the main DQN Adam must also + // reset every auxiliary optimizer, otherwise fold N's momentum + // drives oversized first-epoch updates in fold N+1's expectile + + // attention backward paths and feeds into the trunk gradient. + // Per `feedback_no_partial_refactor.md` — same shared contract, + // every consumer migrates together. + if let Err(e) = self.gpu_tlob.reset_adam_state() { + tracing::warn!("Fold-boundary TLOB Adam reset failed (non-fatal): {e}"); + } else { + tracing::info!("Fold-boundary TLOB Adam state reset"); + } + if let Err(e) = self.gpu_iql.reset_adam_state() { + tracing::warn!("Fold-boundary IQL Adam reset failed (non-fatal): {e}"); + } else { + tracing::info!("Fold-boundary IQL Adam state reset"); + } + if let Err(e) = self.gpu_iql_low.reset_adam_state() { + tracing::warn!("Fold-boundary IQL-low Adam reset failed (non-fatal): {e}"); + } else { + tracing::info!("Fold-boundary IQL-low Adam state reset"); + } + // FoldReset state (eval_q_mean_ema, eval_q_std_ema, isv_v_range_slots, // isv_learning_health, isv_sharpe_ema, isv_q_means) is now handled by // StateResetRegistry-driven dispatch in DQNTrainer::reset_for_fold diff --git a/docs/dqn-wire-up-audit.md b/docs/dqn-wire-up-audit.md index 6bdb4ebd6..2c401e93e 100644 --- a/docs/dqn-wire-up-audit.md +++ b/docs/dqn-wire-up-audit.md @@ -2,6 +2,46 @@ **Status:** Populated during Plan 1 Task 6 (A.5 orphan audit). Updated on every commit per Invariant 7. +Fold-boundary reset gap — auxiliary Adam state resets (2026-04-28, +Fix 2 of 3): the same fold-boundary contract that resets the main +DQN's `m_buf` / `v_buf` / `adam_step` (gpu_dqn_trainer.rs:3428) was +never extended to the seven auxiliary optimizers added later — five +co-located inside `GpuDqnTrainer` (q_attn at lines 1961-1962, sel +at 1969-1970, denoise at 3061-3063, mamba2 at 3175-3176, ofi_embed +at 3347-3349) plus two on separate types (`GpuTlob` adam_m/v/step +at gpu_tlob.rs:143-145, `GpuIqlTrainer` m_buf/v_buf/adam_step at +gpu_iql_trainer.rs:215-216,253). Without the matching reset each of +these enters fold N+1 with fold N's momentum, producing oversized +Adam steps in the first epochs whose effect compounds through the +corresponding backward pass — Q-attention into the trunk, IQL into +the per-sample-support / advantage-weight stream that C51 + IQN +consume, TLOB into the QKV/output projections feeding trunk +gradient. Fold-1 grad-explosion symptoms are dominated by IQN +target lag (Fix 1) but the aux Adam states are on the same shared +fold-boundary contract (`feedback_no_partial_refactor.md`) and +must migrate together. **Fix**: extend +`GpuDqnTrainer::reset_adam_state` with `memset_zeros` of the five +aux (m, v) buffer pairs and step counter resets; add +`pub(crate) fn reset_adam_state` to `GpuTlob` (memsets adam_m/v, +zeroes adam_step + t_pinned via the device-mapped page) and +`pub fn reset_adam_state` to `GpuIqlTrainer` (same pattern, m_buf/ +v_buf/adam_step/t_pinned). Wire all three from +`FusedTrainingCtx::reset_for_fold` immediately after the existing +IQN reset block — both `gpu_iql` and `gpu_iql_low` call the new +helper. All three new helpers strictly mirror the main DQN's +`reset_adam_state`: DtoD memset-zero only, pinned host counter +written through the device-mapped page (no HtoD/HtoH/DtoH). + +Touched: `gpu_dqn_trainer.rs` (`reset_adam_state` extended with five +aux blocks, no signature change), `gpu_tlob.rs` (new +`pub(crate) fn reset_adam_state` after `adam_step`), `gpu_iql_trainer.rs` +(new `pub fn reset_adam_state` before the `adam_step()` getter), +`fused_training.rs` (three new invocations after the existing IQN +Adam reset block in `reset_for_fold`). cargo check -p ml --lib clean +at 13 warnings (workspace baseline). cargo test -p ml --lib --no-run +clean. No fingerprint change. No new HtoD/HtoH/DtoH paths — only +DtoD memsets and existing pinned-write patterns. + 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