From 9c3ddf8b37b571e4529cbc1d3d3a7f0f5b0a84a8 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Thu, 23 Apr 2026 17:45:37 +0200 Subject: [PATCH] =?UTF-8?q?fix(dqn):=20hard-copy=20online=E2=86=92target?= =?UTF-8?q?=20at=20fold=20boundaries?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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) --- crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs | 18 ++++++++++++++++++ crates/ml/src/trainers/dqn/fused_training.rs | 8 ++++++++ 2 files changed, 26 insertions(+) diff --git a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs index 61cf7dae7..4414969e2 100644 --- a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs +++ b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs @@ -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::(); + 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, MLError> { let len = self.target_params_buf.len(); diff --git a/crates/ml/src/trainers/dqn/fused_training.rs b/crates/ml/src/trainers/dqn/fused_training.rs index 3d8000d49..2495fff9a 100644 --- a/crates/ml/src/trainers/dqn/fused_training.rs +++ b/crates/ml/src/trainers/dqn/fused_training.rs @@ -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()?;