8 bite-sized tasks with TDD approach, exact file paths, and complete code. Covers: DeviceConfig relocation, GpuCapabilities detection, ModelMemoryEstimate per-model profiles, and wiring into PPO trainer and hyperopt campaign. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
977 lines
30 KiB
Markdown
977 lines
30 KiB
Markdown
# 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`:
|
||
|
||
```rust
|
||
//! 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):
|
||
|
||
```rust
|
||
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**
|
||
|
||
```bash
|
||
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**
|
||
|
||
```rust
|
||
//! 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**
|
||
|
||
```bash
|
||
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**
|
||
|
||
```rust
|
||
//! 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**
|
||
|
||
```bash
|
||
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):
|
||
|
||
```rust
|
||
/// 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**
|
||
|
||
```bash
|
||
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):
|
||
```rust
|
||
use crate::gpu::DeviceConfig;
|
||
```
|
||
|
||
Replace line 40:
|
||
```rust
|
||
// 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):
|
||
```rust
|
||
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:
|
||
```rust
|
||
use crate::gpu::DeviceConfig;
|
||
```
|
||
|
||
Replace line 42:
|
||
```rust
|
||
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:
|
||
```rust
|
||
use crate::gpu::DeviceConfig;
|
||
```
|
||
|
||
Replace line 59:
|
||
```rust
|
||
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:
|
||
```rust
|
||
use crate::gpu::DeviceConfig;
|
||
```
|
||
|
||
Replace line 56:
|
||
```rust
|
||
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**
|
||
|
||
```bash
|
||
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:
|
||
```rust
|
||
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):
|
||
```rust
|
||
// 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**
|
||
|
||
```bash
|
||
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:
|
||
```rust
|
||
use crate::gpu::capabilities::GpuCapabilities;
|
||
use crate::gpu::memory_profile::{self, resolve_batch_size};
|
||
```
|
||
|
||
Replace the `dqn_default()` method (lines 34-45):
|
||
```rust
|
||
/// 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):
|
||
```rust
|
||
/// 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):
|
||
```rust
|
||
#[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**
|
||
|
||
```bash
|
||
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:
|
||
```rust
|
||
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**
|
||
|
||
```bash
|
||
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.
|