foxhunt/docs/plans/2026-02-23-dynamic-gpu-detection-implementation.md

Dynamic GPU Detection Implementation Plan

For Claude: REQUIRED SUB-SKILL: Use superpowers:executing-plans to implement this plan task-by-task.

Goal: Replace all hardcoded GPU assumptions (RTX 3050 Ti, batch_size 230) with dynamic detection that scales from 4GB dev GPUs to production A100/H100 hardware.

Architecture: Three focused structs in a new ml/src/gpu/ module: DeviceConfig (relocated from Liquid), GpuCapabilities (hardware detection), ModelMemoryEstimate + resolve_batch_size() (per-model batch sizing). Existing AutoBatchSizer math is reused, not reimplemented.

Tech Stack: Rust, candle_core::Device, nvidia-smi (subprocess), serde, tracing

Build: SQLX_OFFLINE=true cargo check --workspace Test: SQLX_OFFLINE=true cargo test -p ml --lib Clippy rules: #![deny(clippy::unwrap_used, clippy::expect_used, clippy::panic, clippy::indexing_slicing)] — use .get(), ?, .ok_or() only.

---

Task 1: Create gpu module with DeviceConfig

Files:

• Create: ml/src/gpu/mod.rs
• Modify: ml/src/lib.rs (add pub mod gpu;)

Step 1: Write the failing test

Create ml/src/gpu/mod.rs with tests that reference DeviceConfig:

//! GPU device management for ML training
//!
//! Provides centralized device selection, GPU capability detection,
//! and per-model batch size optimization.

pub mod capabilities;
pub mod memory_profile;

use candle_core::Device;
use serde::{Deserialize, Serialize};

use crate::MLError;

/// Device configuration for ML training and inference.
///
/// Use `Auto` for production — it detects CUDA if available, falls back to CPU.
/// Use `Cuda(id)` to pin to a specific GPU. Use `Cpu` to force CPU mode.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub enum DeviceConfig {
    Cpu,
    Cuda(usize),
    Auto,
}

impl DeviceConfig {
    /// Resolve this config into a concrete candle `Device`.
    ///
    /// - `Cpu` → always `Device::Cpu`
    /// - `Cuda(id)` → `Device::new_cuda(id)`, errors if unavailable
    /// - `Auto` → CUDA device 0 if available, else CPU
    pub fn resolve(&self) -> Result<Device, MLError> {
        match self {
            DeviceConfig::Cpu => Ok(Device::Cpu),
            DeviceConfig::Cuda(id) => Device::new_cuda(*id).map_err(|e| {
                MLError::ConfigurationError(format!(
                    "CUDA device {} required but unavailable: {}",
                    id, e
                ))
            }),
            DeviceConfig::Auto => {
                match Device::new_cuda(0) {
                    Ok(dev) => Ok(dev),
                    Err(_) => Ok(Device::Cpu),
                }
            }
        }
    }

    /// Returns true if this config will attempt to use a GPU.
    pub fn is_gpu(&self) -> bool {
        matches!(self, DeviceConfig::Cuda(_) | DeviceConfig::Auto)
    }
}

