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) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-04-04 15:32:59 +02:00
parent be8a5c3331
commit 6565221b9f
2 changed files with 9 additions and 5 deletions

View File

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

View File

@@ -176,10 +176,11 @@ impl ReplayBufferType {
batch_size: usize,
stream: &Arc<CudaStream>,
) -> Result<Self, MLError> {
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