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

fix: GPU training hang — VRAM oversubscription, counterfactual data loss, event leak

Browse Source 
Three bugs causing the baseline RL training to hang on first epoch while
smoketest completes in 1.2s:

1. VRAM oversubscription (hang root cause): detect_gpu_hardware auto-scaling
   computed n_episodes without accounting for 2x counterfactual doubling or
   dtod_clone allocations, causing 4x actual memory vs budget. Replaced with
   configurable gpu_n_episodes field (smoketest=32, localdev=128, prod=4096).

2. Counterfactual experiences silently dropped: build_next_states_f32 received
   n_episodes instead of n_episodes*2, and PER insert used base count instead
   of doubled count — ~50% of augmented training data was generated then lost.

3. CudaEvent leak in hot loop: record_event(None) created+destroyed 8000 events
   per epoch in forward_online_raw/f32. Pre-allocated 4 events in CublasForward
   struct, eliminating driver overhead and handle leaks during CUDA Graph capture.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-04-02 23:08:33 +02:00
parent 529e22d8c4
commit 4c5f4049e5
10 changed files with 50 additions and 61 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
config/training/dqn-localdev.toml
View File

@@ -45,6 +45,7 @@ min_epochs_before_stopping = 80
[experience]
gpu_timesteps_per_episode = 200
gpu_n_episodes = 128
min_hold_bars = 3
[experience.fill_simulation]

1
config/training/dqn-production.toml
View File

@@ -51,6 +51,7 @@ patience = 20
min_epochs_before_stopping = 80
[experience]
gpu_n_episodes = 4096
initial_capital = 35000.0
tx_cost_multiplier = 0.0001
min_hold_bars = 5

1
config/training/dqn-smoketest.toml
View File

@@ -49,6 +49,7 @@ min_epochs_before_stopping = 5
[experience]
gpu_timesteps_per_episode = 100
gpu_n_episodes = 32
min_hold_bars = 3
[experience.fill_simulation]

39
crates/ml/src/cuda_pipeline/batched_forward.rs
View File

