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

fix(gpu): prevent OOM from overflow in replay buffer, dynamic shmem tiles

Browse Source 
- 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 <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-03-10 17:26:43 +01:00
parent e2fa781cb2
commit d3ced2da44
5 changed files with 53 additions and 11 deletions
Download Patch File Download Diff File Expand all files Collapse all files

14
crates/ml-core/src/memory_optimization/auto_batch_size.rs
View File

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

16
crates/ml-dqn/src/gpu_replay_buffer.rs
View File

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

8
crates/ml-dqn/src/replay_buffer_type.rs
View File

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

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

@@ -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<CudaFunction>,
/// 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::<f32>() as u32;
let n = n_episodes as u32;

7
crates/ml/src/hyperopt/adapters/dqn.rs
View File

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