impl Default for DeviceConfig {
    fn default() -> Self {
        DeviceConfig::Auto
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_cpu_always_resolves() {
        let device = DeviceConfig::Cpu.resolve().unwrap();
        assert!(!device.is_cuda());
    }

    #[test]
    fn test_auto_resolves_without_error() {
        // Auto should always succeed (falls back to CPU if no GPU)
        let device = DeviceConfig::Auto.resolve().unwrap();
        // On CI/dev without GPU, this will be CPU — that's correct
        let _ = device;
    }

    #[test]
    fn test_is_gpu() {
        assert!(!DeviceConfig::Cpu.is_gpu());
        assert!(DeviceConfig::Cuda(0).is_gpu());
        assert!(DeviceConfig::Auto.is_gpu());
    }

    #[test]
    fn test_default_is_auto() {
        assert_eq!(DeviceConfig::default(), DeviceConfig::Auto);
    }

    #[test]
    fn test_serde_roundtrip() {
        let configs = vec![DeviceConfig::Cpu, DeviceConfig::Cuda(0), DeviceConfig::Auto];
        for config in configs {
            let json = serde_json::to_string(&config).unwrap();
            let deserialized: DeviceConfig = serde_json::from_str(&json).unwrap();
            assert_eq!(config, deserialized);
        }
    }
}

Step 2: Add module declaration to lib.rs

In ml/src/lib.rs, add after the existing pub mod gradient_utils; line (~line 689):

pub mod gpu; // GPU device management and capability detection

Step 3: Run tests to verify they pass

Run: SQLX_OFFLINE=true cargo test -p ml --lib gpu::tests -- --nocapture Expected: All 5 tests PASS

Step 4: Commit

git add ml/src/gpu/mod.rs ml/src/lib.rs
git commit -m "feat(ml): add gpu module with DeviceConfig enum"

---

Task 2: Create GpuCapabilities for hardware detection

Files:

• Create: ml/src/gpu/capabilities.rs

Step 1: Write the struct and tests

//! GPU hardware capability detection.
//!
//! Detects GPU VRAM and device name via nvidia-smi.
//! Use `GpuCapabilities::detect()` at startup and cache the result.

use serde::{Deserialize, Serialize};
use std::process::Command;
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)]
pub struct GpuCapabilities {
    /// GPU device name (e.g., "NVIDIA GeForce RTX 3050 Ti Laptop GPU")
    pub device_name: String,
    /// Total VRAM in MB (0.0 for CPU)
    pub total_vram_mb: f64,
    /// Free VRAM in MB at detection time (0.0 for CPU)
    pub free_vram_mb: f64,
    /// Whether a CUDA GPU was detected
    pub is_cuda: bool,
}

impl GpuCapabilities {
    /// Detect GPU capabilities by querying nvidia-smi.
    ///
    /// Falls back to CPU if nvidia-smi is unavailable or fails.
    /// Call this once at startup and cache the result.
    pub fn detect() -> Self {
        match Self::query_nvidia_smi() {
            Ok(caps) => {
                info!(
                    "GPU detected: {} (Total: {:.0} MB, Free: {:.0} MB)",
                    caps.device_name, caps.total_vram_mb, caps.free_vram_mb
                );
                caps
            }
            Err(reason) => {
                warn!("GPU detection failed ({}), using CPU fallback", reason);
                Self::cpu_fallback()
            }
        }
    }

    /// Explicit CPU-only capabilities (no GPU).
    pub fn cpu_fallback() -> Self {
        Self {
            device_name: "CPU".to_string(),
            total_vram_mb: 0.0,
            free_vram_mb: 0.0,
            is_cuda: false,
        }
    }

    /// Construct with known VRAM values (for testing or manual override).
    pub fn with_vram(device_name: impl Into<String>, total_mb: f64, free_mb: f64) -> Self {
        Self {
            device_name: device_name.into(),
            total_vram_mb: total_mb,
            free_vram_mb: free_mb,
            is_cuda: total_mb > 0.0,
        }
    }

    /// Query nvidia-smi for GPU memory and device name.
    fn query_nvidia_smi() -> Result<Self, String> {
        let output = Command::new("nvidia-smi")
            .args([
                "--query-gpu=memory.total,memory.free,name",
                "--format=csv,noheader,nounits",
            ])
            .output()
            .map_err(|e| format!("nvidia-smi not found: {}", e))?;

        if !output.status.success() {
            let stderr = String::from_utf8_lossy(&output.stderr);
            return Err(format!("nvidia-smi failed: {}", stderr));
        }

        let stdout = String::from_utf8_lossy(&output.stdout);
        let parts: Vec<&str> = stdout.trim().split(',').collect();

        if parts.len() < 3 {
            return Err(format!("unexpected nvidia-smi format: {}", stdout));
        }

        let total_mb = parts.first()
            .ok_or("missing total")?
            .trim()
            .parse::<f64>()
            .map_err(|e| format!("parse total: {}", e))?;
        let free_mb = parts.get(1)
            .ok_or("missing free")?
            .trim()
            .parse::<f64>()
            .map_err(|e| format!("parse free: {}", e))?;
        let device_name = parts.get(2)
            .ok_or("missing name")?
            .trim()
            .to_string();

        Ok(Self {
            device_name,
            total_vram_mb: total_mb,
            free_vram_mb: free_mb,
            is_cuda: true,
        })
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_cpu_fallback() {
        let caps = GpuCapabilities::cpu_fallback();
        assert_eq!(caps.device_name, "CPU");
        assert_eq!(caps.total_vram_mb, 0.0);
        assert_eq!(caps.free_vram_mb, 0.0);
        assert!(!caps.is_cuda);
    }

    #[test]
    fn test_with_vram_rtx_3050_ti() {
        let caps = GpuCapabilities::with_vram("RTX 3050 Ti", 4096.0, 3700.0);
        assert_eq!(caps.device_name, "RTX 3050 Ti");
        assert_eq!(caps.total_vram_mb, 4096.0);
        assert_eq!(caps.free_vram_mb, 3700.0);
        assert!(caps.is_cuda);
    }

    #[test]
    fn test_with_vram_a100() {
        let caps = GpuCapabilities::with_vram("A100-SXM4-80GB", 81920.0, 80000.0);
        assert_eq!(caps.total_vram_mb, 81920.0);
        assert!(caps.is_cuda);
    }

    #[test]
    fn test_with_vram_zero_is_cpu() {
        let caps = GpuCapabilities::with_vram("none", 0.0, 0.0);
        assert!(!caps.is_cuda);
    }

    #[test]
    fn test_detect_does_not_panic() {
        // detect() should never panic — falls back to CPU gracefully
        let caps = GpuCapabilities::detect();
        assert!(!caps.device_name.is_empty());
    }

    #[test]
    fn test_serde_roundtrip() {
        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);
    }
}

