From e8f2df34f0ef759f641a045de235ecd0ed2d4230 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Thu, 2 Apr 2026 01:55:16 +0200 Subject: [PATCH] =?UTF-8?q?fix:=20first-call-only=20sync=20for=20EMA=20and?= =?UTF-8?q?=20attention=20=E2=80=94=20unblocks=20training=20loop?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The full cuStreamSynchronize removal hung the training loop because cudarc's EventTrackingGuard + check_err need one initial sync to clear stale events from CUDA Graph capture. After the first call, stream ordering is sufficient. - target_ema_update: sync on first call only (target_params_initialized) - apply_attention_forward: sync on first call only (attention_initialized) - All subsequent steps: zero sync, pure async kernel dispatch Co-Authored-By: Claude Opus 4.6 (1M context) --- .../ml/src/cuda_pipeline/gpu_dqn_trainer.rs | 22 +++++++++++-------- 1 file changed, 13 insertions(+), 9 deletions(-) diff --git a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs index 99cf28c90..914038133 100644 --- a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs +++ b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs @@ -581,6 +581,7 @@ pub struct GpuDqnTrainer { total_params: usize, params_initialized: bool, target_params_initialized: bool, + attention_initialized: bool, // ── CUDA Graphs (split: forward+backward, then optimizer) ────── // Graph A: zero → cuBLAS forward → C51 loss → C51 grad → cuBLAS backward @@ -1990,11 +1991,15 @@ impl GpuDqnTrainer { let sh2 = self.config.shared_h2; let n_bytes = batch_size * sh2 * f32_size; - // Stream ordering ensures graph replay completed before touching save_h_s2. - // Disable event tracking for all buffer access in this method. // After CUDA Graph capture, cudarc's device_ptr() records write events // that fail with CUDA_ERROR_INVALID_VALUE on stale captured graph events. + // First call syncs to clear stale events; subsequent calls rely on stream ordering. + if !self.attention_initialized { + unsafe { cudarc::driver::sys::cuStreamSynchronize(self.stream.cu_stream()); } + let _ = self.stream.context().check_err(); + self.attention_initialized = true; + } let _evt_guard = EventTrackingGuard::new(self.stream.context()); // Run attention forward: save_h_s2 → attention output buffer @@ -2647,6 +2652,7 @@ impl GpuDqnTrainer { total_params, params_initialized: false, target_params_initialized: false, + attention_initialized: false, graph_forward: None, graph_adam: None, upload_staging_buf, @@ -5485,14 +5491,12 @@ impl GpuDqnTrainer { tau: f32, ) -> Result<(), MLError> { // Sync stream and clear any stale errors from graph capture phase. - let _eg = EventTrackingGuard::new(self.stream.context()); - // cudarc stores errors from cuStreamWaitEvent on disabled events during - // graph capture. check_err() consumes them so bind_to_thread() succeeds. - let _ = self.stream.context().check_err(); - - // First call: flatten target weights into flat buffer (20 DtoD copies, once only). - // Subsequent calls reuse the flat buffer which is kept in sync by unflatten at end. + // First call only: clear stale cudarc events from graph capture phase. + // Subsequent calls skip the sync — stream ordering is sufficient. if !self.target_params_initialized { + let _eg = EventTrackingGuard::new(self.stream.context()); + unsafe { cudarc::driver::sys::cuStreamSynchronize(self.stream.cu_stream()); } + let _ = self.stream.context().check_err(); self.flatten_target_weights(target_d, target_b)?; self.target_params_initialized = true; }