From 4ab6975c38798c15b5004f6388ca7388d040dfe1 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Mon, 23 Feb 2026 11:49:22 +0100 Subject: [PATCH] fix(ml): wire activation_multiplier, cache GPU detection, fix stale comments - Scale model_memory_mb by activation_multiplier in resolve_batch_size() so TFT (2.5x) gets proportionally smaller batches than DQN (1.0x) - Add cached_capabilities() with OnceLock to avoid spawning nvidia-smi on every call to CampaignConfig::dqn_default/ppo_default/PpoTrainer::new - Add PartialEq to GpuCapabilities and simplify serde roundtrip test - Update stale "RTX 3050 Ti" and "Exceeds GPU limit (230)" comments to reflect dynamic detection - Document Auto variant silent CPU fallback behavior (no callers affected) Co-Authored-By: Claude Opus 4.6 --- ml/src/gpu/capabilities.rs | 22 +++++++++++++++++++--- ml/src/gpu/memory_profile.rs | 4 +++- ml/src/gpu/mod.rs | 1 + ml/src/hyperopt/campaign.rs | 10 +++++----- ml/src/trainers/ppo.rs | 12 ++++++------ 5 files changed, 34 insertions(+), 15 deletions(-) diff --git a/ml/src/gpu/capabilities.rs b/ml/src/gpu/capabilities.rs index a78dca491..142de3ea8 100644 --- a/ml/src/gpu/capabilities.rs +++ b/ml/src/gpu/capabilities.rs @@ -5,12 +5,13 @@ use serde::{Deserialize, Serialize}; use std::process::Command; +use std::sync::OnceLock; use tracing::{info, warn}; /// Detected GPU hardware capabilities. /// /// Construct via `detect()` for real hardware or `cpu_fallback()` / `with_vram()` for testing. -#[derive(Debug, Clone, Serialize, Deserialize)] +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] pub struct GpuCapabilities { /// GPU device name (e.g., "NVIDIA GeForce RTX 3050 Ti Laptop GPU") pub device_name: String, @@ -112,6 +113,15 @@ impl GpuCapabilities { } } +/// Cached GPU capabilities -- detected once, reused across all callers. +/// +/// Use this instead of `GpuCapabilities::detect()` when you don't need fresh data. +/// The first call runs nvidia-smi; subsequent calls return the cached result. +pub fn cached_capabilities() -> &'static GpuCapabilities { + static CAPS: OnceLock = OnceLock::new(); + CAPS.get_or_init(GpuCapabilities::detect) +} + #[cfg(test)] mod tests { use super::*; @@ -159,7 +169,13 @@ mod tests { let caps = GpuCapabilities::with_vram("RTX 4090", 24576.0, 23000.0); let json = serde_json::to_string(&caps).unwrap(); let deserialized: GpuCapabilities = serde_json::from_str(&json).unwrap(); - assert_eq!(caps.device_name, deserialized.device_name); - assert_eq!(caps.total_vram_mb, deserialized.total_vram_mb); + assert_eq!(caps, deserialized); + } + + #[test] + fn test_cached_capabilities_returns_same_instance() { + let caps1 = cached_capabilities(); + let caps2 = cached_capabilities(); + assert_eq!(caps1, caps2); } } diff --git a/ml/src/gpu/memory_profile.rs b/ml/src/gpu/memory_profile.rs index 611119785..d792d1d4f 100644 --- a/ml/src/gpu/memory_profile.rs +++ b/ml/src/gpu/memory_profile.rs @@ -160,7 +160,9 @@ pub fn resolve_batch_size( return requested_batch_size; } - let model_size_mb = estimate.estimated_size_mb(); + // Scale model memory by activation multiplier to account for attention layers etc. + // Models with attention (TFT=2.5x, TLOB=2.0x) need more memory for intermediate activations. + let model_size_mb = estimate.estimated_size_mb() * estimate.activation_multiplier; let batch_config = BatchSizeConfig { model_memory_mb: model_size_mb, diff --git a/ml/src/gpu/mod.rs b/ml/src/gpu/mod.rs index 7d62b799c..181b5f4cb 100644 --- a/ml/src/gpu/mod.rs +++ b/ml/src/gpu/mod.rs @@ -19,6 +19,7 @@ use crate::MLError; pub enum DeviceConfig { Cpu, Cuda(usize), + /// Auto-detect: CUDA device 0 if available, else CPU (silent fallback -- never errors). Auto, } diff --git a/ml/src/hyperopt/campaign.rs b/ml/src/hyperopt/campaign.rs index 773aa3dee..2dd38b013 100644 --- a/ml/src/hyperopt/campaign.rs +++ b/ml/src/hyperopt/campaign.rs @@ -8,7 +8,7 @@ use std::path::PathBuf; use serde::{Deserialize, Serialize}; -use crate::gpu::capabilities::GpuCapabilities; +use crate::gpu::capabilities::cached_capabilities; use crate::gpu::memory_profile::{self, resolve_batch_size}; // Use canonical ModelType from crate root (re-exported from common) @@ -36,8 +36,8 @@ pub struct CampaignConfig { impl CampaignConfig { /// DQN campaign defaults (50 trials, SHA with η=3, GPU-adaptive batch size). pub fn dqn_default() -> Self { - let caps = GpuCapabilities::detect(); - let max_batch = resolve_batch_size(&caps, &memory_profile::estimates::DQN, 512); + let caps = cached_capabilities(); + let max_batch = resolve_batch_size(caps, &memory_profile::estimates::DQN, 512); Self { model_type: ModelType::DQN, num_trials: 50, @@ -51,8 +51,8 @@ impl CampaignConfig { /// PPO campaign defaults (30 trials, Hyperband, GPU-adaptive batch size). pub fn ppo_default() -> Self { - let caps = GpuCapabilities::detect(); - let max_batch = resolve_batch_size(&caps, &memory_profile::estimates::PPO, 512); + let caps = cached_capabilities(); + let max_batch = resolve_batch_size(caps, &memory_profile::estimates::PPO, 512); Self { model_type: ModelType::PPO, num_trials: 30, diff --git a/ml/src/trainers/ppo.rs b/ml/src/trainers/ppo.rs index 50500b92a..c1eed2744 100644 --- a/ml/src/trainers/ppo.rs +++ b/ml/src/trainers/ppo.rs @@ -2,7 +2,7 @@ //! //! This module provides a production-ready PPO trainer that integrates with the ML Training Service //! gRPC interface. It supports: -//! - GPU acceleration (RTX 3050 Ti) +//! - GPU acceleration (auto-detected) //! - Optional vectorized environments for 2-3x speedup //! - Hyperparameter configuration //! checkpoint management, and comprehensive metrics reporting. @@ -17,7 +17,7 @@ use tracing::{debug, info, warn}; use crate::dqn::TradingAction; use crate::gpu::DeviceConfig; -use crate::gpu::capabilities::GpuCapabilities; +use crate::gpu::capabilities::cached_capabilities; use crate::gpu::memory_profile::{self, resolve_batch_size}; use crate::ppo::gae::GAEConfig; use crate::ppo::ppo::{PPOConfig, PPO}; @@ -200,7 +200,7 @@ impl PpoTrainer { /// * `hyperparams` - Training hyperparameters from gRPC request /// * `state_dim` - State vector dimension (inferred from data) /// * `checkpoint_dir` - Directory for saving model checkpoints - /// * `use_gpu` - Whether to use GPU acceleration (RTX 3050 Ti) + /// * `use_gpu` - Whether to use GPU acceleration (auto-detected) /// * `num_envs` - Optional number of parallel environments (None/Some(1) = standard, Some(n>1) = vectorized for 2-3x speedup) pub fn new( mut hyperparams: PpoHyperparameters, @@ -226,9 +226,9 @@ impl PpoTrainer { // Dynamic GPU validation: detect hardware and auto-shrink batch size if needed let (device, effective_batch_size) = if use_gpu { - let caps = GpuCapabilities::detect(); + let caps = cached_capabilities(); let max_batch = resolve_batch_size( - &caps, + caps, &memory_profile::estimates::PPO, hyperparams.batch_size, ); @@ -1128,7 +1128,7 @@ mod tests { #[tokio::test] async fn test_ppo_trainer_gpu_batch_limit() { let mut params = create_test_params(); - params.batch_size = 300; // Exceeds GPU limit (230) + params.batch_size = 300; // Exceeds auto-detected GPU limit let trainer = PpoTrainer::new( params,