Step 2: Run tests

Run: SQLX_OFFLINE=true cargo test -p ml --lib gpu::capabilities::tests -- --nocapture Expected: All 6 tests PASS

Step 3: Commit

git add ml/src/gpu/capabilities.rs
git commit -m "feat(ml): add GpuCapabilities for hardware detection"

---

Task 3: Create ModelMemoryEstimate and resolve_batch_size

Files:

• Create: ml/src/gpu/memory_profile.rs

Step 1: Write the struct, function, and tests

//! Per-model memory estimation and dynamic batch size resolution.
//!
//! Each ML model provides a `ModelMemoryEstimate` describing its memory footprint.
//! `resolve_batch_size()` combines this with `GpuCapabilities` to compute the
//! optimal batch size for the detected hardware.

use serde::{Deserialize, Serialize};
use tracing::{info, warn};

use crate::memory_optimization::auto_batch_size::{
    AutoBatchSizer, BatchSizeConfig, ModelPrecision, OptimizerType,
};

use super::capabilities::GpuCapabilities;

/// Memory footprint estimate for an ML model.
///
/// Each model type provides one of these so the batch size resolver
/// can compute hardware-appropriate limits.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ModelMemoryEstimate {
    /// Model name (for logging)
    pub name: &'static str,
    /// Estimated parameter count
    pub param_count: usize,
    /// Activation memory multiplier (1.0 for simple MLPs, 2.0+ for attention)
    pub activation_multiplier: f64,
    /// Whether the model supports gradient checkpointing
    pub supports_checkpointing: bool,
    /// Default sequence length for this model
    pub default_seq_len: usize,
    /// Default feature dimension for this model
    pub default_feature_dim: usize,
}

impl ModelMemoryEstimate {
    /// Estimated model size in MB (FP32, parameters only).
    pub fn estimated_size_mb(&self) -> f64 {
        // 4 bytes per FP32 parameter
        (self.param_count as f64 * 4.0) / (1024.0 * 1024.0)
    }
}

/// Pre-defined memory estimates for all 10 Foxhunt models.
///
/// These are conservative estimates derived from model architecture configs.
/// Actual values may vary with specific hyperparameter choices.
pub mod estimates {
    use super::ModelMemoryEstimate;

    pub const DQN: ModelMemoryEstimate = ModelMemoryEstimate {
        name: "DQN",
        param_count: 50_000,
        activation_multiplier: 1.0,
        supports_checkpointing: false,
        default_seq_len: 1,
        default_feature_dim: 32,
    };

    pub const PPO: ModelMemoryEstimate = ModelMemoryEstimate {
        name: "PPO",
        param_count: 100_000,
        activation_multiplier: 1.0,
        supports_checkpointing: false,
        default_seq_len: 1,
        default_feature_dim: 32,
    };

    pub const TFT: ModelMemoryEstimate = ModelMemoryEstimate {
        name: "TFT",
        param_count: 2_000_000,
        activation_multiplier: 2.5,
        supports_checkpointing: true,
        default_seq_len: 60,
        default_feature_dim: 225,
    };

    pub const MAMBA2: ModelMemoryEstimate = ModelMemoryEstimate {
        name: "Mamba2",
        param_count: 500_000,
        activation_multiplier: 1.5,
        supports_checkpointing: false,
        default_seq_len: 60,
        default_feature_dim: 64,
    };