@@ -57,7 +57,7 @@ use std::mem::ManuallyDrop;
use cudarc::cublas::result as cublas_result;
use cudarc::cublas::sys as cublas_sys;
use cudarc::driver::{CudaFunction, CudaSlice, CudaStream, DevicePtr, LaunchConfig, PushKernelArg};
use cudarc::driver::{CudaEvent, CudaFunction, CudaSlice, CudaStream, DevicePtr, LaunchConfig, PushKernelArg};
use crate::MLError;
@@ -132,6 +132,11 @@ pub struct CublasForward {
/// Each branch (exposure, order, urgency) submits GEMMs to its own stream,
/// then the main stream joins all three before consuming the logits.
branch_streams: [Arc<CudaStream>; 3],
// ── Pre-allocated CUDA events for fork-join synchronization ──
/// Reused across forward passes — avoids cuEventCreate/Destroy per call.
trunk_done_event: CudaEvent,
branch_done_events: [CudaEvent; 3],
}
impl CublasForward {
@@ -193,6 +198,16 @@ impl CublasForward {
stream.fork().map_err(|e| MLError::DeviceError(format!("fork branch stream 2: {e}")))?,
];
// ── Pre-allocate CUDA events for fork-join (reused across all forward passes) ──
let ctx = stream.context();
let trunk_done_event = ctx.new_event(None)
.map_err(|e| MLError::DeviceError(format!("alloc trunk_done_event: {e}")))?;
let branch_done_events = [
ctx.new_event(None).map_err(|e| MLError::DeviceError(format!("alloc branch_done_event 0: {e}")))?,
ctx.new_event(None).map_err(|e| MLError::DeviceError(format!("alloc branch_done_event 1: {e}")))?,
ctx.new_event(None).map_err(|e| MLError::DeviceError(format!("alloc branch_done_event 2: {e}")))?,
];
Ok(Self {
handle: SendSyncCublasHandle(raw_handle),
_workspace_buf: workspace_buf,
@@ -214,6 +229,8 @@ impl CublasForward {
branch_1_size,
branch_2_size,
branch_streams,
trunk_done_event,
branch_done_events,
})
}
@@ -296,8 +313,8 @@ impl CublasForward {
// on a non-captured stream makes it join the capture graph, and the
// fork/join pattern is recorded as graph dependencies.
// Record trunk completion on the main stream.
let trunk_done = stream.record_event(None)
// Record trunk completion on the main stream (pre-allocated event).
self.trunk_done_event.record(stream)
.map_err(|e| MLError::ModelError(format!("trunk event record: {e}")))?;
let mut logit_byte_offset: u64 = 0;
@@ -307,7 +324,7 @@ impl CublasForward {
let w_fc_idx = branch_w_base[d];
// Branch stream waits for trunk completion.
bs.wait(&trunk_done)
bs.wait(&self.trunk_done_event)
.map_err(|e| MLError::ModelError(format!("branch {d} wait trunk: {e}")))?;
// Redirect cuBLAS handle to this branch stream.
@@ -336,11 +353,11 @@ impl CublasForward {
.map_err(|e| MLError::ModelError(format!("cublasSetStream restore: {e:?}")))?;
}
// Join: main stream waits for all 3 branches to complete.
// Join: main stream waits for all 3 branches to complete (pre-allocated events).
for d in 0..3 {
let branch_done = self.branch_streams[d].record_event(None)
self.branch_done_events[d].record(&self.branch_streams[d])
.map_err(|e| MLError::ModelError(format!("branch {d} done event: {e}")))?;
stream.wait(&branch_done)
stream.wait(&self.branch_done_events[d])
.map_err(|e| MLError::ModelError(format!("main wait branch {d}: {e}")))?;
}
} else {
@@ -405,7 +422,7 @@ impl CublasForward {
if distinct_branches {
// Multi-stream branch dispatch (f32 path, experience collection).
let trunk_done = stream.record_event(None)
self.trunk_done_event.record(stream)
.map_err(|e| MLError::ModelError(format!("f32 trunk event record: {e}")))?;
let mut logit_byte_offset: u64 = 0;
@@ -414,7 +431,7 @@ impl CublasForward {
let n_d = branch_sizes[d];
let w_fc_idx = branch_w_base[d];
bs.wait(&trunk_done)
bs.wait(&self.trunk_done_event)
.map_err(|e| MLError::ModelError(format!("f32 branch {d} wait trunk: {e}")))?;
unsafe {
@@ -440,9 +457,9 @@ impl CublasForward {
}
for d in 0..3 {
let branch_done = self.branch_streams[d].record_event(None)
self.branch_done_events[d].record(&self.branch_streams[d])
.map_err(|e| MLError::ModelError(format!("f32 branch {d} done event: {e}")))?;
stream.wait(&branch_done)
stream.wait(&self.branch_done_events[d])
.map_err(|e| MLError::ModelError(format!("f32 main wait branch {d}: {e}")))?;
}
} else {

3
crates/ml/src/cuda_pipeline/gpu_experience_collector.rs
View File

@@ -1385,8 +1385,9 @@ impl GpuExperienceCollector {
// Normalizing would scramble the C51 atom distribution.
// Build next_states on GPU (episode-aware shift by n_steps)
// Must use total (includes counterfactual) so next_states covers all experiences.
let next_states = build_next_states_f32(
&self.stream, &states, n_episodes, timesteps, sd, n_steps as usize,
&self.stream, &states, n_episodes * 2, timesteps, sd, n_steps as usize,
)?;
// ── Episode ID fill (GPU-native, zero CPU) ────────────────────

4
crates/ml/src/trainers/dqn/config.rs
View File

@@ -1175,6 +1175,9 @@ pub struct DQNHyperparameters {
/// Timesteps per episode in GPU kernel (default: 500, max: 1000)
/// Higher values collect more experience per launch at the cost of VRAM
pub gpu_timesteps_per_episode: usize,
/// Number of parallel episodes for GPU experience collection.
/// Controls VRAM usage: total = gpu_n_episodes × gpu_timesteps_per_episode × 2 (counterfactual) × state_dim × 4 bytes.
pub gpu_n_episodes: usize,
/// Average bid-ask spread for GPU portfolio simulation (default: 0.0001 = 1bp)
pub avg_spread: f64,
@@ -1677,6 +1680,7 @@ impl DQNHyperparameters {
// Phase 3: GPU experience collection
gpu_timesteps_per_episode: 500, // Default: 500 timesteps per episode
gpu_n_episodes: 256, // Default: 256 parallel episodes
avg_spread: 0.0001, // Default: 1bp (ES/NQ futures)
// GPU-dynamic network sizing — None means auto-detect from hardware.

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

@@ -729,7 +729,6 @@ impl DQNTrainer {
training_guard: None,
gpu_monitoring: None,
reduction_kernels: None,
cached_n_episodes: None,
// GPU pipeline: staging buffer pool (auto-initialized on CUDA devices)
buffer_pool,

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

@@ -265,9 +265,6 @@ pub struct DQNTrainer {
/// Multi-GPU configuration for data-parallel training (None = single GPU)
pub(crate) multi_gpu: Option<crate::cuda_pipeline::multi_gpu::MultiGpuConfig>,
/// Cached GPU n_episodes (computed once from nvidia-smi, reused across epochs)
/// Avoids forking nvidia-smi subprocess every epoch (~5-10ms per fork).
pub(crate) cached_n_episodes: Option<i32>,
/// Pre-computed OFI features per bar (indexed by global bar position).
/// Populated during data loading when MBP-10 order book data is available.

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

@@ -764,7 +764,7 @@ impl DQNTrainer {
cfg.value_hidden_dim,
cfg.branch_hidden_dim,
);
let alloc_episodes = self.compute_alloc_episodes(state_dim);
let alloc_episodes = self.compute_alloc_episodes();
Some(GpuExperienceCollector::new(
stream,
online_br.vars(),
@@ -789,7 +789,7 @@ impl DQNTrainer {
cfg.value_hidden_dim,
cfg.advantage_hidden_dim,
);
let alloc_episodes = self.compute_alloc_episodes(state_dim);
let alloc_episodes = self.compute_alloc_episodes();
Some(GpuExperienceCollector::new(
stream,
online.store(),
@@ -814,7 +814,7 @@ impl DQNTrainer {
cfg.value_hidden_dim,
cfg.advantage_hidden_dim,
);
let alloc_episodes = self.compute_alloc_episodes(state_dim);
let alloc_episodes = self.compute_alloc_episodes();
Some(GpuExperienceCollector::new(
stream,
online.vars(),
@@ -897,20 +897,9 @@ impl DQNTrainer {
Ok(())
}
/// Compute allocation episode count for GPU buffers.
/// Always auto-scales from VRAM — no manual override.
fn compute_alloc_episodes(&self, state_dim: usize) -> usize {
use ml_core::memory_optimization::detect_gpu_hardware;
match detect_gpu_hardware() {
Ok(hw) => {
let optimal = hw.optimal_n_episodes(
state_dim,
self.hyperparams.gpu_timesteps_per_episode,
);
optimal.max(32).min(16384)
}
Err(_) => 256, // safe fallback
}
/// Episode count for GPU buffer allocation — uses configured gpu_n_episodes.
fn compute_alloc_episodes(&self) -> usize {
self.hyperparams.gpu_n_episodes.max(32)
}
// ═══════════════════════════════════════════════════════════════════════
@@ -961,32 +950,7 @@ impl DQNTrainer {
use crate::cuda_pipeline::gpu_experience_collector::ExperienceCollectorConfig;
let raw_sd = if !self.hyperparams.mbp10_data_dir.is_empty() { 53 } else { 45 };
let aligned_sd = (raw_sd + 7) & !7;
// Cache n_episodes on first epoch — always auto-scaled from VRAM.
let n_episodes = if let Some(cached) = self.cached_n_episodes {
cached
} else {
use ml_core::memory_optimization::detect_gpu_hardware;
let computed = match detect_gpu_hardware() {
Ok(hw) => {
let optimal = hw.optimal_n_episodes(
aligned_sd,
self.hyperparams.gpu_timesteps_per_episode,
);
let chosen = optimal.max(32).min(16384) as i32;
info!(
"GPU auto-scaled n_episodes: {} (SMs={}, VRAM={:.0}MB)",
chosen, hw.sm_count, hw.free_memory_mb
);
chosen
}
Err(_) => 256_i32, // safe fallback
};
self.cached_n_episodes = Some(computed);
computed
};
let n_episodes = self.hyperparams.gpu_n_episodes.max(32) as i32;
// Domain randomization: all randomization is GPU-native (zero CPU RNG)
// Domain randomization, mirror, causal, vaccine: ALL always active (one production path)
@@ -1179,7 +1143,7 @@ impl DQNTrainer {
"GPU zero-roundtrip collection FAILED (no CPU fallback): {e}"
))?;
let count = gpu_batch.n_episodes * gpu_batch.timesteps;
let count = gpu_batch.n_episodes * gpu_batch.timesteps * 2; // counterfactual doubles experiences
// Record CUDA event after experience collection completes on the forked stream.
// This replaces stream.synchronize() — the event is checked lazily at the start
@@ -1238,7 +1202,7 @@ impl DQNTrainer {
if count > 0 {
// Direct CudaSlice insertion — zero Candle, zero CPU roundtrip.
let total = gpu_batch.n_episodes * gpu_batch.timesteps;
let total = gpu_batch.n_episodes * gpu_batch.timesteps * 2; // counterfactual doubles experiences
// actions: CudaSlice<i32> → CudaSlice<u32> (safe reinterpret for values [0..44])
let actions_u32: &CudaSlice<u32> = unsafe {

4
crates/ml/src/training_profile.rs
View File

@@ -257,6 +257,7 @@ pub struct RewardSection {
#[derive(Debug, Clone, Deserialize, Default)]
pub struct ExperienceSection {
pub gpu_timesteps_per_episode: Option<usize>,
pub gpu_n_episodes: Option<usize>,
pub initial_capital: Option<f64>,
/// Maps to `DQNHyperparameters::transaction_cost_multiplier`.
pub tx_cost_multiplier: Option<f64>,
@@ -933,6 +934,9 @@ impl DqnTrainingProfile {
if let Some(v) = ex.gpu_timesteps_per_episode {
hp.gpu_timesteps_per_episode = v;
}
if let Some(v) = ex.gpu_n_episodes {
hp.gpu_n_episodes = v;
}
if let Some(v) = ex.initial_capital {
hp.initial_capital = v as f32;
}