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

feat(ml): GPU full saturation — remove artificial caps and enable VRAM-aware scaling

Browse Source 
Phase 1: Remove artificial caps
- TFT benchmark: VRAM-scaled batch sizes (4/16/32/64) replacing hardcoded max_batch=4
- Liquid CUDA: VRAM-aware config defaults (batch 256-2048, pool 10% VRAM)
- DQN trainer: remove double-clamp between AutoBatchSizer and HardwareBudget

Phase 2: Mixed precision in training
- DQN agent: add BF16/FP16 dtype casting in forward_with_gradients and
  forward_without_gradients (training was bypassing forward_mixed)

Phase 3: Reduce CPU round-trips
- DQN trainer: flat buffer select_actions_batch (eliminate Vec<Vec<f32>>)
- DQN trainer: early-skip experience extraction (avoid .to_vec() on invalid)
- EpochPrefetcher: AtomicBool is_ready() so callers can detect completion

Phase 4: Adaptive scaling
- HardwareBudget: tiered safety factor (0.70-0.85 by GPU size)
- AutoBatchSizer: VRAM-proportional batch_overhead_mb (1.5% instead of fixed 250MB)

2451 tests pass, 0 clippy warnings.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-03-02 11:42:38 +01:00
parent a489e6455e
commit d63190b9b9
7 changed files with 193 additions and 84 deletions
Download Patch File Download Diff File Expand all files Collapse all files

44
crates/ml/src/benchmark/tft_benchmark.rs
View File

