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

fix(cuda): unify all GPU ops onto single forked CudaStream

Browse Source 
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) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-03-17 10:42:56 +01:00
parent fa739ba40c
commit 44ddd38112
6 changed files with 51 additions and 24 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -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<CudaStream> 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) =

17
crates/ml/src/trainers/dqn/fused_training.rs
View File

@@ -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<candle_core::cuda_backend::cudarc::driver::CudaStream>,
) -> Result<Self> {
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();

17
crates/ml/src/trainers/dqn/trainer/constructor.rs
View File

@@ -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<Arc<candle_core::cuda_backend::cudarc::driver::CudaStream>> =
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(),

6
crates/ml/src/trainers/dqn/trainer/mod.rs
View File

@@ -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<Arc<CudaStream>>,
/// Training metrics
pub(crate) metrics: Arc<RwLock<TrainingMetrics>>,
/// Loss history for plateau detection

7
crates/ml/src/trainers/dqn/trainer/train_step.rs
View File

@@ -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);

14
crates/ml/src/trainers/dqn/trainer/training_loop.rs
View File

@@ -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;