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

fix: single-stream DDQN for cuBLAS determinism

Browse Source 
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) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-04-13 10:36:18 +02:00
parent 13524dd3a5
commit e3f7ab3347
1 changed files with 13 additions and 24 deletions
Download Patch File Download Diff File Expand all files Collapse all files

37
crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs
View File

@@ -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, &param_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,