diff --git a/config/gpu/a100.toml b/config/gpu/a100.toml index 32cff0ca8..1bda11d4e 100644 --- a/config/gpu/a100.toml +++ b/config/gpu/a100.toml @@ -3,6 +3,8 @@ batch_size = 512 num_atoms = 51 buffer_size = 200_000 +hidden_dim_base = 256 +replay_buffer_vram_fraction = 0.70 [experience] gpu_n_episodes = 256 diff --git a/config/gpu/default.toml b/config/gpu/default.toml index 5819f43b5..4bf24c935 100644 --- a/config/gpu/default.toml +++ b/config/gpu/default.toml @@ -3,6 +3,8 @@ batch_size = 256 num_atoms = 21 buffer_size = 50_000 +hidden_dim_base = 256 +replay_buffer_vram_fraction = 0.50 [experience] gpu_n_episodes = 64 diff --git a/config/gpu/h100.toml b/config/gpu/h100.toml index 228b492c8..6a4c43d9d 100644 --- a/config/gpu/h100.toml +++ b/config/gpu/h100.toml @@ -3,6 +3,8 @@ batch_size = 1024 num_atoms = 51 buffer_size = 500_000 +hidden_dim_base = 256 +replay_buffer_vram_fraction = 0.70 [experience] gpu_n_episodes = 256 diff --git a/config/gpu/rtx3050.toml b/config/gpu/rtx3050.toml index ff0c399aa..e0d19348c 100644 --- a/config/gpu/rtx3050.toml +++ b/config/gpu/rtx3050.toml @@ -3,6 +3,8 @@ batch_size = 64 num_atoms = 11 buffer_size = 5_000 +hidden_dim_base = 256 +replay_buffer_vram_fraction = 0.40 [experience] gpu_n_episodes = 16 diff --git a/crates/ml-core/src/gpu/memory_profile.rs b/crates/ml-core/src/gpu/memory_profile.rs index 7a34d92ae..489bb8b0b 100644 --- a/crates/ml-core/src/gpu/memory_profile.rs +++ b/crates/ml-core/src/gpu/memory_profile.rs @@ -5,7 +5,6 @@ //! optimal batch size for the detected hardware. use serde::{Deserialize, Serialize}; -use tracing::info; /// Memory footprint estimate for an ML model. /// @@ -129,99 +128,10 @@ pub mod estimates { }; } -/// Resolve default hidden_dim_base from available GPU VRAM. -/// -/// Returns a conservative base dimension scaled to GPU capacity. -/// Tiers: -/// - ≤4 GB (RTX 3050 Ti): 256 -/// - 5–12 GB (RTX 3060–4080): 512 -/// - 13–25 GB (L4, RTX 3090): 768 -/// - 26+ GB (L40S, H100): 2048 -pub fn resolve_hidden_dim_base(gpu_vram_mb: f64) -> usize { - match gpu_vram_mb as u64 { - 0..=4096 => 256, - 4097..=12288 => 512, - 12289..=25600 => 768, - _ => 2048, - } -} - -/// Detect available GPU VRAM in MB. -/// -/// Resolution order: -/// 1. `GPU_MEMORY_MB` env var (for K8s containers where nvidia-smi may not reflect limits) -/// 2. Cached `GpuCapabilities` from nvidia-smi -/// 3. Falls back to 0 on CPU -pub fn detect_vram_mb() -> usize { - // Environment override — K8s resource limits may not be visible to nvidia-smi - if let Ok(vram_str) = std::env::var("GPU_MEMORY_MB") { - if let Ok(vram) = vram_str.parse::() { - info!("GPU VRAM from GPU_MEMORY_MB env: {} MB", vram); - return vram; - } - } - let caps = super::capabilities::cached_capabilities(); - if caps.is_cuda { - caps.total_vram_mb as usize - } else { - 0 - } -} - -/// Scale hidden dimensions for a model based on available GPU VRAM. -/// -/// Returns a vector of hidden layer sizes appropriate for the detected GPU. -/// Falls back to conservative defaults on CPU or small GPUs. -/// -/// Model-specific topologies: -/// - **DQN**: 3-layer MLP `[base, base/2, base/4]` — deeper networks saturate GPU compute -/// - **PPO policy**: 2-layer MLP `[base, base/2]` — lightweight for fast rollouts -/// - **PPO value**: 3-layer MLP `[base, base*3/4, base/2]` — larger critic for stable training -/// -/// VRAM tiers: -/// | VRAM | GPU examples | DQN base | PPO policy base | PPO value base | -/// |-------------|---------------------|----------|----------------|---------------| -/// | >= 40 GB | H100, L40S | 2048 | 1024 | 2048 | -/// | 16–39 GB | A100 40GB, L4 | 1024 | 512 | 1024 | -/// | 8–15 GB | RTX 4090 | 512 | 256 | 512 | -/// | < 8 GB | CPU / RTX 3050 Ti | 256 | 128 | 256 | -pub fn vram_scaled_hidden_dims(vram_mb: usize, model_type: &str) -> Vec { - match model_type { - "dqn" => match vram_mb { - v if v >= 40_000 => vec![2048, 1024, 512], - v if v >= 16_000 => vec![1024, 512, 256], - v if v >= 8_000 => vec![512, 256, 128], - _ => vec![256, 128, 64], - }, - "ppo_policy" => match vram_mb { - v if v >= 40_000 => vec![1024, 512], - v if v >= 16_000 => vec![512, 256], - v if v >= 8_000 => vec![256, 128], - _ => vec![128, 64], - }, - "ppo_value" => match vram_mb { - v if v >= 40_000 => vec![2048, 1536, 1024], - v if v >= 16_000 => vec![1024, 768, 512], - v if v >= 8_000 => vec![512, 384, 256], - _ => vec![256, 192, 128], - }, - _ => vec![256, 128, 64], - } -} - -/// Return the VRAM-appropriate `hidden_dim_base` for a model type. -/// -/// This is the scalar base dimension that trainers expand into layer-specific -/// widths via their own formulas (e.g. DQN: `[base, base/2, base/4]`). -/// -/// Differs from [`resolve_hidden_dim_base`] by being model-aware: DQN and PPO -/// value networks can use larger bases than PPO policy networks because their -/// architectures have different parameter efficiency. -pub fn vram_scaled_base_dim(vram_mb: usize, model_type: &str) -> usize { - let dims = vram_scaled_hidden_dims(vram_mb, model_type); - // Base dim is the first (largest) hidden layer - dims.first().copied().unwrap_or(256) -} +// detect_vram_mb, vram_scaled_hidden_dims, vram_scaled_base_dim, +// resolve_hidden_dim_base — ALL REMOVED. +// Replaced by GpuProfile::load() from gpu::profile (TOML config files). +// Replaced by GpuProfile::load() from gpu::profile (TOML config files). /// Compute total parameter count for a fully-connected network given layer dimensions. /// @@ -257,15 +167,6 @@ mod tests { assert!(size > 180.0 && size < 200.0, "TFT should be ~190.7 MB, got {}", size); } - #[test] - fn test_resolve_hidden_dim_base_tiers() { - assert_eq!(resolve_hidden_dim_base(3700.0), 256); // 4GB GPU - assert_eq!(resolve_hidden_dim_base(11000.0), 512); // 12GB GPU - assert_eq!(resolve_hidden_dim_base(23000.0), 768); // 24GB GPU - assert_eq!(resolve_hidden_dim_base(48000.0), 2048); // 48GB GPU - assert_eq!(resolve_hidden_dim_base(80000.0), 2048); // 80GB GPU - } - #[test] fn test_network_param_count() { // [54, 256, 128, 64, 45] → 54*256+256 + 256*128+128 + 128*64+64 + 64*45+45 @@ -296,146 +197,4 @@ mod tests { } } - // ----------------------------------------------------------------------- - // vram_scaled_hidden_dims tests - // ----------------------------------------------------------------------- - - #[test] - fn test_vram_scaled_hidden_dims_dqn_h100() { - let dims = vram_scaled_hidden_dims(80_000, "dqn"); - assert_eq!(dims, vec![2048, 1024, 512]); - } - - #[test] - fn test_vram_scaled_hidden_dims_dqn_l40s() { - let dims = vram_scaled_hidden_dims(48_000, "dqn"); - assert_eq!(dims, vec![2048, 1024, 512]); - } - - #[test] - fn test_vram_scaled_hidden_dims_dqn_a100_40gb() { - let dims = vram_scaled_hidden_dims(40_000, "dqn"); - assert_eq!(dims, vec![2048, 1024, 512]); - } - - #[test] - fn test_vram_scaled_hidden_dims_dqn_l4_24gb() { - let dims = vram_scaled_hidden_dims(24_000, "dqn"); - assert_eq!(dims, vec![1024, 512, 256]); - } - - #[test] - fn test_vram_scaled_hidden_dims_dqn_rtx4090() { - let dims = vram_scaled_hidden_dims(16_000, "dqn"); - assert_eq!(dims, vec![1024, 512, 256]); - } - - #[test] - fn test_vram_scaled_hidden_dims_dqn_small_gpu() { - let dims = vram_scaled_hidden_dims(4_000, "dqn"); - assert_eq!(dims, vec![256, 128, 64]); - } - - #[test] - fn test_vram_scaled_hidden_dims_dqn_cpu() { - let dims = vram_scaled_hidden_dims(0, "dqn"); - assert_eq!(dims, vec![256, 128, 64]); - } - - #[test] - fn test_vram_scaled_hidden_dims_ppo_policy_h100() { - let dims = vram_scaled_hidden_dims(80_000, "ppo_policy"); - assert_eq!(dims, vec![1024, 512]); - } - - #[test] - fn test_vram_scaled_hidden_dims_ppo_policy_cpu() { - let dims = vram_scaled_hidden_dims(0, "ppo_policy"); - assert_eq!(dims, vec![128, 64]); - } - - #[test] - fn test_vram_scaled_hidden_dims_ppo_value_h100() { - let dims = vram_scaled_hidden_dims(80_000, "ppo_value"); - assert_eq!(dims, vec![2048, 1536, 1024]); - } - - #[test] - fn test_vram_scaled_hidden_dims_ppo_value_cpu() { - let dims = vram_scaled_hidden_dims(0, "ppo_value"); - assert_eq!(dims, vec![256, 192, 128]); - } - - #[test] - fn test_vram_scaled_hidden_dims_unknown_model_type() { - let dims = vram_scaled_hidden_dims(80_000, "unknown"); - assert_eq!(dims, vec![256, 128, 64], "unknown model type should return conservative defaults"); - } - - #[test] - fn test_vram_scaled_hidden_dims_monotonic_scaling() { - // Larger VRAM should yield >= dimensions for every model type - let vram_tiers = [0, 4_000, 8_000, 16_000, 40_000, 80_000]; - for model in &["dqn", "ppo_policy", "ppo_value"] { - let mut prev_first = 0usize; - for &vram in &vram_tiers { - let dims = vram_scaled_hidden_dims(vram, model); - let first = dims.first().copied().unwrap_or(0); - assert!( - first >= prev_first, - "model={} vram={}: first dim {} < previous {}", - model, vram, first, prev_first - ); - prev_first = first; - } - } - } - - #[test] - fn test_vram_scaled_hidden_dims_descending_layers() { - // All model types should have descending layer sizes (wider to narrower) - for model in &["dqn", "ppo_policy", "ppo_value"] { - for &vram in &[0, 8_000, 48_000, 80_000] { - let dims = vram_scaled_hidden_dims(vram, model); - for pair in dims.windows(2) { - assert!( - pair.first().copied().unwrap_or(0) >= pair.last().copied().unwrap_or(0), - "model={} vram={}: non-descending dims {:?}", - model, vram, dims - ); - } - } - } - } - - #[test] - fn test_vram_scaled_base_dim_matches_first_hidden() { - for model in &["dqn", "ppo_policy", "ppo_value"] { - for &vram in &[0, 8_000, 48_000, 80_000] { - let dims = vram_scaled_hidden_dims(vram, model); - let base = vram_scaled_base_dim(vram, model); - assert_eq!( - base, - dims.first().copied().unwrap_or(256), - "base_dim should equal first hidden dim for model={} vram={}", - model, vram - ); - } - } - } - - #[test] - fn test_vram_scaled_base_dim_dqn_tiers() { - assert_eq!(vram_scaled_base_dim(0, "dqn"), 256); - assert_eq!(vram_scaled_base_dim(8_000, "dqn"), 512); - assert_eq!(vram_scaled_base_dim(16_000, "dqn"), 1024); - assert_eq!(vram_scaled_base_dim(48_000, "dqn"), 2048); - } - - #[test] - fn test_detect_vram_mb_no_gpu() { - // In CI without GPU, detect_vram_mb returns 0 - // We just verify it doesn't panic - let _ = detect_vram_mb(); - } } diff --git a/crates/ml-core/src/gpu/profile.rs b/crates/ml-core/src/gpu/profile.rs index d21c85bbf..64fcb656e 100644 --- a/crates/ml-core/src/gpu/profile.rs +++ b/crates/ml-core/src/gpu/profile.rs @@ -39,6 +39,8 @@ pub struct TrainingProfile { pub batch_size: usize, pub num_atoms: usize, pub buffer_size: usize, + pub hidden_dim_base: usize, + pub replay_buffer_vram_fraction: f64, } /// GPU experience collection parameters. @@ -199,6 +201,8 @@ impl GpuProfile { batch_size: 256, num_atoms: 21, buffer_size: 50_000, + hidden_dim_base: 256, + replay_buffer_vram_fraction: 0.50, }, experience: ExperienceProfile { gpu_n_episodes: 64, diff --git a/crates/ml/examples/train_baseline_rl.rs b/crates/ml/examples/train_baseline_rl.rs index 4574ae662..b2bb34ddf 100644 --- a/crates/ml/examples/train_baseline_rl.rs +++ b/crates/ml/examples/train_baseline_rl.rs @@ -113,8 +113,7 @@ use baseline_common::completion::{write_failure_marker, write_success_marker, Co use baseline_common::{load_all_bars, spread_cost_bps}; use common::metrics::{server as metrics_server, training_metrics as metrics}; use ml::features::extraction::{extract_ml_features, FeatureVector}; -use ml::gpu::memory_profile::{detect_vram_mb, vram_scaled_base_dim}; -use ml::gpu::profile::GpuProfile; +use ml_core::gpu::profile::GpuProfile; use ml::types::OHLCVBar; use ml::walk_forward::{generate_walk_forward_windows, NormStats, WalkForwardConfig}; @@ -482,17 +481,15 @@ fn train_dqn_fold( // hidden_dim_base: hyperopt JSON takes priority. When absent, use VRAM-aware // scaling so that large GPUs (L40S 48GB, H100 80GB) get proportionally wider // networks instead of the tiny CPU defaults. - let hp_hidden_base = hp_usize(hp, "hidden_dim_base"); - let dqn_hidden_base = hp_hidden_base.or_else(|| { - let vram = detect_vram_mb(); - let base = vram_scaled_base_dim(vram, "dqn"); - info!(" [DQN] VRAM-scaled hidden_dim_base: {} (VRAM: {} MB)", base, vram); - Some(base) - }); - // Load GPU profile for VRAM-aware defaults (replaces scattered if/else chains) let gpu_profile = GpuProfile::load(); + let hp_hidden_base = hp_usize(hp, "hidden_dim_base"); + let dqn_hidden_base = hp_hidden_base.or_else(|| { + info!(" [DQN] hidden_dim_base: {} (from GPU profile)", gpu_profile.training.hidden_dim_base); + Some(gpu_profile.training.hidden_dim_base) + }); + // When noisy_nets are enabled (the hyperopt default), epsilon-greedy must be // low or zero — otherwise epsilon=1.0 forces pure random actions for the entire // run, preventing the model from ever learning. Read exploration params from @@ -556,7 +553,7 @@ fn train_dqn_fold( // GPU PER is mandatory on CUDA. VRAM fraction controls AutoReplaySizer. // Small GPUs (RTX 3050 profile: buffer_size=5000 < 100K threshold) bypass // AutoReplaySizer entirely, so VRAM fraction is irrelevant for them. - replay_buffer_vram_fraction: if detect_vram_mb() >= 8192 { 0.70 } else { 0.0 }, + replay_buffer_vram_fraction: gpu_profile.training.replay_buffer_vram_fraction, // GPU experience collection: profile-driven defaults replace VRAM if/else chains gpu_n_episodes: gpu_profile.experience.gpu_n_episodes, gpu_timesteps_per_episode: gpu_profile.experience.gpu_timesteps_per_episode, @@ -672,10 +669,9 @@ fn train_ppo_fold( // the value network from the same base (value: [4*base, 3*base, 2*base, base, base/2]). let hp_ppo_hidden_base = hp_usize(hp, "hidden_dim_base"); let ppo_hidden_base = hp_ppo_hidden_base.or_else(|| { - let vram = detect_vram_mb(); - let base = vram_scaled_base_dim(vram, "ppo_policy"); - info!(" [PPO] VRAM-scaled hidden_dim_base: {} (VRAM: {} MB)", base, vram); - Some(base) + let profile = ml_core::gpu::profile::GpuProfile::load(); + info!(" [PPO] hidden_dim_base: {} (from GPU profile)", profile.training.hidden_dim_base); + Some(profile.training.hidden_dim_base) }); let hyperparams = PpoHyperparameters { diff --git a/crates/ml/src/trainers/dqn/trainer/constructor.rs b/crates/ml/src/trainers/dqn/trainer/constructor.rs index 29fb08f63..03b7b8546 100644 --- a/crates/ml/src/trainers/dqn/trainer/constructor.rs +++ b/crates/ml/src/trainers/dqn/trainer/constructor.rs @@ -59,10 +59,8 @@ impl DQNTrainer { vec![b, b] // Constant-width: no tapering, no silently-discarded narrow layer } None => { - let caps = crate::gpu::capabilities::cached_capabilities(); - let base = crate::gpu::memory_profile::resolve_hidden_dim_base( - caps.free_vram_mb, - ); + let profile = ml_core::gpu::profile::GpuProfile::load(); + let base = profile.training.hidden_dim_base; let b = crate::cuda_pipeline::align_to_tensor_cores(base); vec![b, b] // Constant-width: no tapering } diff --git a/crates/ml/src/trainers/ppo.rs b/crates/ml/src/trainers/ppo.rs index 57651ca89..665145268 100644 --- a/crates/ml/src/trainers/ppo.rs +++ b/crates/ml/src/trainers/ppo.rs @@ -278,10 +278,10 @@ impl PpoTrainer { // Resolve hidden_dim_base from GPU VRAM when not explicitly set if hyperparams.hidden_dim_base.is_none() && use_gpu { - let caps = cached_capabilities(); - let base = memory_profile::resolve_hidden_dim_base(caps.free_vram_mb); + let profile = ml_core::gpu::profile::GpuProfile::load(); + let base = profile.training.hidden_dim_base; hyperparams.hidden_dim_base = Some(base); - info!("PPO hidden_dim_base resolved from VRAM: {} (GPU: {})", base, caps.device_name); + info!("PPO hidden_dim_base resolved from GPU profile: {}", base); } // Compute accurate model overhead from actual network dimensions diff --git a/crates/ml/tests/smoke_test_real_data.rs b/crates/ml/tests/smoke_test_real_data.rs index 656af5040..a63004b04 100644 --- a/crates/ml/tests/smoke_test_real_data.rs +++ b/crates/ml/tests/smoke_test_real_data.rs @@ -803,15 +803,9 @@ async fn smoke_e2e_dqn_training_loop() { // finite loss, gradient flow, and action diversity without 400s/epoch overhead. hyperparams.max_training_steps_per_epoch = 8; // Curiosity disabled: kernel crashes on RTX 3050 (needs investigation) - // Dynamic C51 atom scaling: training kernel allocates DIST_SIZE(HIDDEN)*NUM_ATOMS - // per-thread arrays. With SHARED_H1=256 and NUM_ATOMS=51, that's 256*51*4=52KB - // per array — way beyond CUDA stack limits on consumer GPUs (4-8GB VRAM). - // Scale atoms to fit: <8GB→11, <16GB→21, >=16GB→51. - if gpu_vram_mb() < 8192 { - hyperparams.num_atoms = 11; - } else if gpu_vram_mb() < 16384 { - hyperparams.num_atoms = 21; - } + // C51 atom count from GPU profile (replaces hardcoded VRAM if/else) + let gpu_profile = ml_core::gpu::profile::GpuProfile::load(); + hyperparams.num_atoms = gpu_profile.training.num_atoms; let checkpoint_dir = tempfile::tempdir().expect("Failed to create temp dir"); let mut trainer = DQNTrainer::new(hyperparams).expect("Failed to create DQN trainer");