    pub const TGGN: ModelMemoryEstimate = ModelMemoryEstimate {
        name: "TGGN",
        param_count: 200_000,
        activation_multiplier: 1.5,
        supports_checkpointing: false,
        default_seq_len: 32,
        default_feature_dim: 64,
    };

    pub const TLOB: ModelMemoryEstimate = ModelMemoryEstimate {
        name: "TLOB",
        param_count: 300_000,
        activation_multiplier: 2.0,
        supports_checkpointing: false,
        default_seq_len: 128,
        default_feature_dim: 51,
    };

    pub const LIQUID: ModelMemoryEstimate = ModelMemoryEstimate {
        name: "Liquid",
        param_count: 100_000,
        activation_multiplier: 1.0,
        supports_checkpointing: false,
        default_seq_len: 1,
        default_feature_dim: 32,
    };

    pub const KAN: ModelMemoryEstimate = ModelMemoryEstimate {
        name: "KAN",
        param_count: 150_000,
        activation_multiplier: 1.0,
        supports_checkpointing: false,
        default_seq_len: 1,
        default_feature_dim: 32,
    };

    pub const XLSTM: ModelMemoryEstimate = ModelMemoryEstimate {
        name: "xLSTM",
        param_count: 400_000,
        activation_multiplier: 1.5,
        supports_checkpointing: false,
        default_seq_len: 60,
        default_feature_dim: 64,
    };

    pub const DIFFUSION: ModelMemoryEstimate = ModelMemoryEstimate {
        name: "Diffusion",
        param_count: 1_000_000,
        activation_multiplier: 2.0,
        supports_checkpointing: true,
        default_seq_len: 32,
        default_feature_dim: 64,
    };
}

