From e7db806d25902cadd91c92a1c69e694577b2656e Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Tue, 7 Apr 2026 12:01:39 +0200 Subject: [PATCH] fix: sync after restore_best_params to prevent cross-stream race MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The evaluator uses its own CUDA stream, but restore_best_params writes weights on the trainer's stream. Without sync, the evaluator's cuBLAS forward reads partially-written weights → CUBLAS_STATUS_EXECUTION_FAILED → CUDA_ERROR_ILLEGAL_ADDRESS → all subsequent trials fail. Added stream.synchronize() after the DtoD copy + unflatten to guarantee weights are globally visible before the evaluator reads them. Co-Authored-By: Claude Opus 4.6 (1M context) --- crates/ml/src/trainers/dqn/fused_training.rs | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/crates/ml/src/trainers/dqn/fused_training.rs b/crates/ml/src/trainers/dqn/fused_training.rs index 2f65b682f..9fd52716f 100644 --- a/crates/ml/src/trainers/dqn/fused_training.rs +++ b/crates/ml/src/trainers/dqn/fused_training.rs @@ -1514,8 +1514,10 @@ impl FusedTrainingCtx { Ok(()) } - /// Restore best-epoch params into online weights. Pure async DtoD + unflatten. + /// Restore best-epoch params into online weights. DtoD + unflatten + sync. /// Called before evaluation to ensure the evaluator sees the best model. + /// Syncs the stream to guarantee writes are globally visible before the + /// evaluator (which may use a different stream) reads them. pub(crate) fn restore_best_params(&mut self) -> Result<()> { let snapshot = self.best_params_snapshot.as_ref() .ok_or_else(|| anyhow::anyhow!("No best params snapshot saved"))?; @@ -1530,6 +1532,9 @@ impl FusedTrainingCtx { // Unflatten: params_bf16 → individual weight tensors self.trainer.unflatten_online_weights(&self.online_dueling, &self.online_branching) .map_err(|e| anyhow::anyhow!("restore_best_params unflatten: {e}"))?; + // Sync: ensure DtoD + unflatten complete before evaluator reads on its stream + self.stream.synchronize() + .map_err(|e| anyhow::anyhow!("restore_best_params sync: {e}"))?; Ok(()) }