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

fix: IQN readiness update moved outside graph capture — cuStreamSynchronize invalidated capture

Browse Source 
read_total_loss() calls cuStreamSynchronize which is illegal during
CUDA Graph stream capture. Was inside submit_aux_ops (captured in
graph_mega/graph_aux). Moved to Step 4 (conditional ops outside graph).

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-04-14 15:26:53 +02:00
parent 7c00793c9f
commit d06a18ef9f
1 changed files with 7 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -1001,6 +1001,13 @@ impl FusedTrainingCtx {
}
self.pending_vaccine_batch = None;
// ── Adaptive IQN lambda: read loss OUTSIDE graph capture ─────────
if let Some(ref iqn) = self.gpu_iqn {
if let Ok(loss) = iqn.read_total_loss() {
self.trainer.update_iqn_readiness(loss);
}
}
// ── Clip grad_buf L2 norm BEFORE Adam (outside graph) ───────────
// Root cause: backward weight gradient dW = activation^T × d_output.
// Spectral norm bounds W, IS-weight clamp bounds d_logits (~2.0),
@@ -1272,11 +1279,6 @@ impl FusedTrainingCtx {
dqn_actions, dqn_rewards, dqn_dones,
) {
Ok(_) => {
// Update adaptive IQN lambda from current loss
if let Ok(loss) = iqn.read_total_loss() {
self.trainer.update_iqn_readiness(loss);
}
let d_h_s2_ptr = iqn.d_h_s2_raw_ptr();
self.trainer.apply_iqn_trunk_gradient(
d_h_s2_ptr, &mut self.online_dueling,