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

fix(dqn): replace hardcoded PER memory cap with dynamic VRAM-proportional budget

Browse Source 
The GPU PER replay buffer had a hardcoded 4 GB MAX_BYTES limit that
rejected the auto-sizer's 10M-entry proposal on H100 (80 GB VRAM).
Now per_max_buffer_bytes() computes 20% of total VRAM (min 1 GB) and
flows through OptimalReplayConfig → DQNConfig → GpuReplayBufferConfig
so both subsystems agree on the budget.

Also fixes misleading regime detection log (indices 211/219 → 40/41)
and renames dqn_config_2025 → dqn_default_config.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-03-11 10:03:21 +01:00
parent 77029165c5
commit b25ce1dcbf
11 changed files with 117 additions and 39 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -427,6 +427,19 @@ pub struct GpuMemoryInfo {
pub used_memory_mb: f64,
}
/// Bundled output from [`GpuHardwareInfo::optimal_replay_config`].
///
/// Ensures callers always receive both the capacity **and** the PER memory
/// budget together, preventing the temporal-coupling bug where the auto-sizer
/// proposes a capacity that the PER pre-flight check rejects.
#[derive(Debug, Clone, Copy)]
pub struct OptimalReplayConfig {
/// VRAM-optimal replay buffer capacity (clamped 100K..10M).
pub capacity: usize,
/// Maximum PER buffer allocation in bytes (20% of total VRAM, min 1 GB).
pub per_max_buffer_bytes: usize,
}
/// Full GPU hardware info for dynamic scaling decisions.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct GpuHardwareInfo {
@@ -492,7 +505,7 @@ impl GpuHardwareInfo {
aligned.min(MAX_EPISODES)
}
/// Calculate VRAM-optimal replay buffer capacity.
/// Calculate VRAM-optimal replay buffer capacity and PER memory budget.
///
/// Allocates `vram_fraction` of free GPU memory to the replay buffer.
/// Each transition stores: 2 × state_dim × f32 (state + next_state)
@@ -501,7 +514,12 @@ impl GpuHardwareInfo {
/// Enforces guardrails:
/// - MIN_REPLAY_CAPACITY = 100_000 (below this, sample efficiency degrades)
/// - MAX_REPLAY_CAPACITY = 10_000_000 (above this, priority staleness dominates)
pub fn optimal_replay_capacity(&self, state_dim: usize, vram_fraction: f64) -> usize {
/// - 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;
@@ -509,7 +527,23 @@ impl GpuHardwareInfo {
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);
raw_capacity.clamp(MIN_REPLAY_CAPACITY, MAX_REPLAY_CAPACITY)
// 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 }
}
/// Maximum PER replay buffer allocation (bytes) proportional to total VRAM.
///
/// Formula: 20% of total GPU memory (min 1 GB).
/// 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);
(max_mb * 1024.0 * 1024.0) as usize
}
}
@@ -1073,9 +1107,10 @@ mod tests {
free_memory_mb: 72_000.0,
sm_count: 132,
};
let cap = hw.optimal_replay_capacity(56, 0.70);
assert!(cap >= 500_000, "H100 should get at least 500K capacity, got {}", cap);
assert!(cap <= 10_000_000, "Should be capped at max, got {}", cap);
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.per_max_buffer_bytes > 0, "PER budget must be positive");
}
#[test]
@@ -1086,7 +1121,9 @@ mod tests {
free_memory_mb: 2_000.0,
sm_count: 20,
};
let cap = hw.optimal_replay_capacity(48, 0.70);
assert!(cap >= 100_000, "Small GPU should get at least min capacity");
let cfg = hw.optimal_replay_config(48, 0.70);
assert!(cfg.capacity >= 100_000, "Small GPU should get at least min capacity");
// 20% of 4096 MB = ~819 MB, but min is 1 GB
assert_eq!(cfg.per_max_buffer_bytes, 1024 * 1024 * 1024, "Small GPU PER budget should be min 1 GB");
}
}

3
crates/ml-core/src/memory_optimization/mod.rs
View File

@@ -11,7 +11,8 @@ pub mod quantization;
pub use auto_batch_size::{
detect_gpu_hardware, detect_gpu_memory, sm_count_from_device_name, AutoBatchSizer,
BatchSizeConfig, GpuHardwareInfo, GpuMemoryInfo, ModelPrecision, OptimizerType,
BatchSizeConfig, GpuHardwareInfo, GpuMemoryInfo, ModelPrecision, OptimalReplayConfig,
OptimizerType,
};
pub use lazy_loader::{LazyCheckpointLoader, LoadStrategy};
pub use oom_detection::{extract_oom_size, is_oom_error};