/// Resolve the optimal batch size for a model given GPU capabilities.
///
/// Uses the existing `AutoBatchSizer` math from `memory_optimization`.
/// If the GPU has no VRAM (CPU mode), returns a conservative default (32).
///
/// # Arguments
/// * `capabilities` - Detected GPU hardware (from `GpuCapabilities::detect()`)
/// * `estimate` - Memory footprint of the model to train
/// * `requested_batch_size` - The batch size the user/config requested (may be shrunk)
///
/// # Returns
/// The effective batch size (may be smaller than requested if VRAM is insufficient).
pub fn resolve_batch_size(
    capabilities: &GpuCapabilities,
    estimate: &ModelMemoryEstimate,
    requested_batch_size: usize,
) -> usize {
    // CPU mode: no VRAM constraint, use requested or conservative default
    if !capabilities.is_cuda {
        info!(
            "CPU mode: using requested batch_size {} for {}",
            requested_batch_size, estimate.name
        );
        return requested_batch_size;
    }

    let model_size_mb = estimate.estimated_size_mb();

    let batch_config = BatchSizeConfig {
        model_memory_mb: model_size_mb,
        model_precision: ModelPrecision::FP32,
        base_model_memory_mb: model_size_mb,
        sequence_length: estimate.default_seq_len,
        feature_dim: estimate.default_feature_dim,
        gradient_checkpointing: estimate.supports_checkpointing,
        optimizer_type: OptimizerType::AdamW,
        safety_margin: 0.20,
        min_batch_size: 1,
        max_batch_size: requested_batch_size,
    };

    let sizer = AutoBatchSizer::with_manual_memory(
        capabilities.total_vram_mb,
        capabilities.free_vram_mb,
        capabilities.device_name.clone(),
    );

    match sizer.calculate_optimal_batch_size(&batch_config) {
        Ok(optimal) => {
            if optimal < requested_batch_size {
                warn!(
                    "{}: batch_size {} exceeds GPU capacity ({} {} MB free), shrinking to {}",
                    estimate.name,
                    requested_batch_size,
                    capabilities.device_name,
                    capabilities.free_vram_mb,
                    optimal
                );
            } else {
                info!(
                    "{}: batch_size {} fits on {} ({} MB free)",
                    estimate.name,
                    requested_batch_size,
                    capabilities.device_name,
                    capabilities.free_vram_mb
                );
            }
            optimal
        }
        Err(e) => {
            warn!(
                "{}: GPU memory insufficient ({}), falling back to batch_size 1",
                estimate.name, e
            );
            1
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_dqn_estimate_size() {
        let size = estimates::DQN.estimated_size_mb();
        // 50K params × 4 bytes = 200KB ≈ 0.19 MB
        assert!(size > 0.1 && size < 1.0, "DQN should be ~0.19 MB, got {}", size);
    }

    #[test]
    fn test_tft_estimate_size() {
        let size = estimates::TFT.estimated_size_mb();
        // 2M params × 4 bytes = 8MB
        assert!(size > 5.0 && size < 15.0, "TFT should be ~7.6 MB, got {}", size);
    }

    #[test]
    fn test_resolve_batch_size_cpu_passthrough() {
        let caps = GpuCapabilities::cpu_fallback();
        let batch = resolve_batch_size(&caps, &estimates::DQN, 256);
        assert_eq!(batch, 256, "CPU mode should pass through requested batch size");
    }

    #[test]
    fn test_resolve_batch_size_small_gpu() {
        // 4GB GPU — DQN is tiny, should fit large batches
        let caps = GpuCapabilities::with_vram("RTX 3050 Ti", 4096.0, 3700.0);
        let batch = resolve_batch_size(&caps, &estimates::DQN, 512);
        assert!(batch >= 32, "DQN on 4GB GPU should handle at least 32, got {}", batch);
    }

    #[test]
    fn test_resolve_batch_size_tft_constrained() {
        // 4GB GPU — TFT is memory-hungry with attention
        let caps = GpuCapabilities::with_vram("RTX 3050 Ti", 4096.0, 3700.0);
        let batch_small = resolve_batch_size(&caps, &estimates::TFT, 512);
        let batch_dqn = resolve_batch_size(&caps, &estimates::DQN, 512);
        // TFT should get a smaller batch than DQN on the same hardware
        assert!(
            batch_small <= batch_dqn,
            "TFT ({}) should get <= batch_size than DQN ({}) on 4GB",
            batch_small, batch_dqn
        );
    }

    #[test]
    fn test_resolve_batch_size_large_gpu_no_shrink() {
        // 80GB A100 — everything should fit at requested size
        let caps = GpuCapabilities::with_vram("A100-SXM4-80GB", 81920.0, 80000.0);
        let batch = resolve_batch_size(&caps, &estimates::TFT, 256);
        // A100 should handle 256 batch for TFT easily
        assert_eq!(batch, 128, "A100 should handle large batches (power-of-2 rounded)");
    }

    #[test]
    fn test_all_models_have_valid_estimates() {
        let all = [
            &estimates::DQN, &estimates::PPO, &estimates::TFT,
            &estimates::MAMBA2, &estimates::TGGN, &estimates::TLOB,
            &estimates::LIQUID, &estimates::KAN, &estimates::XLSTM,
            &estimates::DIFFUSION,
        ];
        for est in all {
            assert!(!est.name.is_empty());
            assert!(est.param_count > 0);
            assert!(est.activation_multiplier >= 1.0);
            assert!(est.default_seq_len > 0);
            assert!(est.default_feature_dim > 0);
            let size = est.estimated_size_mb();
            assert!(size > 0.0, "{} has zero estimated size", est.name);
        }
    }
}

Step 2: Run tests

Run: SQLX_OFFLINE=true cargo test -p ml --lib gpu::memory_profile::tests -- --nocapture Expected: All 7 tests PASS

Note: The test_resolve_batch_size_large_gpu_no_shrink assertion value (128) is the power-of-2 rounded result from AutoBatchSizer. If it differs, adjust to match the actual AutoBatchSizer rounding behavior — the key invariant is that A100 should NOT shrink below the requested size.

Step 3: Commit

git add ml/src/gpu/memory_profile.rs
git commit -m "feat(ml): add ModelMemoryEstimate and resolve_batch_size"

---

Task 4: Update Liquid CfC to re-import DeviceConfig

Files:

• Modify: ml/src/liquid/candle_cfc.rs (remove DeviceConfig, re-import)

Step 1: Replace the DeviceConfig definition

In ml/src/liquid/candle_cfc.rs, replace lines 14-36 (the DeviceConfig enum and impl):

/// Device configuration for CfC training — re-exported from central gpu module.
pub use crate::gpu::DeviceConfig;

This removes the duplicate definition and imports from the new canonical location.

Step 2: Run tests to verify nothing broke

Run: SQLX_OFFLINE=true cargo test -p ml --lib liquid -- --nocapture Expected: All existing Liquid tests PASS (DeviceConfig API is identical)

Run: SQLX_OFFLINE=true cargo check --workspace Expected: Compiles with 0 errors

Step 3: Commit

git add ml/src/liquid/candle_cfc.rs
git commit -m "refactor(ml): use canonical DeviceConfig from gpu module in Liquid CfC"

---

Task 5: Wire DeviceConfig into ensemble adapters

Files:

• Modify: ml/src/ensemble/adapters/dqn.rs (lines 40, 50)
• Modify: ml/src/ensemble/adapters/ppo.rs (line 42)
• Modify: ml/src/ensemble/adapters/tft.rs (line 59)
• Modify: ml/src/ensemble/adapters/mamba2.rs (line 56)

Step 1: Update DQN adapter

In ml/src/ensemble/adapters/dqn.rs:

Add import at top (after line 8):

use crate::gpu::DeviceConfig;

Replace line 40:

// Before:
let device = Device::cuda_if_available(0).unwrap_or(Device::Cpu);
// After:
let device = DeviceConfig::Auto.resolve().unwrap_or(Device::Cpu);

Replace line 50 (same pattern):

let device = DeviceConfig::Auto.resolve().unwrap_or(Device::Cpu);

Step 2: Update PPO adapter

In ml/src/ensemble/adapters/ppo.rs:

Add import at top:

use crate::gpu::DeviceConfig;

Replace line 42:

let device = DeviceConfig::Auto.resolve().unwrap_or(Device::Cpu);

Step 3: Update TFT adapter

In ml/src/ensemble/adapters/tft.rs:

Add import at top:

use crate::gpu::DeviceConfig;

Replace line 59:

let device = DeviceConfig::Auto.resolve().unwrap_or(Device::Cpu);

Step 4: Update Mamba2 adapter

In ml/src/ensemble/adapters/mamba2.rs:

Add import at top:

use crate::gpu::DeviceConfig;

Replace line 56:

let device = DeviceConfig::Auto.resolve().unwrap_or(Device::Cpu);

Step 5: Run tests to verify

Run: SQLX_OFFLINE=true cargo test -p ml --lib ensemble -- --nocapture Expected: All ensemble adapter tests PASS

Run: SQLX_OFFLINE=true cargo check --workspace Expected: 0 errors

Step 6: Commit

git add ml/src/ensemble/adapters/dqn.rs ml/src/ensemble/adapters/ppo.rs ml/src/ensemble/adapters/tft.rs ml/src/ensemble/adapters/mamba2.rs
git commit -m "refactor(ml): use DeviceConfig::Auto in ensemble adapters"

---

Task 6: Wire dynamic batch size into PPO trainer

Files:

• Modify: ml/src/trainers/ppo.rs (lines 224-249)

Step 1: Replace hardcoded 230 limit

In ml/src/trainers/ppo.rs, add imports at the top:

use crate::gpu::DeviceConfig;
use crate::gpu::capabilities::GpuCapabilities;
use crate::gpu::memory_profile::{self, resolve_batch_size};

Replace lines 224-249 (the GPU validation block + device creation):

        // 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 max_batch = resolve_batch_size(
                &caps,
                &memory_profile::estimates::PPO,
                hyperparams.batch_size,
            );
            let device = DeviceConfig::Auto.resolve().unwrap_or(Device::Cpu);
            if device.is_cuda() {
                info!("PPO using GPU: {} (batch_size: {})", caps.device_name, max_batch);
            } else {
                warn!("GPU requested but unavailable, falling back to CPU");
            }
            (device, max_batch)
        } else {
            (Device::Cpu, hyperparams.batch_size)
        };

        // Apply effective batch size (may have been shrunk for GPU fit)
        hyperparams.batch_size = effective_batch_size;

