From d06a18ef9f1bed74c52ca375d7c013bb6e80b092 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Tue, 14 Apr 2026 15:26:53 +0200 Subject: [PATCH] =?UTF-8?q?fix:=20IQN=20readiness=20update=20moved=20outsi?= =?UTF-8?q?de=20graph=20capture=20=E2=80=94=20cuStreamSynchronize=20invali?= =?UTF-8?q?dated=20capture?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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) --- crates/ml/src/trainers/dqn/fused_training.rs | 12 +++++++----- 1 file changed, 7 insertions(+), 5 deletions(-) diff --git a/crates/ml/src/trainers/dqn/fused_training.rs b/crates/ml/src/trainers/dqn/fused_training.rs index 2ff145c58..74df8bb36 100644 --- a/crates/ml/src/trainers/dqn/fused_training.rs +++ b/crates/ml/src/trainers/dqn/fused_training.rs @@ -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,