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

feat: GPU TOML profile system — remove ALL hardcoded VRAM if/else chains

Browse Source 
Created config/gpu/{default,rtx3050,h100,a100}.toml with all GPU-specific
parameters: batch_size, num_atoms, buffer_size, hidden_dim_base,
replay_buffer_vram_fraction, gpu_n_episodes, gpu_timesteps_per_episode,
cuda_stack_bytes.

GpuProfile::load() auto-detects GPU by device name, falls back to
embedded defaults (include_str!). Override via FOXHUNT_GPU_PROFILE env.

Removed dead code:
- detect_vram_mb(), vram_scaled_hidden_dims(), vram_scaled_base_dim(),
  resolve_hidden_dim_base() + 18 tests for these functions

All callers updated: train_baseline_rl, DQNTrainer constructor,
PPO trainer, smoke tests, pipeline tests.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-03-21 11:39:40 +01:00
parent d4c9c9c34e
commit e4b7d2ffb0
10 changed files with 35 additions and 276 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
config/gpu/a100.toml
View File

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

2
config/gpu/default.toml
View File

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

2
config/gpu/h100.toml
View File

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

2
config/gpu/rtx3050.toml
View File

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

249
crates/ml-core/src/gpu/memory_profile.rs
View File

@@ -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::<usize>() {
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<usize> {
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();
}
}

4
crates/ml-core/src/gpu/profile.rs
View File

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

26
crates/ml/examples/train_baseline_rl.rs
View File

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

6
crates/ml/src/trainers/dqn/trainer/constructor.rs
View File

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

6
crates/ml/src/trainers/ppo.rs
View File

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

12
crates/ml/tests/smoke_test_real_data.rs
View File

@@ -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");