From 44ddd38112f8fbf12222540f7561c74fbaea3cbd Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Tue, 17 Mar 2026 10:42:56 +0100 Subject: [PATCH] fix(cuda): unify all GPU ops onto single forked CudaStream MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Fork a dedicated CudaStream once in DQNTrainer constructor and share it via Arc::clone with all GPU components (experience collector, fused trainer, portfolio sim, monitoring). This eliminates: 1. Candle's default stream (stream 0) from the training hot path 2. Dual-stream cudarc event tracking conflicts during CUDA Graph capture 3. Implicit serialization points between default and forked streams The fused training context no longer forks its own stream — it receives the trainer's unified stream. With all GPU work on a single stream, cudarc event tracking is safely disabled in GpuDqnTrainer::new(), removing the cuStreamWaitEvent/cuEventRecord overhead that broke CUDA Graph capture. Co-Authored-By: Claude Opus 4.6 (1M context) --- crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs | 14 +++++++++----- crates/ml/src/trainers/dqn/fused_training.rs | 17 +++++------------ .../ml/src/trainers/dqn/trainer/constructor.rs | 17 +++++++++++++++++ crates/ml/src/trainers/dqn/trainer/mod.rs | 6 ++++++ .../ml/src/trainers/dqn/trainer/train_step.rs | 7 ++++++- .../src/trainers/dqn/trainer/training_loop.rs | 14 ++++++++------ 6 files changed, 51 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 b18824c9a..1aca3461c 100644 --- a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs +++ b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs @@ -383,11 +383,15 @@ impl GpuDqnTrainer { let b = config.batch_size; let total_params = compute_total_params(&config); - // TODO: Unify all GPU ops onto this forked stream. Currently the - // experience collector + curiosity trainer use Candle's default stream - // (stream 0) while this trainer uses the forked stream. The dual-stream - // architecture causes cudarc event tracking conflicts. - // Fix: store this Arc in DQNTrainer, pass to all GPU components. + // All GPU components now share this single forked CudaStream (stored in + // DQNTrainer). Disable cudarc's automatic event tracking: it injects + // cuStreamWaitEvent/cuEventRecord into every buffer access, creating + // cross-stream event references that break CUDA Graph capture. Safe + // because all GPU work runs on this single stream — no cross-stream + // synchronization needed. + // + // SAFETY: single-stream architecture, all work synchronized on `stream`. + unsafe { stream.context().disable_event_tracking(); } // ── Compile all 5 training kernels from same module ────────── let (forward_loss_kernel, backward_kernel, grad_norm_kernel, adam_update_kernel, f32_to_bf16_kernel) = diff --git a/crates/ml/src/trainers/dqn/fused_training.rs b/crates/ml/src/trainers/dqn/fused_training.rs index 49d6ad204..5cdf0ea3b 100644 --- a/crates/ml/src/trainers/dqn/fused_training.rs +++ b/crates/ml/src/trainers/dqn/fused_training.rs @@ -82,10 +82,11 @@ impl FusedTrainingCtx { /// /// Only valid when `device.is_cuda() && agent.is_using_branching()`. pub(crate) fn new( - device: &Device, + _device: &Device, agent: &DQNAgentType, hyperparams: &DQNHyperparameters, batch_size: usize, + stream: Arc, ) -> Result { let dqn = match agent { DQNAgentType::Standard(d) => d, @@ -103,17 +104,9 @@ impl FusedTrainingCtx { anyhow::anyhow!("Fused CUDA training requires branching target network") })?; - // Get a dedicated CudaStream for CUDA Graph capture. - // The default stream from Candle is the legacy default stream (stream 0) - // which does NOT support CUDA Graph capture (begin_capture returns - // CUDA_ERROR_STREAM_CAPTURE_UNSUPPORTED). We fork a new non-default - // stream that supports capture and all async operations. - let cuda_dev = match device { - Device::Cuda(dev) => dev, - _ => return Err(anyhow::anyhow!("Fused CUDA training requires CUDA device")), - }; - let stream = cuda_dev.cuda_stream().fork() - .map_err(|e| anyhow::anyhow!("Failed to fork CUDA stream for graph capture: {e}"))?; + // Use the caller-provided forked CudaStream. All GPU components (experience + // collector, fused trainer, portfolio sim, monitoring) share this single + // stream via Arc::clone — no more dual-stream event tracking conflicts. // Build config from DQN network dimensions let (shared_h1, shared_h2, value_h, adv_h) = agent.network_dims(); diff --git a/crates/ml/src/trainers/dqn/trainer/constructor.rs b/crates/ml/src/trainers/dqn/trainer/constructor.rs index 7831f6c37..2a4365526 100644 --- a/crates/ml/src/trainers/dqn/trainer/constructor.rs +++ b/crates/ml/src/trainers/dqn/trainer/constructor.rs @@ -126,6 +126,22 @@ impl DQNTrainer { .map_err(|e| anyhow::anyhow!("Failed to initialize device: {}", e))? }; + // Fork a dedicated CudaStream once. All GPU components (experience collector, + // fused trainer, portfolio sim, monitoring) share this single stream via + // Arc::clone. This eliminates: + // 1. Candle's default stream (stream 0) from the hot path + // 2. Dual-stream cudarc event tracking conflicts during CUDA Graph capture + // 3. Implicit serialization points between default and forked streams + let cuda_stream: Option> = + if let candle_core::Device::Cuda(ref cuda_dev) = device { + let stream = cuda_dev.cuda_stream().fork() + .map_err(|e| anyhow::anyhow!("Failed to fork CUDA stream: {e}"))?; + info!("Forked dedicated CudaStream for all GPU components"); + Some(stream) + } else { + None + }; + // Dynamic replay buffer sizing: scale replay capacity to available VRAM. // Only activates when replay_buffer_vram_fraction > 0 and GPU is detected. // Also computes the PER memory budget from actual VRAM — no hardcoded caps. @@ -555,6 +571,7 @@ impl DQNTrainer { agent: Arc::new(RwLock::new(agent)), hyperparams, device, + cuda_stream, metrics: Arc::new(RwLock::new(TrainingMetrics::new())), loss_history: Vec::new(), q_value_history: Vec::new(), diff --git a/crates/ml/src/trainers/dqn/trainer/mod.rs b/crates/ml/src/trainers/dqn/trainer/mod.rs index a00914751..b26290d7b 100644 --- a/crates/ml/src/trainers/dqn/trainer/mod.rs +++ b/crates/ml/src/trainers/dqn/trainer/mod.rs @@ -8,6 +8,7 @@ use std::sync::Arc; use anyhow::{Context, Result}; use candle_core::{Device, Tensor}; +use candle_core::cuda_backend::cudarc::driver::CudaStream; use crate::cuda_pipeline::DqnGpuData; use ml_core::fill_simulator::FillSimulator; use risk::drawdown_monitor::DrawdownMonitor; @@ -50,6 +51,11 @@ pub struct DQNTrainer { pub(crate) hyperparams: DQNHyperparameters, /// Device (GPU or CPU) pub(crate) device: Device, + /// Single forked CudaStream shared by ALL GPU components (experience collector, + /// fused trainer, portfolio sim, monitoring). Forked once in constructor — + /// eliminates dual-stream event tracking conflicts and Candle default stream + /// usage in the hot path. None on CPU devices. + pub(crate) cuda_stream: Option>, /// Training metrics pub(crate) metrics: Arc>, /// Loss history for plateau detection diff --git a/crates/ml/src/trainers/dqn/trainer/train_step.rs b/crates/ml/src/trainers/dqn/trainer/train_step.rs index e92ee3b65..453547cc5 100644 --- a/crates/ml/src/trainers/dqn/trainer/train_step.rs +++ b/crates/ml/src/trainers/dqn/trainer/train_step.rs @@ -610,13 +610,18 @@ impl DQNTrainer { if !needs_init { return; } + // Fused context requires the unified forked CudaStream + let stream = match self.cuda_stream { + Some(ref s) => std::sync::Arc::clone(s), + None => return, + }; if self.fused_ctx.is_some() { info!("Fused CUDA context: batch_size changed, recreating"); self.fused_ctx = None; } let agent = self.agent.read().await; match super::super::fused_training::FusedTrainingCtx::new( - &self.device, &*agent, &self.hyperparams, self.current_batch_size, + &self.device, &*agent, &self.hyperparams, self.current_batch_size, stream, ) { Ok(ctx) => { info!("Fused CUDA training initialized (batch_size={})", self.current_batch_size); diff --git a/crates/ml/src/trainers/dqn/trainer/training_loop.rs b/crates/ml/src/trainers/dqn/trainer/training_loop.rs index 409dcfda1..c59023fdd 100644 --- a/crates/ml/src/trainers/dqn/trainer/training_loop.rs +++ b/crates/ml/src/trainers/dqn/trainer/training_loop.rs @@ -13,6 +13,8 @@ //! - `log_epoch_metrics_and_financials`: logging, Prometheus, QuestDB, VaR/CVaR //! - `handle_epoch_checkpoints_and_early_stopping`: best-checkpoint + early stopping +use std::sync::Arc; + use anyhow::{Context, Result}; use candle_core::Tensor; use common::metrics::{questdb_sink, training_metrics}; @@ -372,11 +374,11 @@ impl DQNTrainer { return Ok(()); } - let candle_core::Device::Cuda(ref cuda_dev) = self.device else { - return Ok(()); + let stream = match self.cuda_stream { + Some(ref s) => Arc::clone(s), + None => return Ok(()), }; - let stream = cuda_dev.cuda_stream(); let target_dim = 4; let feature_dim = 42; let num_bars = training_data.len(); @@ -458,12 +460,12 @@ impl DQNTrainer { return Ok(()); } - let candle_core::Device::Cuda(ref cuda_dev) = self.device else { - return Ok(()); + let stream = match self.cuda_stream { + Some(ref s) => Arc::clone(s), + None => return Ok(()), }; use crate::cuda_pipeline::gpu_experience_collector::GpuExperienceCollector; - let stream = cuda_dev.cuda_stream(); // Read-lock agent to extract dueling network weights for GPU collector. let agent = self.agent.read().await;