10
crates/ml-dqn/src/dqn.rs
View File

@@ -107,6 +107,11 @@ pub struct DQNConfig {
pub per_beta_max: f64,
/// Number of steps to anneal beta from start to max
pub per_beta_annealing_steps: usize,
/// Maximum memory allocation for GPU/CPU PER replay buffer (bytes).
/// Computed at runtime from detected GPU VRAM via
/// `GpuHardwareInfo::per_max_buffer_bytes()`.
/// Default: 4 GB (for tests/CPU). Overridden by `DqnTrainer` on CUDA.
pub per_max_memory_bytes: usize,
// Dueling Networks configuration
/// Whether to use dueling network architecture
@@ -260,6 +265,7 @@ impl Default for DQNConfig {
per_beta_start: 0.4,
per_beta_max: 1.0,
per_beta_annealing_steps: 100_000,
per_max_memory_bytes: 4 * 1024 * 1024 * 1024,
use_dueling: true,
dueling_hidden_dim: 128,
use_distributional: true, // BUG #36 FIXED: C51 re-enabled
@@ -556,6 +562,7 @@ impl DQNConfig {
per_beta_start: 0.4,
per_beta_max: 1.0,
per_beta_annealing_steps: 100000,
per_max_memory_bytes: 4 * 1024 * 1024 * 1024,
use_dueling: true,
dueling_hidden_dim: 512, // Wave 11.4: Large hidden dim for aggressive config
use_distributional: true,
@@ -625,6 +632,7 @@ impl DQNConfig {
per_beta_start: 0.4,
per_beta_max: 1.0,
per_beta_annealing_steps: 100000,
per_max_memory_bytes: 4 * 1024 * 1024 * 1024,
use_dueling: false,
dueling_hidden_dim: 64,
use_distributional: false,
@@ -703,6 +711,7 @@ impl DQNConfig {
per_beta_start: 0.4,
per_beta_max: 1.0,
per_beta_annealing_steps: 100000,
per_max_memory_bytes: 4 * 1024 * 1024 * 1024,
use_dueling: false,
dueling_hidden_dim: 64,
use_distributional: false,
@@ -1289,6 +1298,7 @@ impl DQN {
config.per_beta_start,
config.per_beta_max,
config.per_beta_annealing_steps,
config.per_max_memory_bytes,
&device,
)?
} else {

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

@@ -292,6 +292,10 @@ pub struct GpuReplayBufferConfig {
pub beta_max: f32,
pub beta_annealing_steps: usize,
pub epsilon: f32,
/// Maximum VRAM allocation for this buffer (bytes).
/// Computed at runtime from detected GPU VRAM via
/// `GpuHardwareInfo::per_max_buffer_bytes()` — no hardcoded cap.
pub max_memory_bytes: usize,
}
/// GPU-resident ring buffer for experience replay.
@@ -333,10 +337,6 @@ impl GpuReplayBuffer {
///
/// All tensors are allocated on the given device (CPU or CUDA) at creation.
/// For 100K capacity × 51 `state_dim`, this uses ~47 MB VRAM.
/// Maximum VRAM budget for the replay buffer (4 GB).
/// Prevents accidental OOM from misconfigured `capacity`/`state_dim`.
const MAX_BYTES: usize = 4 * 1024 * 1024 * 1024;
pub fn new(config: GpuReplayBufferConfig, device: &Device) -> Result<Self, MLError> {
let cap = config.capacity;
let sdim = config.state_dim;
@@ -345,11 +345,11 @@ impl GpuReplayBuffer {
// states + next_states: 2 * cap * sdim * 4 bytes
// actions: cap * 4 (u32), rewards + dones + priorities: 3 * cap * 4
let bytes_needed = 2 * cap * sdim * 4 + 4 * cap * 4;
if bytes_needed > Self::MAX_BYTES {
if bytes_needed > config.max_memory_bytes {
#[allow(clippy::integer_division)]
let needed_mb = bytes_needed / (1024 * 1024);
#[allow(clippy::integer_division)]
let limit_mb = Self::MAX_BYTES / (1024 * 1024);
let limit_mb = config.max_memory_bytes / (1024 * 1024);
return Err(MLError::ModelError(format!(
"GPU replay buffer would need {needed_mb} MB (limit {limit_mb} MB). Reduce capacity or state_dim.",
)));
@@ -795,6 +795,7 @@ mod tests {
beta_max: 1.0,
beta_annealing_steps: 100_000,
epsilon: 1e-6,
max_memory_bytes: 4 * 1024 * 1024 * 1024,
};
let buf = GpuReplayBuffer::new(config, &Device::Cpu);
assert!(buf.is_ok());
@@ -815,6 +816,7 @@ mod tests {
beta_max: 1.0,
beta_annealing_steps: 1000,
epsilon: 1e-6,
max_memory_bytes: 4 * 1024 * 1024 * 1024,
};
let mut buf = GpuReplayBuffer::new(config, &Device::Cpu).expect("buf");
@@ -845,6 +847,7 @@ mod tests {
beta_max: 1.0,
beta_annealing_steps: 1000,
epsilon: 1e-6,
max_memory_bytes: 4 * 1024 * 1024 * 1024,
};
let mut buf = GpuReplayBuffer::new(config, &Device::Cpu).expect("buf");
buf.step();
@@ -864,6 +867,7 @@ mod tests {
beta_max: 1.0,
beta_annealing_steps: 1000,
epsilon: 1e-6,
max_memory_bytes: 4 * 1024 * 1024 * 1024,
};
let device = Device::Cpu;
let mut buf = GpuReplayBuffer::new(config, &device).expect("buf");
@@ -896,6 +900,7 @@ mod tests {
beta_max: 1.0,
beta_annealing_steps: 1000,
epsilon: 1e-6,
max_memory_bytes: 4 * 1024 * 1024 * 1024,
};
let device = Device::Cpu;
let mut buf = GpuReplayBuffer::new(config, &device).expect("buf");
@@ -938,6 +943,7 @@ mod tests {
beta_max: 1.0,
beta_annealing_steps: 1000,
epsilon: 1e-6,
max_memory_bytes: 4 * 1024 * 1024 * 1024,
};
let device = Device::Cpu;
let mut buf = GpuReplayBuffer::new(config, &device).expect("buf");

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

@@ -186,6 +186,7 @@ impl ReplayBufferType {
beta: f64,
beta_max: f64,
beta_annealing_steps: usize,
max_memory_bytes: usize,
device: &candle_core::Device,
) -> Result<Self, MLError> {
use crate::gpu_replay_buffer::{GpuReplayBuffer, GpuReplayBufferConfig};
@@ -198,6 +199,7 @@ impl ReplayBufferType {
beta_max: beta_max as f32,
beta_annealing_steps,
epsilon: 1e-6,
max_memory_bytes,
};
let buffer = GpuReplayBuffer::new(config, device)?;
@@ -219,9 +221,10 @@ impl ReplayBufferType {
beta: f64,
beta_max: f64,
beta_annealing_steps: usize,
max_memory_bytes: usize,
device: &candle_core::Device,
) -> Result<Self, MLError> {
match Self::new_gpu_prioritized(capacity, state_dim, alpha, beta, beta_max, beta_annealing_steps, device) {
match Self::new_gpu_prioritized(capacity, state_dim, alpha, beta, beta_max, beta_annealing_steps, max_memory_bytes, device) {
Ok(buf) => {
tracing::info!(
"GPU PER replay buffer allocated ({} capacity, {} state_dim)",
@@ -238,8 +241,7 @@ impl ReplayBufferType {
// Pre-flight: estimate CPU PER memory before attempting fallback.
// SegmentTree allocates 2 * next_power_of_two(capacity) f32 entries.
// PrioritizedReplayBuffer allocates capacity Option<Experience> slots.
const MAX_CPU_PER_BYTES: usize = 4 * 1024 * 1024 * 1024; // 4 GB
// Uses the same dynamic memory limit as GPU PER — no hardcoded cap.
let tree_elems = capacity
.checked_next_power_of_two()
.and_then(|p| p.checked_mul(2))
@@ -251,13 +253,13 @@ impl ReplayBufferType {
);
let estimated_bytes = tree_bytes.saturating_add(exp_bytes);
if estimated_bytes > MAX_CPU_PER_BYTES {
if estimated_bytes > max_memory_bytes {
return Err(MLError::ModelError(format!(
"GPU PER failed ({}) and CPU PER fallback would need ~{} MB (limit {} MB). \
Reduce capacity or state_dim.",
e,
estimated_bytes / (1024 * 1024),
MAX_CPU_PER_BYTES / (1024 * 1024),
max_memory_bytes / (1024 * 1024),
)));
}

25
crates/ml/src/trainers/dqn/config.rs
View File

@@ -1264,10 +1264,10 @@ use crate::dqn::dqn::DQNConfig;
///
/// # Not Recommended For
///
/// - HFT (use dqn_config_2025_hft with attention layers enabled)
/// - Very small datasets (use dqn_config_2025_conservative with reduced capacity)
/// - HFT (use dqn_hft_config with attention layers enabled)
/// - Very small datasets (use dqn_conservative_config with reduced capacity)
/// - Offline RL (disable exploration)
pub(crate) fn dqn_config_2025() -> DQNConfig {
pub(crate) fn dqn_default_config() -> DQNConfig {
DQNConfig {
// Network Architecture
state_dim: 48, // 42 market + 3 portfolio = 45, aligned to 48 for tensor cores
@@ -1293,6 +1293,7 @@ pub(crate) fn dqn_config_2025() -> DQNConfig {
per_beta_start: 0.4,
per_beta_max: 1.0,
per_beta_annealing_steps: 100_000,
per_max_memory_bytes: 4 * 1024 * 1024 * 1024,
// Exploration Strategy
epsilon_start: 1.0,
@@ -1347,13 +1348,13 @@ pub(crate) fn dqn_config_2025() -> DQNConfig {
/// HFT-optimized variant with attention layers and faster updates
///
/// Differences from standard 2025 config:
/// Differences from default config:
/// - Attention layers enabled for temporal pattern recognition (requires network config)
/// - Faster target updates (tau = 0.005)
/// - Smaller warmup (1000 steps)
/// - Larger batch size (512) for throughput
pub(crate) fn dqn_config_2025_hft() -> DQNConfig {
let mut config = dqn_config_2025();
pub(crate) fn dqn_hft_config() -> DQNConfig {
let mut config = dqn_default_config();
// HFT-specific overrides
config.warmup_steps = 1000;
@@ -1365,13 +1366,13 @@ pub(crate) fn dqn_config_2025_hft() -> DQNConfig {
/// Conservative variant for small datasets or cautious exploration
///
/// Differences from standard 2025 config:
/// Differences from default config:
/// - Smaller network (256-128-64)
/// - Lower learning rate (5e-5)
/// - More aggressive exploration decay
/// - Smaller batch size (256)
pub(crate) fn dqn_config_2025_conservative() -> DQNConfig {
let mut config = dqn_config_2025();
pub(crate) fn dqn_conservative_config() -> DQNConfig {
let mut config = dqn_default_config();
// Conservative overrides
config.hidden_dims = vec![256, 128, 64];
@@ -1384,13 +1385,13 @@ pub(crate) fn dqn_config_2025_conservative() -> DQNConfig {
/// Aggressive variant for maximum exploration and learning
///
/// Differences from standard 2025 config:
/// Differences from default config:
/// - Larger network (768-512-256)
/// - Higher learning rate (3e-4)
/// - Slower exploration decay
/// - Larger batch size (2048)
pub(crate) fn dqn_config_2025_aggressive() -> DQNConfig {
let mut config = dqn_config_2025();
pub(crate) fn dqn_aggressive_config() -> DQNConfig {
let mut config = dqn_default_config();
// Aggressive overrides
config.hidden_dims = vec![768, 512, 256];

1
crates/ml/src/trainers/dqn/smoke_tests/gpu_residency.rs
View File

@@ -25,6 +25,7 @@ fn test_buffer_config(capacity: usize, state_dim: usize) -> GpuReplayBufferConfi
beta_max: 1.0,
beta_annealing_steps: 1000,
epsilon: 1e-6,
max_memory_bytes: 4 * 1024 * 1024 * 1024,
}
}

1
crates/ml/src/trainers/dqn/smoke_tests/performance.rs
View File

@@ -62,6 +62,7 @@ async fn test_per_sample_latency() -> anyhow::Result<()> {
beta_max: 1.0,
beta_annealing_steps: 1000,
epsilon: 1e-6,
max_memory_bytes: 4 * 1024 * 1024 * 1024,
};
let mut buf = GpuReplayBuffer::new(config, &device)?;

2
crates/ml/src/trainers/dqn/smoke_tests/training_stability.rs
View File

@@ -135,6 +135,7 @@ async fn test_per_weights_in_valid_range() -> anyhow::Result<()> {
beta_max: 1.0,
beta_annealing_steps: 1000,
epsilon: 1e-6,
max_memory_bytes: 4 * 1024 * 1024 * 1024,
};
let mut buf = GpuReplayBuffer::new(config, &device)?;
@@ -175,6 +176,7 @@ async fn test_per_indices_in_valid_range() -> anyhow::Result<()> {
beta_max: 1.0,
beta_annealing_steps: 1000,
epsilon: 1e-6,
max_memory_bytes: 4 * 1024 * 1024 * 1024,
};
let mut buf = GpuReplayBuffer::new(config, &device)?;

29
crates/ml/src/trainers/dqn/trainer.rs
View File

@@ -394,25 +394,29 @@ impl DQNTrainer {
// Dynamic replay buffer sizing: scale replay capacity to available VRAM.
// Only activates when replay_buffer_vram_fraction > 0 and GPU is detected.
// Also computes the PER memory budget from actual VRAM — no hardcoded caps.
let mut per_max_memory_bytes: usize = 4 * 1024 * 1024 * 1024; // CPU fallback: 4 GB
if hyperparams.replay_buffer_vram_fraction > 0.0 && device.is_cuda() {
use ml_core::memory_optimization::detect_gpu_hardware;
match detect_gpu_hardware() {
Ok(hw) => {
let raw_sd = if hyperparams.mbp10_data_dir.is_some() { 53 } else { 45 };
let aligned_sd = crate::dqn::mixed_precision::align_dim_for_tensor_cores(raw_sd, &device);
let optimal = hw.optimal_replay_capacity(
let replay_cfg = hw.optimal_replay_config(
aligned_sd,
hyperparams.replay_buffer_vram_fraction,
);
if optimal != hyperparams.buffer_size {
per_max_memory_bytes = replay_cfg.per_max_buffer_bytes;
if replay_cfg.capacity != hyperparams.buffer_size {
info!(
"AutoReplaySizer: replay buffer {} -> {} (VRAM={:.0}MB, fraction={:.0}%)",
"AutoReplaySizer: replay buffer {} -> {} (VRAM={:.0}MB, fraction={:.0}%, PER budget={:.0}MB)",
hyperparams.buffer_size,
optimal,
replay_cfg.capacity,
hw.free_memory_mb,
hyperparams.replay_buffer_vram_fraction * 100.0,
per_max_memory_bytes as f64 / (1024.0 * 1024.0),
);
hyperparams.buffer_size = optimal;
hyperparams.buffer_size = replay_cfg.capacity;
}
}
Err(e) => {
@@ -422,6 +426,16 @@ impl DQNTrainer {
);
}
}
} else if device.is_cuda() {
// No auto-sizer, but still compute PER budget from VRAM
use ml_core::memory_optimization::detect_gpu_hardware;
if let Ok(hw) = detect_gpu_hardware() {
per_max_memory_bytes = hw.per_max_buffer_bytes();
} else {
// GPU detection failed — keep CPU default (4 GB)
}
} else {
// CPU device — keep default 4 GB PER budget
}
info!(
@@ -494,6 +508,7 @@ impl DQNTrainer {
per_beta_start: hyperparams.per_beta_start,
per_beta_max: 1.0,
per_beta_annealing_steps: hyperparams.epochs * 2000, // ~2000 steps/epoch (130k bars / ~64 batch_size)
per_max_memory_bytes,
// Wave 2.1: Dueling Networks (ENABLED BY DEFAULT - Wave 6.4)
use_dueling: hyperparams.use_dueling,
@@ -547,8 +562,8 @@ impl DQNTrainer {
// Create DQN agent
let agent = if hyperparams.enable_regime_qnetwork {
info!("Creating regime-conditional DQN with 3 heads (Trending, Ranging, Volatile)");
info!(" - Regime detection: ADX (index 211) + Entropy (index 219)");
info!(" - Classification: Trending (ADX>25), Volatile (ADX≤25 & Entropy>0.7), Ranging (ADX≤25 & Entropy≤0.7)");
info!(" - Regime detection: ADX (raw index 40) + CUSUM direction (raw index 41)");
info!(" - Classification: Trending (ADX>0.25), Volatile (ADX≤0.25 & |CUSUM|>0.7), Ranging (otherwise)");
let regime_agent = RegimeConditionalDQN::new_on_device(config, device.clone())
.map_err(|e| anyhow::anyhow!("Failed to create regime-conditional DQN: {}", e))?;
DQNAgentType::RegimeConditional(regime_agent)

2
crates/ml/tests/gpu_per_integration_test.rs
View File

@@ -20,6 +20,7 @@ fn test_config(capacity: usize, state_dim: usize) -> GpuReplayBufferConfig {
beta_max: 1.0,
beta_annealing_steps: 1000,
epsilon: 1e-6,
max_memory_bytes: 4 * 1024 * 1024 * 1024,
}
}
@@ -252,6 +253,7 @@ fn test_gpu_per_oom_rejects_absurd_capacity() {
0.4,
1.0,
1000,
4 * 1024 * 1024 * 1024, // 4 GB limit for test
&device,
);