gitadmin/foxhunt
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

fix: sync after restore_best_params to prevent cross-stream race

Browse Source 
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) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-04-07 12:01:39 +02:00
parent 39c2520fca
commit e7db806d25
1 changed files with 6 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

7
crates/ml/src/trainers/dqn/fused_training.rs
View File

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