jgrusewski fa1a94bf9e fix(dqn): reset IQN Adam state at fold boundaries ...

GpuIqnHead carries its own m_buf/v_buf/adam_step — separate from
GpuDqnTrainer::reset_adam_state, which only zeroes the legacy
iqn_trunk_* buffers. Without a fold-boundary reset, the IQN
optimizer enters fold N+1 with fold N's momentum, producing
oversized Adam steps that compound through the IQN backward
pass into the runaway gradients we observed in both train-7rgqd
(crashed fold 1 ep 52) and train-5gzpn (NaN'd fold 1 ep 17).

- Add GpuIqnHead::reset_adam_state() — zero m_buf, v_buf,
  adam_step, and the pinned t counter.
- Call it from FusedTraining::reset_for_fold after the trainer's
  main Adam reset, gated by gpu_iqn.is_some(). Non-fatal warn on
  failure to match the surrounding shrink-and-perturb pattern.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

2026-04-23 17:16:52 +02:00

93 B

Raw Blame History

View Raw