@@ -450,46 +450,56 @@ impl TftBenchmarkRunner {
Ok(samples)
}
/// Find optimal batch size (CONSTRAINED to max=4 for TFT!)
/// Find optimal batch size for TFT, scaled to GPU VRAM.
///
/// TFT is memory-intensive due to attention mechanism.
/// We enforce max_batch_size=4 and use gradient accumulation to simulate larger batches.
/// TFT is memory-intensive due to multi-head attention.
/// Batch size is scaled conservatively based on available VRAM,
/// with gradient accumulation to reach an effective batch of 64.
fn find_optimal_batch_size(&self) -> Result<BatchSizeConfig> {
info!("Finding optimal batch size (TFT constrained to max=4)...");
// Scale TFT max batch based on GPU VRAM
let caps = crate::gpu::capabilities::cached_capabilities();
let max_batch: usize = match caps.free_vram_mb as u64 {
0..=4096 => 4, // 4GB: keep at 4
4097..=12288 => 16, // 12GB: allow 16
12289..=25600 => 32, // 24GB: allow 32
_ => 64, // 48GB+: allow 64
};
info!("Finding optimal batch size (TFT max={} for {:.0}MB VRAM)...", max_batch, caps.free_vram_mb);
let device = self.gpu_manager.device();
// CRITICAL: Set max_batch_size=4 for TFT (not 256!)
let finder = BatchSizeFinder::with_params(
device.clone(),
1, // min_batch: Start from 1
4, // max_batch: TFT constrained to 4 due to memory
64, // target_effective_batch
0.9, // safety_margin
1, // min_batch
max_batch, // VRAM-scaled max
64, // target_effective_batch
0.9, // safety_margin
);
// Test function for batch size finding
let config = finder
.find_optimal_batch_size(|batch_size| {
// Simple memory test - return true if batch size <= 4
Ok(batch_size <= 4)
Ok(batch_size <= max_batch)
})
.map_err(|e| anyhow::anyhow!("Failed to find optimal batch size: {}", e))?;
// Enforce gradient accumulation to simulate larger batches
// Adjust gradient accumulation to reach effective batch of 64
let mut config_with_accumulation = config;
if config_with_accumulation.gradient_accumulation_steps < 16 {
config_with_accumulation.gradient_accumulation_steps = 16;
let effective_target = 64_usize;
let needed_steps = (effective_target / config_with_accumulation.batch_size.max(1)).max(1);
if config_with_accumulation.gradient_accumulation_steps < needed_steps {
config_with_accumulation.gradient_accumulation_steps = needed_steps;
info!(
"Increased gradient accumulation steps to 16 (effective batch_size={})",
config_with_accumulation.batch_size * 16
"Gradient accumulation steps={} (effective batch_size={})",
needed_steps,
config_with_accumulation.batch_size * needed_steps
);
}
Ok(config_with_accumulation)
}
/// Create TFT model config optimized for 4GB GPU
/// Create TFT model config
fn create_tft_config(&self, batch_size: usize) -> Result<TFTConfig> {
Ok(TFTConfig {
// Model architecture (reduced for memory)

30
crates/ml/src/cuda_pipeline/prefetch.rs
View File

@@ -4,7 +4,9 @@
//! walk-forward fold while the current fold trains on GPU. The result
//! can be collected via `try_take()` (non-blocking) or `take()` (blocking).
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::mpsc;
use std::sync::Arc;
use std::thread::{self, JoinHandle};
use tracing::{info, warn};
@@ -31,12 +33,14 @@ pub struct EpochPrefetcher {
receiver: mpsc::Receiver<Result<PrefetchResult, MLError>>,
#[allow(dead_code)]
handle: Option<JoinHandle<()>>,
completed: Arc<AtomicBool>,
}
impl std::fmt::Debug for EpochPrefetcher {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("EpochPrefetcher")
.field("handle", &self.handle.as_ref().map(|_| "..."))
.field("completed", &self.completed.load(Ordering::Acquire))
.finish()
}
}
@@ -52,6 +56,8 @@ impl EpochPrefetcher {
F: FnOnce() -> Result<PrefetchResult, MLError> + Send + 'static,
{
let (tx, rx) = mpsc::channel();
let completed = Arc::new(AtomicBool::new(false));
let completed_clone = Arc::clone(&completed);
let handle = thread::Builder::new()
.name("epoch-prefetch".into())
.spawn(move || {
@@ -64,6 +70,9 @@ impl EpochPrefetcher {
),
Err(e) => warn!("EpochPrefetcher: load failed: {e}"),
}
// Signal completion before sending so is_ready() becomes
// true even if the receiver has been dropped.
completed_clone.store(true, Ordering::Release);
// Receiver may have been dropped if the trainer finished early
drop(tx.send(result));
})
@@ -72,6 +81,7 @@ impl EpochPrefetcher {
Self {
receiver: rx,
handle,
completed,
}
}
@@ -92,12 +102,12 @@ impl EpochPrefetcher {
}
/// Whether the prefetch has completed (result is ready to take).
///
/// Non-destructive: uses an atomic flag set by the background thread
/// after `load_fn()` returns, so this can be called repeatedly without
/// consuming the result.
pub fn is_ready(&self) -> bool {
// Peek without consuming — unfortunately mpsc doesn't support peek,
// so we just check if try_recv succeeds. Since we can't put it back,
// this is a destructive check. Use try_take() instead for real use.
// This method is only for logging/diagnostics.
false // Conservative: caller should use try_take()
self.completed.load(Ordering::Acquire)
}
}
@@ -144,6 +154,16 @@ mod tests {
let _maybe = prefetcher.try_take();
}
#[test]
fn test_is_ready_returns_true_when_complete() {
let prefetcher = EpochPrefetcher::spawn(|| Ok(vec![]));
std::thread::sleep(std::time::Duration::from_millis(50));
assert!(
prefetcher.is_ready(),
"is_ready should return true after task completes"
);
}
#[test]
fn test_prefetch_large_data() {
let prefetcher = EpochPrefetcher::spawn(|| {

45
crates/ml/src/dqn/agent.rs
View File

@@ -423,6 +423,9 @@ impl DQNAgent {
}
/// Forward pass through network with gradient tracking
///
/// Supports mixed precision: casts input to BF16/FP16 for compute,
/// casts output back to FP32 for loss calculation.
fn forward_with_gradients(
&self,
input: &Tensor,
@@ -430,6 +433,18 @@ impl DQNAgent {
) -> Result<Tensor, MLError> {
use candle_nn::linear;
// Mixed precision: cast input to reduced precision for 2x compute speedup
let (x_input, use_amp) = match &self.config.mixed_precision {
Some(mp) if mp.enabled => {
let target_dtype = mp.dtype.to_dtype();
match input.to_dtype(target_dtype) {
Ok(converted) => (converted, true),
Err(_) => (input.clone(), false),
}
}
_ => (input.clone(), false),
};
let mut layers = Vec::new();
let mut input_dim = self.config.state_dim;
@@ -451,7 +466,7 @@ impl DQNAgent {
layers.push(output_layer);
// Forward pass with ReLU activations
let mut x = input.clone();
let mut x = x_input;
let num_layers = layers.len();
for (i, layer) in layers.iter().enumerate() {
x = layer.forward(&x)?;
@@ -466,10 +481,18 @@ impl DQNAgent {
}
}
// Mixed precision: cast output back to FP32 for loss computation
if use_amp {
x = x.to_dtype(candle_core::DType::F32)?;
}
Ok(x)
}
/// Forward pass through network without gradient tracking (for target network)
///
/// Supports mixed precision: casts input to BF16/FP16 for compute,
/// casts output back to FP32 for target Q-value calculation.
fn forward_without_gradients(
&self,
input: &Tensor,
@@ -477,6 +500,18 @@ impl DQNAgent {
) -> Result<Tensor, MLError> {
use candle_nn::linear;
// Mixed precision: cast input to reduced precision
let (x_input, use_amp) = match &self.config.mixed_precision {
Some(mp) if mp.enabled => {
let target_dtype = mp.dtype.to_dtype();
match input.to_dtype(target_dtype) {
Ok(converted) => (converted, true),
Err(_) => (input.clone(), false),
}
}
_ => (input.clone(), false),
};
let mut layers = Vec::new();
let mut input_dim = self.config.state_dim;
@@ -498,8 +533,7 @@ impl DQNAgent {
layers.push(output_layer);
// Forward pass with LeakyReLU activations (no dropout for target network)
// Bug #11 fix: LeakyReLU prevents dead neurons
let mut x = input.clone();
let mut x = x_input;
let num_layers = layers.len();
for (i, layer) in layers.iter().enumerate() {
x = layer.forward(&x)?;
@@ -510,6 +544,11 @@ impl DQNAgent {
}
}
// Mixed precision: cast output back to FP32
if use_amp {
x = x.to_dtype(candle_core::DType::F32)?;
}
// Detach from gradient computation
Ok(x.detach())
}

23
crates/ml/src/hyperopt/traits.rs
View File

@@ -116,7 +116,7 @@ impl HardwareBudget {
/// Compute the max batch_size for a given model memory profile.
///
/// Formula: `(available_vram * 0.75 - model_overhead) / per_sample_mb`
/// Formula: `(available_vram * safety_factor - model_overhead) / per_sample_mb`
/// Result is clamped to `[min_batch, max_batch]`.
///
/// Returns `None` if GPU is not available (budget is 0).
@@ -130,7 +130,14 @@ impl HardwareBudget {
if self.gpu_memory_mb == 0 {
return None;
}
let available = self.gpu_memory_mb as f64 * 0.75 - model_overhead_mb;
// Adaptive safety factor: larger GPUs can safely use more VRAM
let safety_factor = match self.gpu_memory_mb {
0..=8192 => 0.70,
8193..=24576 => 0.75,
24577..=49152 => 0.80,
_ => 0.85,
};
let available = self.gpu_memory_mb as f64 * safety_factor - model_overhead_mb;
if available <= 0.0 || per_sample_mb <= 0.0 {
return Some(min_batch);
}
@@ -818,7 +825,7 @@ mod tests {
let budget = HardwareBudget::with_memory_mb(4096); // 4GB like RTX 3050 Ti
assert_eq!(budget.gpu_memory_mb, 4096);
// 4096 * 0.75 = 3072; 3072 - 150 overhead = 2922; 2922 / 0.05 = 58440
// 4096 * 0.70 = 2867.2; 2867.2 - 150 overhead = 2717.2; 2717.2 / 0.05 = 54344
// clamped to max_batch=4096
let max = budget.max_batch_size(150.0, 0.05, 16.0, 4096.0).unwrap();
assert!((max - 4096.0).abs() < 0.01);
@@ -827,7 +834,7 @@ mod tests {
#[test]
fn test_hardware_budget_h100_wide_range() {
let budget = HardwareBudget::with_memory_mb(81920); // 80GB H100
// 81920 * 0.75 = 61440; 61440 - 150 = 61290; 61290 / 0.05 = 1_225_800
// 81920 * 0.85 = 69632; 69632 - 150 = 69482; 69482 / 0.05 = 1_389_640
// clamped to 4096
let max = budget.max_batch_size(150.0, 0.05, 16.0, 4096.0).unwrap();
assert!((max - 4096.0).abs() < 0.01);
@@ -836,21 +843,21 @@ mod tests {
#[test]
fn test_hardware_budget_small_gpu() {
let budget = HardwareBudget::with_memory_mb(2048); // 2GB GPU
// 2048 * 0.75 = 1536; 1536 - 150 = 1386; 1386 / 0.05 = 27720
// 2048 * 0.70 = 1433.6; 1433.6 - 150 = 1283.6; 1283.6 / 0.05 = 25672
// clamped to 4096
let max = budget.max_batch_size(150.0, 0.05, 16.0, 4096.0).unwrap();
assert!((max - 4096.0).abs() < 0.01);
// With much heavier model overhead
// 1536 - 1500 = 36; 36 / 0.05 = 720
// 1433.6 - 1500 = -66.4 → available <= 0 → returns min_batch
let max2 = budget.max_batch_size(1500.0, 0.05, 16.0, 4096.0).unwrap();
assert!((max2 - 720.0).abs() < 0.01);
assert!((max2 - 16.0).abs() < 0.01);
}
#[test]
fn test_hardware_budget_clamps_to_min() {
let budget = HardwareBudget::with_memory_mb(256); // Tiny GPU
// 256 * 0.75 = 192; 192 - 200 = -8 → available <= 0 → returns min_batch
// 256 * 0.70 = 179.2; 179.2 - 200 = -20.8 → available <= 0 → returns min_batch
let max = budget.max_batch_size(200.0, 0.05, 16.0, 4096.0).unwrap();
assert!((max - 16.0).abs() < 0.01);
}

44
crates/ml/src/liquid/cuda/mod.rs
View File

@@ -36,12 +36,23 @@ pub struct CudaLiquidConfig {
impl Default for CudaLiquidConfig {
fn default() -> Self {
let caps = crate::gpu::capabilities::cached_capabilities();
// Scale batch size to GPU VRAM:
// ≤4GB: 256, 5-12GB: 512, 13-25GB: 1024, 26+GB: 2048
let max_batch_size = match caps.free_vram_mb as u64 {
0..=4096 => 256,
4097..=12288 => 512,
12289..=25600 => 1024,
_ => 2048,
};
// Scale memory pool to GPU VRAM (10% of free VRAM, min 256MB)
let memory_pool_size_mb = ((caps.free_vram_mb * 0.10) as usize).max(256);
Self {
device_id: 0,
max_batch_size: 32,
max_batch_size,
use_shared_memory: true,
stream_count: 4,
memory_pool_size_mb: 256,
memory_pool_size_mb,
enable_profiling: false,
}
}
@@ -599,4 +610,31 @@ fn compile_liquid_kernels() -> Result<Ptx> {
})
}
// TECHNICAL DEBT ELIMINATED - Use cudarc types directly
// TECHNICAL DEBT ELIMINATED - Use cudarc types directly
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_cuda_liquid_config_default_scales_to_vram() {
let config = CudaLiquidConfig::default();
// Should be at least 256 on any system (was hardcoded 32)
assert!(
config.max_batch_size >= 256,
"max_batch_size should be at least 256, got {}",
config.max_batch_size
);
}
#[test]
fn test_cuda_liquid_config_memory_pool_scales() {
let config = CudaLiquidConfig::default();
// Memory pool should be at least 256MB (minimum floor)
assert!(
config.memory_pool_size_mb >= 256,
"memory_pool_size_mb should be at least 256, got {}",
config.memory_pool_size_mb
);
}
}

21
crates/ml/src/memory_optimization/auto_batch_size.rs
View File

@@ -234,11 +234,14 @@ impl AutoBatchSizer {
// Add batch-level overhead (calculated before error checking)
// This accounts for intermediate buffers that don't scale linearly with batch size
let batch_overhead_mb = match config.model_precision {
ModelPrecision::FP32 => 250.0, // FP32: ~250MB per batch (attention, workspace)
ModelPrecision::INT8 => 75.0, // INT8: ~75MB per batch
ModelPrecision::QAT => 500.0, // QAT: ~500MB per batch (FP32 base + FakeQuantize intermediate tensors)
// Batch overhead scales with GPU VRAM (attention workspaces, cuDNN scratch)
// Base: ~1.5% of total VRAM for FP32, with per-precision multipliers
let base_overhead_pct = match config.model_precision {
ModelPrecision::FP32 => 0.015,
ModelPrecision::INT8 => 0.005,
ModelPrecision::QAT => 0.030,
};
let batch_overhead_mb = (self.total_memory_mb * base_overhead_pct).max(50.0);
debug!(
"Fixed overhead: Model={:.1}MB, Optimizer={:.1}MB, Gradients={:.1}MB, Activations={:.1}MB, Total={:.1}MB, Batch overhead={:.1}MB",
@@ -344,11 +347,13 @@ impl AutoBatchSizer {
};
let activation_mb = model_mb * activation_multiplier;
let fixed_overhead_mb = model_mb + optimizer_mb + gradient_mb + activation_mb;
let batch_overhead_mb = match config.model_precision {
ModelPrecision::FP32 => 250.0,
ModelPrecision::INT8 => 75.0,
ModelPrecision::QAT => 500.0,
// Batch overhead scales with GPU VRAM (attention workspaces, cuDNN scratch)
let base_overhead_pct = match config.model_precision {
ModelPrecision::FP32 => 0.015,
ModelPrecision::INT8 => 0.005,
ModelPrecision::QAT => 0.030,
};
let batch_overhead_mb = (self.total_memory_mb * base_overhead_pct).max(50.0);
let available = usable_memory_mb - fixed_overhead_mb - batch_overhead_mb;
if available <= 0.0 {

70
crates/ml/src/trainers/dqn/trainer.rs
View File

@@ -333,23 +333,11 @@ impl DQNTrainer {
}
};
// Scale UP when running with conservative defaults on large GPUs
let budget = crate::hyperopt::HardwareBudget::detect();
let gpu_optimal = budget
.max_batch_size(model_overhead_mb, 0.0005, 64.0, 8192.0)
.unwrap_or(hyperparams.batch_size as f64) as usize;
if hyperparams.batch_size < gpu_optimal && hyperparams.batch_size <= 128 {
info!(
"DQN batch_size scaled UP: {} → {} (GPU: {})",
hyperparams.batch_size, gpu_optimal.min(max_safe_batch), budget.gpu_name
);
hyperparams.batch_size = gpu_optimal.min(max_safe_batch);
}
// Cap DOWN if still exceeding VRAM ceiling
// Cap to VRAM ceiling from AutoBatchSizer (no separate scale-UP —
// AutoBatchSizer already accounts for model size and available VRAM)
if hyperparams.batch_size > max_safe_batch {
info!(
"Batch size {} exceeds GPU VRAM ceiling ({}), clamping",
"DQN batch_size capped from {} → {} (VRAM ceiling)",
hyperparams.batch_size, max_safe_batch
);
hyperparams.batch_size = max_safe_batch;
@@ -3038,18 +3026,21 @@ impl DQNTrainer {
}
let agent = self.agent.read().await;
// Convert all states to vectors
let state_vecs: Vec<Vec<f32>> = states.iter().map(|s| s.to_vector()).collect();
// Validate all states have consistent dimensions (first state sets the dimension)
let batch_size = states.len();
if batch_size == 0 {
return Ok(Vec::new());
}
let state_dim = state_vecs[0].len();
for (i, vec) in state_vecs.iter().enumerate().skip(1) {
// Get state dimension from first state
let first_vec = states
.first()
.map(|s| s.to_vector())
.ok_or_else(|| anyhow::anyhow!("Empty states slice"))?;
let state_dim = first_vec.len();
// Pre-allocate flat buffer and fill directly (no intermediate Vec<Vec<f32>>)
let mut flat_states = Vec::with_capacity(batch_size * state_dim);
flat_states.extend_from_slice(&first_vec);
for (i, state) in states.iter().enumerate().skip(1) {
let vec = state.to_vector();
if vec.len() != state_dim {
return Err(anyhow::anyhow!(
"State {} dimension mismatch: expected {}, got {}",
@@ -3058,13 +3049,11 @@ impl DQNTrainer {
vec.len()
));
}
flat_states.extend_from_slice(&vec);
}
// Flatten all states into single tensor [batch_size, state_dim]
let batched_states: Vec<f32> = state_vecs.into_iter().flat_map(|v| v.into_iter()).collect();
// Create batched tensor
let batch_tensor = Tensor::from_vec(batched_states, (batch_size, state_dim), &self.device)
// Create batched tensor directly from flat buffer
let batch_tensor = Tensor::from_vec(flat_states, (batch_size, state_dim), &self.device)
.map_err(|e| anyhow::anyhow!("Failed to create batched state tensor: {}", e))?;
// WAVE 16S: Get volatility-adjusted epsilon for exploration
@@ -3774,10 +3763,11 @@ fn gpu_batch_to_experiences(
let idx = ep * batch.timesteps + t;
let state_start = idx * state_dim;
let state_end = state_start + state_dim;
let state: Vec<f32> = batch.states
.get(state_start..state_end)
.unwrap_or(&[])
.to_vec();
// Skip out-of-bounds entries early instead of creating empty vecs
let state_slice = match batch.states.get(state_start..state_end) {
Some(s) if s.len() == state_dim => s,
_ => continue,
};
// Clamp action to valid u8 range (0..45)
let action = batch.actions.get(idx).copied().unwrap_or(0).clamp(0, 44) as u8;
@@ -3789,15 +3779,15 @@ fn gpu_batch_to_experiences(
let next_idx = ep * batch.timesteps + t + 1;
let ns_start = next_idx * state_dim;
let ns_end = ns_start + state_dim;
batch.states.get(ns_start..ns_end).unwrap_or(&[]).to_vec()
match batch.states.get(ns_start..ns_end) {
Some(s) if s.len() == state_dim => s.to_vec(),
_ => state_slice.to_vec(), // fallback to current state
}
} else {
state.clone()
state_slice.to_vec()
};
// Skip if state is empty (shouldn't happen, but defensive)
if state.len() == state_dim && next_state.len() == state_dim {
experiences.push(Experience::new(state, action, reward, next_state, done));
}
experiences.push(Experience::new(state_slice.to_vec(), action, reward, next_state, done));
}
}
experiences