From e3f7ab3347047a34f558dceebfde3c808b562b4b Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Mon, 13 Apr 2026 10:36:18 +0200 Subject: [PATCH] fix: single-stream DDQN for cuBLAS determinism Move DDQN forward (Pass 3) from double_dqn_stream to main stream. NVIDIA docs: "bit-wise reproducibility is valid only when a single CUDA stream is active." The double_dqn_stream fork is kept but no longer used for cuBLAS operations. Combined with AlgoGetIds (verified: same algo_id=16 every run), this satisfies all documented NVIDIA preconditions for bit-wise reproducibility. Epochs 1-3 are identical across runs. Epoch 4+ divergence remains under investigation. Co-Authored-By: Claude Opus 4.6 (1M context) --- .../ml/src/cuda_pipeline/gpu_dqn_trainer.rs | 37 +++++++------------ 1 file changed, 13 insertions(+), 24 deletions(-) diff --git a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs index 2c3277d98..2b034032f 100644 --- a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs +++ b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs @@ -2808,8 +2808,10 @@ impl GpuDqnTrainer { let double_dqn_stream = stream.fork() .map_err(|e| MLError::DeviceError(format!("fork double_dqn_stream: {e}")))?; + // Use MAIN stream for DDQN too — NVIDIA bit-wise reproducibility + // requires a single active CUDA stream. let cublas_forward_ddqn = CublasForward::new( - &double_dqn_stream, + &stream, config.batch_size, config.state_dim, config.shared_h1, @@ -3941,30 +3943,13 @@ impl GpuDqnTrainer { } if step < 3 { tracing::info!(step, "replay_forward: main graph launched, recording pass1_event"); } - // 2. Record event after main graph completes - self.pass1_event.record(&self.stream) - .map_err(|e| MLError::ModelError(format!("pass1 event record: {e}")))?; - - // 3. double_dqn_stream waits for main stream - self.double_dqn_stream.wait(&self.pass1_event) - .map_err(|e| MLError::ModelError(format!("ddqn wait pass1: {e}")))?; - - // 4. Replay Pass 3 on double_dqn_stream - if step < 3 { tracing::info!(step, "replay_forward: launching ddqn graph"); } + // 2. Replay Pass 3 (DDQN) on the SAME main stream — sequential, deterministic. + // NVIDIA: "bit-wise reproducibility is valid only when a single CUDA stream is active" if let Some(ref g) = self.graph_forward_ddqn { g.0.launch().map_err(|e| { MLError::ModelError(format!("CUDA graph_forward_ddqn replay: {e}")) })?; } - if step < 3 { tracing::info!(step, "replay_forward: ddqn launched, recording pass3_event"); } - - // 5. Main stream waits for Pass 3 to complete - self.pass3_event.record(&self.double_dqn_stream) - .map_err(|e| MLError::ModelError(format!("pass3 event record: {e}")))?; - self.stream.wait(&self.pass3_event) - .map_err(|e| MLError::ModelError(format!("join pass3: {e}")))?; - - if step < 3 { tracing::info!(step, "replay_forward: complete"); } Ok(()) } @@ -4601,8 +4586,9 @@ impl GpuDqnTrainer { MLError::ModelError("graph_forward_mse capture returned None".into()) })?; - // ── Capture graph_forward_ddqn (Pass 3 on double_dqn_stream) ── - let begin_ddqn = self.double_dqn_stream.begin_capture( + // ── Capture graph_forward_ddqn (on main stream for determinism) ── + // NVIDIA: "bit-wise reproducibility is valid only when a single CUDA stream is active" + let begin_ddqn = self.stream.begin_capture( cudarc::driver::sys::CUstreamCaptureMode::CU_STREAM_CAPTURE_MODE_THREAD_LOCAL, ); if let Err(e) = begin_ddqn { @@ -4611,7 +4597,7 @@ impl GpuDqnTrainer { return Err(MLError::ModelError(format!("graph_forward_ddqn begin_capture: {e}"))); } let submit_ddqn_result = self.submit_forward_ops_ddqn(); - let graph_ddqn_result = self.double_dqn_stream.end_capture( + let graph_ddqn_result = self.stream.end_capture( cudarc::driver::sys::CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH, ); submit_ddqn_result?; @@ -4951,8 +4937,11 @@ impl GpuDqnTrainer { let param_sizes = compute_param_sizes(&self.config); let on_w_ptrs = f32_weight_ptrs_from_base(self.ptrs.params_ptr, ¶m_sizes); + // Run DDQN on the MAIN stream for deterministic cuBLAS. + // NVIDIA guarantees bit-wise reproducibility only with a single active stream. + // The double_dqn_stream parallelism is sacrificed for determinism. self.cublas_forward_ddqn.forward_online_raw( - &self.double_dqn_stream, self.ptrs.next_states_buf, &on_w_ptrs, + &self.stream, self.ptrs.next_states_buf, &on_w_ptrs, self.ptrs.on_next_h_s1_scratch, self.ptrs.on_next_h_s2_scratch, self.ptrs.on_next_h_v_scratch, self.ptrs.on_next_h_b_scratch, self.ptrs.on_next_h_b_scratch,