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:
@@ -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();
|
||||
|
||||
@@ -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()?;
|
||||
|
||||
Reference in New Issue
Block a user