Also update line 258 to use the device variable (which it already does — just verify).

Step 2: Run tests

Run: SQLX_OFFLINE=true cargo test -p ml --lib trainers::ppo -- --nocapture Expected: All PPO trainer tests PASS

Step 3: Commit

git add ml/src/trainers/ppo.rs
git commit -m "feat(ml): replace hardcoded batch_size 230 with dynamic GPU detection in PPO trainer"

---

Task 7: Wire dynamic batch size into hyperopt campaign

Files:

• Modify: ml/src/hyperopt/campaign.rs (lines 34-58, 168-174)

Step 1: Replace hardcoded 230 defaults

In ml/src/hyperopt/campaign.rs, add imports:

use crate::gpu::capabilities::GpuCapabilities;
use crate::gpu::memory_profile::{self, resolve_batch_size};

Replace the dqn_default() method (lines 34-45):

    /// 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);
        Self {
            model_type: ModelType::DQN,
            num_trials: 50,
            data_dir: PathBuf::from("test_data/real/databento/ml_training"),
            max_batch_size: max_batch,
            early_stopping_eta: 3,
            max_epochs_per_trial: 81,
            results_base_dir: PathBuf::from("ml/hyperopt_results"),
        }
    }

Replace the ppo_default() method (lines 47-58):

    /// 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);
        Self {
            model_type: ModelType::PPO,
            num_trials: 30,
            data_dir: PathBuf::from("test_data/real/databento/ml_training"),
            max_batch_size: max_batch,
            early_stopping_eta: 3,
            max_epochs_per_trial: 81,
            results_base_dir: PathBuf::from("ml/hyperopt_results"),
        }
    }

