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

fix(adam): reset all 7 auxiliary Adam states at fold boundary — close partial-refactor gap

Browse Source 
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) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-04-28 19:15:08 +02:00
parent 7c19b59033
commit e1efa8994f
5 changed files with 148 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

40
crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs
View File

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

23
crates/ml/src/cuda_pipeline/gpu_iql_trainer.rs
View File

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

21
crates/ml/src/cuda_pipeline/gpu_tlob.rs
View File

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

24
crates/ml/src/trainers/dqn/fused_training.rs
View File

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

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

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