Audit of all Adam optimizer state in crates/ml/src/ against the existing
reset_adam_state calls at the fold-boundary in
fused_training.rs::reset_for_fold:
- DQN main (line 971) — covers m_buf/v_buf, IQN-trunk, q_attn, sel,
denoise, mamba2, ofi_embed, PopArt, Q-div EMA via
GpuDqnTrainer::reset_adam_state. Also covers VSN params (in main
params buffer slots) and aux next_bar/regime heads (slots 119-126).
- IQN head (line 1011)
- TLOB (line 1026)
- IQL-high (line 1031)
- IQL-low (line 1036)
Missing wire identified and added:
- GpuAttention (4-head feature attention on h_s2): owns a separate
(attn_m, attn_v, attn_adam_step) tuple in gpu_attention.rs, called
every training step via attn.adam_step at fused_training.rs:1962/1995,
but was never zeroed at fold boundary. Same pathology as IQN/TLOB/IQL
— fold N momentum oversizes fold N+1's first-epoch SDP + output
projection updates, compounding through the trunk gradient.
Implementation:
- Add GpuAttention::reset_adam_state mirroring the gpu_iqn_head and
gpu_tlob pattern (memset_zeros m/v + zero step counter + write 0
through pinned t_pinned for next adam_step launch).
- Wire it under `if let Some(ref mut attn) = self.gpu_attention` in
reset_for_fold, immediately after the IQL-low reset, with the same
warn-on-failure / info-on-success log pattern.
- Append SP3 Task B5 entry to docs/dqn-wire-up-audit.md documenting
the audit + fix (Invariant 7 contract).
Out-of-scope optimizers (documented in audit, not changed):
- DecisionTransformer: scoped within DT pretrain block, dropped at end
of pretrain (no fold leak possible).
- GpuCuriosityTrainer: owned by GpuExperienceCollector, external to
FusedTrainingCtx. Reset belongs at the collector layer if needed.
- MetaQNetwork: host-side observability MLP for collapse prediction
(does not feed LearningHealth or trunk gradients).
- GpuMoeHead: forward-only, no Adam state of its own.
Comprehensive coverage now per feedback_no_partial_refactor — every
in-context Adam state buffer zeros at fold transition. SP3 Mech 4
closes the "Adam EMA accumulation drives weight pathology" pathway.
ef429c25d8