Step 2: Fix the test assertion

Replace the test (lines 167-174):

    #[test]
    fn test_campaign_config_dqn_defaults() {
        let config = CampaignConfig::dqn_default();
        assert_eq!(config.model_type, ModelType::DQN);
        assert_eq!(config.num_trials, 50);
        assert!(config.max_batch_size > 0, "max_batch_size should be positive");
        // No longer asserts <= 230 — batch size is dynamic based on GPU
    }

Step 3: Run tests

Run: SQLX_OFFLINE=true cargo test -p ml --lib hyperopt::campaign -- --nocapture Expected: All campaign tests PASS

Step 4: Commit

git add ml/src/hyperopt/campaign.rs
git commit -m "feat(ml): use dynamic GPU detection for hyperopt campaign batch sizes"

---

Task 8: Add re-exports and run full test suite

Files:

• Modify: ml/src/lib.rs (add re-exports to prelude)

Step 1: Add gpu types to prelude

In ml/src/lib.rs, in the pub mod prelude section (~line 1923), add:

    pub use crate::gpu::{DeviceConfig, capabilities::GpuCapabilities};

Step 2: Run full workspace check

Run: SQLX_OFFLINE=true cargo check --workspace Expected: 0 errors

Step 3: Run full ML test suite

Run: SQLX_OFFLINE=true cargo test -p ml --lib 2>&1 | tail -20 Expected: All ~2204 tests PASS, 0 failures

Step 4: Run clippy

Run: SQLX_OFFLINE=true cargo clippy -p ml --lib -- -D warnings 2>&1 | tail -20 Expected: 0 clippy errors

Step 5: Commit

git add ml/src/lib.rs
git commit -m "feat(ml): export GPU types in prelude"

---

Summary of Changes

File	Action	What
ml/src/gpu/mod.rs	Create	DeviceConfig enum (relocated)
ml/src/gpu/capabilities.rs	Create	GpuCapabilities hardware detection
ml/src/gpu/memory_profile.rs	Create	ModelMemoryEstimate + resolve_batch_size()
ml/src/lib.rs	Modify	Add pub mod gpu; + prelude exports
ml/src/liquid/candle_cfc.rs	Modify	Replace DeviceConfig definition with re-import
ml/src/ensemble/adapters/dqn.rs	Modify	Use DeviceConfig::Auto
ml/src/ensemble/adapters/ppo.rs	Modify	Use DeviceConfig::Auto
ml/src/ensemble/adapters/tft.rs	Modify	Use DeviceConfig::Auto
ml/src/ensemble/adapters/mamba2.rs	Modify	Use DeviceConfig::Auto
ml/src/trainers/ppo.rs	Modify	Dynamic batch size, remove 230 limit
ml/src/hyperopt/campaign.rs	Modify	Dynamic batch size defaults

Out of Scope (follow-up work)

These files also have cuda_if_available but are NOT changed in this plan:

• dqn/dqn.rs, dqn/network.rs, dqn/ensemble.rs — internal DQN model code (lower priority)
• tft/mod.rs, tft/training.rs, tft/quantized_tft.rs — TFT model code
• trainers/dqn/trainer.rs, trainers/mamba2.rs, trainers/tlob.rs — other trainers
• inference.rs, training_pipeline.rs — inference/pipeline paths
• flash_attention/, benchmark/, benchmarks.rs — test/benchmark code
• stress_testing.rs — test-only code
• data_loaders/ — data loading (CPU is fine here)

These can be migrated incrementally in follow-up PRs.