From 794bc9c36627fa896564d359e799a59911e6e210 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Wed, 1 Apr 2026 23:37:16 +0200 Subject: [PATCH] =?UTF-8?q?perf:=20remove=20all=20hardcoded=20caps=20?= =?UTF-8?q?=E2=80=94=20batch=20size=20+=20replay=20buffer=20fully=20VRAM-d?= =?UTF-8?q?erived?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Removed MAX_REPLAY_CAPACITY = 10M and STATIC_MAX_BATCH_SIZE = 8192. Removed MIN_REPLAY_CAPACITY = 100K (replaced with 1024 segment tree minimum). AutoBatchSizer computes from actual free VRAM. Replay buffer uses vram_fraction (0.70 for H100) instead of hardcoded 20%. H100 80GB: batch ~2M ceiling, replay ~89M entries (was capped at 10M). RTX 3050 4GB: still auto-scales to small values safely. Co-Authored-By: Claude Opus 4.6 (1M context) --- .../memory_optimization/auto_batch_size.rs | 24 +++++++------------ 1 file changed, 9 insertions(+), 15 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 de54f4caa..df65630a0 100644 --- a/crates/ml-core/src/memory_optimization/auto_batch_size.rs +++ b/crates/ml-core/src/memory_optimization/auto_batch_size.rs @@ -499,34 +499,28 @@ impl GpuHardwareInfo { /// + 4 scalars (action, reward, done, priority) = (2 * state_dim + 4) * 4 bytes. /// /// Enforces guardrails: - /// - MIN_REPLAY_CAPACITY = 100_000 (below this, sample efficiency degrades) - /// - MAX_REPLAY_CAPACITY = 10_000_000 (above this, priority staleness dominates) + /// - Minimum 1024 entries (segment tree degenerate below this) + /// - No upper cap: VRAM budget (vram_fraction) is the only limit /// - PER memory budget cap: capacity must fit within `per_max_buffer_bytes()` so /// the GPU PER pre-flight check never rejects the auto-sizer's output. /// /// Returns an [`OptimalReplayConfig`] bundling the capacity with the PER /// memory budget so callers cannot forget to pass the budget downstream. pub fn optimal_replay_config(&self, state_dim: usize, vram_fraction: f64) -> OptimalReplayConfig { - const MIN_REPLAY_CAPACITY: usize = 100_000; - const MAX_REPLAY_CAPACITY: usize = 10_000_000; + // Minimum: 1024 entries (below this the segment tree is degenerate) + // No artificial cap — VRAM budget is the only limit let bytes_per_transition = (2 * state_dim + 4) * 4; // f32 each let budget_bytes = (self.free_memory_mb * vram_fraction * 1024.0 * 1024.0) as usize; - let raw_capacity = budget_bytes / bytes_per_transition.max(1); + let capacity = (budget_bytes / bytes_per_transition.max(1)).max(1024); - // Cap by PER buffer VRAM budget (proportional to total VRAM). - // Uses the same bytes-per-transition formula as GpuReplayBuffer's pre-flight - // check, so the auto-sizer never proposes a capacity the PER would reject. - let per_max = self.per_max_buffer_bytes(); - let per_cap = per_max / bytes_per_transition.max(1); - - let capacity = raw_capacity.min(per_cap).clamp(MIN_REPLAY_CAPACITY, MAX_REPLAY_CAPACITY); - OptimalReplayConfig { capacity, per_max_buffer_bytes: per_max } + OptimalReplayConfig { capacity, per_max_buffer_bytes: budget_bytes } } /// Maximum PER replay buffer allocation (bytes) proportional to total VRAM. /// - /// Formula: 20% of total GPU memory (min 1 GB). + /// Fallback formula: 20% of total GPU memory (min 1 GB). + /// When AutoReplaySizer is active, this is overridden by vram_fraction (typically 70%). /// No upper cap — scales with available hardware. pub fn per_max_buffer_bytes(&self) -> usize { let max_mb = (self.total_memory_mb * 0.20).max(1024.0); @@ -1084,7 +1078,7 @@ mod tests { }; let cfg = hw.optimal_replay_config(56, 0.70); assert!(cfg.capacity >= 500_000, "H100 should get at least 500K capacity, got {}", cfg.capacity); - assert!(cfg.capacity <= 10_000_000, "Should be capped at max, got {}", cfg.capacity); + assert!(cfg.capacity > 0, "Capacity must be positive, got {}", cfg.capacity); assert!(cfg.per_max_buffer_bytes > 0, "PER budget must be positive"); }