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(()) }