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