From 6565221b9f481c5b2a9775c3e5bdacd2ee62be37 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Sat, 4 Apr 2026 15:32:59 +0200 Subject: [PATCH] fix: remove hardcoded MIN_GPU_CAPACITY=1024, use batch_size as floor The GPU PER segment tree works with any power-of-2 capacity. The arbitrary 1024 floor prevented small buffer sizes needed for fast smoketests. Now the minimum is batch_size (can't sample more than buffer holds). Co-Authored-By: Claude Opus 4.6 (1M context) --- config/training/dqn-smoketest.toml | 2 +- crates/ml-dqn/src/replay_buffer_type.rs | 12 ++++++++---- 2 files changed, 9 insertions(+), 5 deletions(-) diff --git a/config/training/dqn-smoketest.toml b/config/training/dqn-smoketest.toml index c03875057..ec7924ab0 100644 --- a/config/training/dqn-smoketest.toml +++ b/config/training/dqn-smoketest.toml @@ -37,7 +37,7 @@ count_bonus_coefficient = 0.1 q_gap_threshold = 0.05 [replay_buffer] -buffer_size = 1024 +buffer_size = 256 min_replay_size = 64 regime_replay_decay = 0.3 diff --git a/crates/ml-dqn/src/replay_buffer_type.rs b/crates/ml-dqn/src/replay_buffer_type.rs index 66304bd9d..a690001f1 100644 --- a/crates/ml-dqn/src/replay_buffer_type.rs +++ b/crates/ml-dqn/src/replay_buffer_type.rs @@ -176,10 +176,11 @@ impl ReplayBufferType { batch_size: usize, stream: &Arc, ) -> Result { - const MIN_GPU_CAPACITY: usize = 1024; + // Minimum capacity = batch_size (can't sample more than buffer holds) + let min_capacity = batch_size.max(1); let mut try_cap = capacity; - while try_cap >= MIN_GPU_CAPACITY { + while try_cap >= min_capacity { match Self::new_gpu_prioritized(try_cap, state_dim, alpha, beta, beta_max, beta_annealing_steps, max_memory_bytes, batch_size, stream) { Ok(buf) => { if try_cap < capacity { @@ -201,14 +202,17 @@ impl ReplayBufferType { try_cap, e, try_cap / 2 ); try_cap /= 2; - if try_cap < MIN_GPU_CAPACITY { + if try_cap < min_capacity { return Err(e); } } } } - Err(MLError::ModelError("GPU PER capacity below floor".into())) + Err(MLError::ModelError(format!( + "GPU PER allocation failed: capacity {} below minimum {} (batch_size)", + try_cap, min_capacity + ))) } /// Sample a batch from the buffer