From d3ced2da440c43ca83490d52744c020903d033ad Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Tue, 10 Mar 2026 17:26:43 +0100 Subject: [PATCH] fix(gpu): prevent OOM from overflow in replay buffer, dynamic shmem tiles - Staging/GPU replay buffer: truncate batch when it exceeds ring buffer capacity (CPU fallback can flush 100K+ experiences at once) - GpuExperienceCollector: compute shmem tile rows dynamically to stay under 48 KB default limit instead of hardcoded 64-row constant - Auto batch size: subtract concurrent VRAM consumers (~530 MB model + optimizer + PER + data) before budgeting experience collector at 40% - Hyperopt DQN adapter: GPU PER always on, include replay buffer in VRAM estimate for small GPUs (was excluded assuming CPU-only PER) Co-Authored-By: Claude Opus 4.6 --- .../memory_optimization/auto_batch_size.rs | 14 +++++++++++--- crates/ml-dqn/src/gpu_replay_buffer.rs | 16 +++++++++++++++- crates/ml-dqn/src/replay_buffer_type.rs | 8 ++++++++ .../cuda_pipeline/gpu_experience_collector.rs | 19 ++++++++++++++++--- crates/ml/src/hyperopt/adapters/dqn.rs | 7 +++---- 5 files changed, 53 insertions(+), 11 deletions(-) diff --git a/crates/ml-core/src/memory_optimization/auto_batch_size.rs b/crates/ml-core/src/memory_optimization/auto_batch_size.rs index 8735c6103..448052bc4 100644 --- a/crates/ml-core/src/memory_optimization/auto_batch_size.rs +++ b/crates/ml-core/src/memory_optimization/auto_batch_size.rs @@ -456,10 +456,18 @@ impl GpuHardwareInfo { const MAX_EPISODES: usize = 8192; // Per-episode output buffer: states + actions + rewards + done + target_q + td_error - let per_episode_bytes = timesteps * (state_dim + 5) * 4; + // Plus per-episode state: portfolio(6) + barrier(10) + diversity(30+2) + rng(1) ≈ 50 floats + let per_episode_bytes = timesteps * (state_dim + 5) * 4 + 50 * 4; - // VRAM cap: use at most 15% of free VRAM for experience buffers - let vram_budget_bytes = (self.free_memory_mb * 1024.0 * 1024.0 * 0.15) as usize; + // VRAM budget for experience collector output buffers. + // Must account for concurrent VRAM consumers: + // ~300 MB model weights + optimizer state + backward graph + // ~30 MB GPU PER replay buffer (50K × 48 × BF16 × 6 tensors) + // ~variable training data upload (up to 200 MB for ~1M bars) + // Subtract this overhead, then use 40% of remaining for collector. + let overhead_mb = 530.0_f64; + let available_mb = (self.free_memory_mb - overhead_mb).max(64.0); + let vram_budget_bytes = (available_mb * 1024.0 * 1024.0 * 0.40) as usize; let vram_cap = if per_episode_bytes > 0 { vram_budget_bytes / per_episode_bytes } else { diff --git a/crates/ml-dqn/src/gpu_replay_buffer.rs b/crates/ml-dqn/src/gpu_replay_buffer.rs index d99caa93a..be50ad685 100644 --- a/crates/ml-dqn/src/gpu_replay_buffer.rs +++ b/crates/ml-dqn/src/gpu_replay_buffer.rs @@ -338,11 +338,25 @@ impl GpuReplayBuffer { return Ok(()); } + let cap = self.config.capacity; + + // When batch exceeds capacity, keep only the last `cap` items. + // This handles the CPU experience collector fallback which may flush + // hundreds of thousands of experiences at once into a smaller ring buffer. + if batch_size > cap { + let offset = batch_size - cap; + let s = states.narrow(0, offset, cap)?; + let ns = next_states.narrow(0, offset, cap)?; + let a = actions.narrow(0, offset, cap)?; + let r = rewards.narrow(0, offset, cap)?; + let d = dones.narrow(0, offset, cap)?; + return self.insert_batch(&s, &ns, &a, &r, &d); + } + // Cast incoming states to match buffer dtype (e.g. F32 → BF16 on CUDA) let cast_states = &states.to_dtype(self.states.dtype())?; let cast_next_states = &next_states.to_dtype(self.next_states.dtype())?; - let cap = self.config.capacity; let cursor = self.write_cursor; // Build priority tensor: all new experiences get max_priority. diff --git a/crates/ml-dqn/src/replay_buffer_type.rs b/crates/ml-dqn/src/replay_buffer_type.rs index add91833f..7ba1890d3 100644 --- a/crates/ml-dqn/src/replay_buffer_type.rs +++ b/crates/ml-dqn/src/replay_buffer_type.rs @@ -64,6 +64,14 @@ impl StagedGpuBuffer { if self.staging.is_empty() { return Ok(()); } + let cap = self.gpu.capacity(); + // When staging exceeds buffer capacity, keep only the last `cap` items. + // This prevents OOM when the CPU experience collector fallback flushes + // hundreds of thousands of experiences that won't fit in the ring buffer. + if self.staging.len() > cap { + let drain_count = self.staging.len() - cap; + self.staging.drain(..drain_count); + } let n = self.staging.len(); // Use the GPU buffer's configured state_dim (tensor-core-aligned) rather // than the raw experience dimension, so the flushed tensors match the diff --git a/crates/ml/src/cuda_pipeline/gpu_experience_collector.rs b/crates/ml/src/cuda_pipeline/gpu_experience_collector.rs index 0370b090f..9be07075f 100644 --- a/crates/ml/src/cuda_pipeline/gpu_experience_collector.rs +++ b/crates/ml/src/cuda_pipeline/gpu_experience_collector.rs @@ -257,6 +257,8 @@ pub struct GpuExperienceCollector { state_dim: usize, /// Network hidden dims — needed for shared memory tile sizing at launch time. network_dims: (usize, usize, usize, usize), + /// Shared memory tile row count (dynamic, fits under 48 KB default limit). + shmem_tile_rows: usize, /// Warp-cooperative kernel for sm_90+ (H100). None if not available. warp_kernel_func: Option, /// Whether to use warp-cooperative kernel (sm_90+ detected). @@ -344,6 +346,15 @@ impl GpuExperienceCollector { let shmem_max_in_dim = state_dim.max(shared_h1).max(shared_h2); // NOISY_MAX_DIM must cover the widest layer input: same as shmem_max_in_dim. let noisy_max_dim = shmem_max_in_dim; + // Tile rows: dynamic to stay under default 48 KB shared memory limit. + // Formula: (tile_rows * shmem_max_in + tile_rows) * 4 < 49152 + // → tile_rows < 49152 / (4 * (shmem_max_in + 1)) + let shmem_tile_rows = { + let max_tile = 49152 / (4 * (shmem_max_in_dim + 1)); + // Round down to power of 2 for clean tiling, min 16 + let pow2 = (max_tile as u32).next_power_of_two() >> 1; + pow2.max(16) as usize + }; let dim_overrides = format!( "#define NOISY_MAX_DIM {noisy_max_dim}\n\ #define STATE_DIM {state_dim}\n\ @@ -354,7 +365,8 @@ impl GpuExperienceCollector { #define VALUE_H {value_h}\n\ #define ADV_H {adv_h}\n\ #define NUM_ATOMS_MAX {num_atoms_max}\n\ - #define SHMEM_MAX_IN_DIM {shmem_max_in_dim}\n" + #define SHMEM_MAX_IN_DIM {shmem_max_in_dim}\n\ + #define SHMEM_TILE_ROWS {shmem_tile_rows}\n" ); let full_source = format!("{common_src}\n{dim_overrides}\n{kernel_src}"); info!( @@ -578,6 +590,7 @@ impl GpuExperienceCollector { kernel_func, state_dim, network_dims, + shmem_tile_rows, warp_kernel_func, use_warp_kernel, alloc_episodes, @@ -870,10 +883,10 @@ impl GpuExperienceCollector { })?; // ---- Step 4: Launch config ---- - const SHMEM_TILE_ROWS: u32 = 64; + let tile_rows = self.shmem_tile_rows as u32; let (sh1, sh2, _, _) = self.network_dims; let shmem_max_in = self.state_dim.max(sh1).max(sh2) as u32; - let shmem_bytes = (SHMEM_TILE_ROWS * shmem_max_in + SHMEM_TILE_ROWS) + let shmem_bytes = (tile_rows * shmem_max_in + tile_rows) * std::mem::size_of::() as u32; let n = n_episodes as u32; diff --git a/crates/ml/src/hyperopt/adapters/dqn.rs b/crates/ml/src/hyperopt/adapters/dqn.rs index 35f171723..69b16a514 100644 --- a/crates/ml/src/hyperopt/adapters/dqn.rs +++ b/crates/ml/src/hyperopt/adapters/dqn.rs @@ -740,7 +740,7 @@ impl ParameterSpace for DQNParams { if small_gpu { // batch_size: [64, 128] — backward graph scales with batch if let Some(b) = bounds.get_mut(1) { b.1 = b.1.min(128.0); } - // buffer_size: cap to ln(50_000) ≈ 10.82 (CPU PER, so less critical) + // buffer_size: cap to ln(50_000) ≈ 10.82 (GPU PER ring buffer fits in ~28 MB) if let Some(b) = bounds.get_mut(3) { b.1 = b.1.min(50_000_f64.ln()); } // num_atoms: fix to 51 (min) — 201 atoms × 5 actions = 1005 output dim is too large if let Some(b) = bounds.get_mut(15) { b.1 = b.1.min(51.0); } @@ -761,11 +761,10 @@ impl ParameterSpace for DQNParams { } } // Cap hidden_dim_base by VRAM using actual DQN architecture params. - // On small GPUs, pass 0 for replay_buffer_capacity because GPU PER is disabled - // and the estimator shouldn't penalize for CPU-side buffers. + // GPU PER is always on — include replay buffer in VRAM estimate. let worst_batch = bounds.get(1).map_or(512, |b| b.1 as usize); let worst_atoms = bounds.get(15).map_or(201, |b| b.1 as usize); - let worst_buffer = if small_gpu { 0 } else { 100_000 }; + let worst_buffer = if small_gpu { 50_000 } else { 100_000 }; let max_base = budget.max_hidden_dim_base_full( 1, // concurrent_trials (1 per GPU) worst_batch, // worst-case batch_size from (capped) search space