fix(dqn): hard-copy online→target at fold boundaries

target_params_buf was initialized once via DtoD copy at first train
step and then only moved toward online via slow Polyak EMA (tau≈0.005).
At fold boundaries the online weights are shrink-and-perturb'd with
alpha=0.8, which modifies params_buf in-place — but target_params_buf
still held the end-of-previous-fold values. The Bellman target would
then use stale weights against freshly perturbed online predictions,
producing a large TD error gap in the first fold-N+1 training steps.
Polyak averaging at tau=0.005 is far too slow to close that gap before
the oversized gradients compound through Adam into runaway updates —
one of the drivers of the fold-1 gradient explosion observed in both
train-7rgqd and train-5gzpn.

- Add GpuDqnTrainer::sync_target_from_online() — DtoD memcpy of the
  full params_buf into target_params_buf.
- Call it from FusedTraining::reset_for_fold right after shrink-and-
  perturb and before reset_adam_state, so target = perturbed online
  and Adam moments zero out from the same starting point.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-04-23 17:45:37 +02:00
parent fa1a94bf9e
commit 9c3ddf8b37
2 changed files with 26 additions and 0 deletions

View File

@@ -9145,6 +9145,24 @@ impl GpuDqnTrainer {
Ok(dst)
}
/// Hard-copy online params into target params. At fold boundaries the online
/// weights are shrink-and-perturb'd; the EMA target still holds the end-of-
/// prior-fold values and would otherwise produce a large TD error gap in the
/// first steps of the new fold, driving runaway gradients before Polyak
/// averaging can close the divergence. DtoD copy — no CPU staging.
pub fn sync_target_from_online(&mut self) -> Result<(), MLError> {
let n_bytes = self.total_params * std::mem::size_of::<f32>();
unsafe {
cudarc::driver::result::memcpy_dtod_async(
self.target_params_buf.raw_ptr(),
self.params_buf.raw_ptr(),
n_bytes,
self.stream.cu_stream(),
).map_err(|e| MLError::ModelError(format!("sync_target_from_online DtoD: {e}")))?;
}
Ok(())
}
/// Clone target_params_buf to a new CudaSlice via DtoD copy — checkpoint stays on GPU, zero CPU.
pub fn clone_target_params_gpu(&self) -> Result<CudaSlice<f32>, MLError> {
let len = self.target_params_buf.len();

View File

@@ -830,6 +830,14 @@ impl FusedTrainingCtx {
} else {
tracing::info!(alpha = sp_alpha, sigma = sp_sigma, "Fold-boundary shrink-and-perturb applied");
}
// Hard-copy the shrink-and-perturb'd online weights into target params.
// Without this, target_params_buf retains end-of-previous-fold weights
// while online was just modified — the Bellman target would use stale
// weights against perturbed online predictions, producing an outsized
// TD error gap in the first fold-N+1 steps. Polyak averaging alone is
// too slow (tau=0.005) to close that gap before the oversized gradients
// compound through Adam. This pairs with reset_adam_state below.
self.trainer.sync_target_from_online()?;
// Reset Adam optimizer state — each fold starts with fresh momentum.
// Stale momentum from a previous fold causes weight explosion → Q-value -1e30.
self.trainer.reset_adam_state()?;