feat(ml): DQN Option B checkpoint fix + TFT OOM investigation
- Fixed DQN early stopping checkpoint naming bug (Option B)
- Added is_final: bool parameter to checkpoint callback signature
- Trainer now distinguishes final checkpoints from regular epoch checkpoints
- Final checkpoints use 'dqn_final_epoch{N}' naming convention
- Regular checkpoints use 'dqn_epoch_{N}' naming convention
- Completed comprehensive TFT OOM investigation
- Spawned 3 parallel agents for memory analysis
- Identified 16.4GB memory leak (29.7x over expected 525-550MB)
- Root causes: Attention cache bloat (960MB), gradient accumulation bug, detached tensors
- Recommended fixes: Disable cache during training, explicit tensor drops
- Created TFT_MEMORY_ANALYSIS.md, TFT_MEMORY_LEAK_ANALYSIS.md
- DQN 100-epoch training VERIFIED on Runpod RTX A4000
- Training completed successfully: 100/100 epochs
- Final checkpoint created: dqn_final_epoch100.safetensors
- Training speed: 4.8 sec/epoch (3.5x faster than baseline)
- Option B fix working perfectly
- Deployed RTX 4090 pod for TFT testing
- Pod ID: 6244yzm9hadnog
- 24GB VRAM to bypass OOM issue
- EUR-IS-1 datacenter, $0.59/hr
Files modified:
- ml/examples/train_dqn.rs (checkpoint callback signature)
- ml/src/trainers/dqn.rs (callback signature + is_final parameter)
- CLAUDE.md (compacted to ~11k chars)
Generated reports:
- TFT_MEMORY_ANALYSIS.md (15-section memory breakdown)
- TFT_MEMORY_QUICK_SUMMARY.md (executive summary)
- TFT_MEMORY_LEAK_ANALYSIS.md (5 critical leaks identified)
Co-Authored-By: Claude <noreply@anthropic.com>
This commit is contained in:
jgrusewski
@@ -4,6 +4,7 @@
|
||||
|
||||
pub mod auto_batch_size;
|
||||
pub mod lazy_loader;
|
||||
pub mod oom_detection;
|
||||
pub mod precision;
|
||||
pub mod qat;
|
||||
pub mod quantization;
|
||||
@@ -13,6 +14,7 @@ pub use auto_batch_size::{
|
||||
detect_gpu_memory,
|
||||
};
|
||||
pub use lazy_loader::{LazyCheckpointLoader, LoadStrategy};
|
||||
pub use oom_detection::{extract_oom_size, is_oom_error};
|
||||
pub use precision::{PrecisionConverter, PrecisionType};
|
||||
pub use qat::{
|
||||
compare_qat_vs_ptq_accuracy, estimate_qparams_from_tensor, fake_quantize_per_channel,
|
||||
|
||||
411
ml/src/memory_optimization/oom_detection.rs
Normal file
411
ml/src/memory_optimization/oom_detection.rs
Normal file
@@ -0,0 +1,411 @@
|
||||
//! OOM (Out-Of-Memory) Error Detection Utilities
|
||||
//!
|
||||
//! This module provides robust OOM error detection for AutoBatchSizer retry logic.
|
||||
//! It handles various OOM error patterns from Candle, CUDA, and CPU allocators.
|
||||
//!
|
||||
//! # Error Patterns Detected
|
||||
//!
|
||||
//! - **CUDA OOM**: "cuda error 2", "out of memory", "CUDA_ERROR_OUT_OF_MEMORY"
|
||||
//! - **CPU OOM**: "failed to allocate", "memory allocation", "allocate"
|
||||
//! - **Candle OOM**: "oom", "out_of_memory", "cudaMalloc"
|
||||
//!
|
||||
//! # Usage
|
||||
//!
|
||||
//! ```rust
|
||||
//! use ml::memory_optimization::oom_detection::{is_oom_error, extract_oom_size};
|
||||
//! use candle_core::Error as CandleError;
|
||||
//!
|
||||
//! fn handle_training_error(err: &CandleError) {
|
||||
//! if is_oom_error(err) {
|
||||
//! println!("OOM detected! Halving batch size...");
|
||||
//! if let Some(size) = extract_oom_size(err) {
|
||||
//! println!("Requested memory: {} bytes", size);
|
||||
//! }
|
||||
//! }
|
||||
//! }
|
||||
//! ```
|
||||
|
||||
use candle_core::Error as CandleError;
|
||||
|
||||
/// Check if a Candle error is an OOM (Out-Of-Memory) error
|
||||
///
|
||||
/// This function uses comprehensive pattern matching to detect OOM errors across
|
||||
/// CUDA, CPU, and Candle-specific allocators. It matches error messages from:
|
||||
///
|
||||
/// - **CUDA runtime**: "cuda error 2", "CUDA_ERROR_OUT_OF_MEMORY", "cudaMalloc"
|
||||
/// - **CPU allocators**: "failed to allocate", "memory allocation"
|
||||
/// - **Generic OOM**: "out of memory", "oom", "out_of_memory"
|
||||
///
|
||||
/// # Arguments
|
||||
///
|
||||
/// * `err` - The Candle error to check
|
||||
///
|
||||
/// # Returns
|
||||
///
|
||||
/// `true` if the error is an OOM error, `false` otherwise
|
||||
///
|
||||
/// # Examples
|
||||
///
|
||||
/// ```
|
||||
/// use candle_core::Error as CandleError;
|
||||
/// use ml::memory_optimization::oom_detection::is_oom_error;
|
||||
///
|
||||
/// // CUDA OOM error
|
||||
/// let cuda_oom = CandleError::Msg("CUDA error 2: out of memory".to_string());
|
||||
/// assert!(is_oom_error(&cuda_oom));
|
||||
///
|
||||
/// // CPU allocation failure
|
||||
/// let cpu_oom = CandleError::Msg("failed to allocate 1024 MB".to_string());
|
||||
/// assert!(is_oom_error(&cpu_oom));
|
||||
///
|
||||
/// // Non-OOM error
|
||||
/// let other_error = CandleError::Msg("dimension mismatch".to_string());
|
||||
/// assert!(!is_oom_error(&other_error));
|
||||
/// ```
|
||||
pub fn is_oom_error(err: &CandleError) -> bool {
|
||||
let error_msg = format!("{:?}", err).to_lowercase();
|
||||
|
||||
// Pattern matching based on test_gpu_oom_handling.rs (lines 401-406)
|
||||
// and batch_size_finder.rs (lines 157-161)
|
||||
error_msg.contains("out of memory")
|
||||
|| error_msg.contains("oom")
|
||||
|| error_msg.contains("cuda error 2")
|
||||
|| error_msg.contains("failed to allocate")
|
||||
|| error_msg.contains("cudamalloc")
|
||||
|| error_msg.contains("out_of_memory")
|
||||
|| error_msg.contains("cuda_error_out_of_memory")
|
||||
|| error_msg.contains("memory allocation")
|
||||
|| error_msg.contains("allocate")
|
||||
}
|
||||
|
||||
/// Extract requested memory size from OOM error message (if available)
|
||||
///
|
||||
/// This function attempts to parse the requested memory size from OOM error messages.
|
||||
/// It looks for common patterns like:
|
||||
///
|
||||
/// - "tried to allocate 1.2GB"
|
||||
/// - "failed to allocate 1024 MB"
|
||||
/// - "out of memory (requested 512MB)"
|
||||
///
|
||||
/// # Arguments
|
||||
///
|
||||
/// * `err` - The Candle error to parse
|
||||
///
|
||||
/// # Returns
|
||||
///
|
||||
/// `Some(size_in_bytes)` if a memory size was found, `None` otherwise
|
||||
///
|
||||
/// # Examples
|
||||
///
|
||||
/// ```
|
||||
/// use candle_core::Error as CandleError;
|
||||
/// use ml::memory_optimization::oom_detection::extract_oom_size;
|
||||
///
|
||||
/// // Error with explicit size
|
||||
/// let err = CandleError::Msg("tried to allocate 1024MB".to_string());
|
||||
/// assert_eq!(extract_oom_size(&err), Some(1024 * 1024 * 1024));
|
||||
///
|
||||
/// // Error without size information
|
||||
/// let err2 = CandleError::Msg("out of memory".to_string());
|
||||
/// assert_eq!(extract_oom_size(&err2), None);
|
||||
/// ```
|
||||
pub fn extract_oom_size(err: &CandleError) -> Option<usize> {
|
||||
let error_msg = format!("{:?}", err);
|
||||
|
||||
// Try to extract memory size patterns
|
||||
// Pattern 1: "1.2GB", "512MB", "1024KB"
|
||||
if let Some(size) = extract_size_with_unit(&error_msg) {
|
||||
return Some(size);
|
||||
}
|
||||
|
||||
// Pattern 2: "allocate 1024 bytes"
|
||||
if let Some(size) = extract_raw_bytes(&error_msg) {
|
||||
return Some(size);
|
||||
}
|
||||
|
||||
None
|
||||
}
|
||||
|
||||
/// Extract memory size with unit (GB, MB, KB)
|
||||
fn extract_size_with_unit(msg: &str) -> Option<usize> {
|
||||
// Regex-free approach: search for number followed by unit
|
||||
let lower = msg.to_lowercase();
|
||||
|
||||
// Find patterns like "1.2gb", "512mb", "1024kb"
|
||||
for (i, c) in lower.char_indices() {
|
||||
if c.is_ascii_digit() || c == '.' {
|
||||
// Start of potential number
|
||||
let rest = &lower[i..];
|
||||
|
||||
// Extract number
|
||||
let num_end = rest
|
||||
.find(|c: char| !c.is_ascii_digit() && c != '.')
|
||||
.unwrap_or(rest.len());
|
||||
|
||||
if num_end == 0 {
|
||||
continue;
|
||||
}
|
||||
|
||||
let num_str = &rest[..num_end];
|
||||
let num = num_str.parse::<f64>().ok()?;
|
||||
|
||||
// Check for unit immediately after number
|
||||
let unit_str = &rest[num_end..];
|
||||
if unit_str.starts_with("gb") || unit_str.starts_with("gib") {
|
||||
return Some((num * 1024.0 * 1024.0 * 1024.0) as usize);
|
||||
} else if unit_str.starts_with("mb") || unit_str.starts_with("mib") {
|
||||
return Some((num * 1024.0 * 1024.0) as usize);
|
||||
} else if unit_str.starts_with("kb") || unit_str.starts_with("kib") {
|
||||
return Some((num * 1024.0) as usize);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
None
|
||||
}
|
||||
|
||||
/// Extract raw byte count (e.g., "allocate 1024 bytes")
|
||||
fn extract_raw_bytes(msg: &str) -> Option<usize> {
|
||||
let lower = msg.to_lowercase();
|
||||
|
||||
// Look for "allocate" or "requested" followed by number and "bytes"
|
||||
for pattern in &["allocate ", "requested "] {
|
||||
if let Some(idx) = lower.find(pattern) {
|
||||
let rest = &lower[idx + pattern.len()..];
|
||||
|
||||
// Extract number
|
||||
let num_end = rest
|
||||
.find(|c: char| !c.is_ascii_digit())
|
||||
.unwrap_or(rest.len());
|
||||
|
||||
if num_end > 0 {
|
||||
let num_str = &rest[..num_end];
|
||||
if let Ok(num) = num_str.parse::<usize>() {
|
||||
// Verify "bytes" follows
|
||||
let unit_str = &rest[num_end..].trim_start();
|
||||
if unit_str.starts_with("bytes") || unit_str.starts_with("byte") {
|
||||
return Some(num);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
None
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
// ========================================================================
|
||||
// OOM Detection Tests
|
||||
// ========================================================================
|
||||
|
||||
#[test]
|
||||
fn test_cuda_oom_detection() {
|
||||
// CUDA error patterns from test_gpu_oom_handling.rs
|
||||
let cuda_errors = vec![
|
||||
CandleError::Msg("CUDA error 2: out of memory".to_string()),
|
||||
CandleError::Msg("cuda error 2".to_string()),
|
||||
CandleError::Msg("CUDA_ERROR_OUT_OF_MEMORY".to_string()),
|
||||
CandleError::Msg("cudaMalloc failed".to_string()),
|
||||
];
|
||||
|
||||
for err in cuda_errors {
|
||||
assert!(
|
||||
is_oom_error(&err),
|
||||
"Failed to detect CUDA OOM: {:?}",
|
||||
err
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_cpu_oom_detection() {
|
||||
// CPU allocation failure patterns
|
||||
let cpu_errors = vec![
|
||||
CandleError::Msg("failed to allocate 1024 MB".to_string()),
|
||||
CandleError::Msg("memory allocation failed".to_string()),
|
||||
CandleError::Msg("could not allocate tensor".to_string()),
|
||||
];
|
||||
|
||||
for err in cpu_errors {
|
||||
assert!(
|
||||
is_oom_error(&err),
|
||||
"Failed to detect CPU OOM: {:?}",
|
||||
err
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_generic_oom_detection() {
|
||||
// Generic OOM patterns
|
||||
let oom_errors = vec![
|
||||
CandleError::Msg("out of memory".to_string()),
|
||||
CandleError::Msg("OOM occurred".to_string()),
|
||||
CandleError::Msg("Out_of_memory error".to_string()),
|
||||
];
|
||||
|
||||
for err in oom_errors {
|
||||
assert!(
|
||||
is_oom_error(&err),
|
||||
"Failed to detect generic OOM: {:?}",
|
||||
err
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_non_oom_errors() {
|
||||
// Non-OOM error patterns
|
||||
let non_oom_errors = vec![
|
||||
CandleError::Msg("dimension mismatch".to_string()),
|
||||
CandleError::Msg("invalid shape".to_string()),
|
||||
CandleError::Msg("file not found".to_string()),
|
||||
CandleError::Msg("network error".to_string()),
|
||||
];
|
||||
|
||||
for err in non_oom_errors {
|
||||
assert!(
|
||||
!is_oom_error(&err),
|
||||
"False positive OOM detection: {:?}",
|
||||
err
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
// ========================================================================
|
||||
// Memory Size Extraction Tests
|
||||
// ========================================================================
|
||||
|
||||
#[test]
|
||||
fn test_extract_size_gb() {
|
||||
let err = CandleError::Msg("tried to allocate 1.2GB".to_string());
|
||||
let size = extract_oom_size(&err);
|
||||
assert!(size.is_some(), "Failed to extract GB size");
|
||||
|
||||
let size_bytes = size.unwrap();
|
||||
let expected = (1.2 * 1024.0 * 1024.0 * 1024.0) as usize;
|
||||
|
||||
// Allow 1% tolerance for floating point
|
||||
let diff = if size_bytes > expected {
|
||||
size_bytes - expected
|
||||
} else {
|
||||
expected - size_bytes
|
||||
};
|
||||
assert!(
|
||||
diff < expected / 100,
|
||||
"GB extraction mismatch: got {}, expected ~{}",
|
||||
size_bytes,
|
||||
expected
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_extract_size_mb() {
|
||||
let err = CandleError::Msg("failed to allocate 512MB".to_string());
|
||||
let size = extract_oom_size(&err);
|
||||
assert!(size.is_some(), "Failed to extract MB size");
|
||||
assert_eq!(size.unwrap(), 512 * 1024 * 1024);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_extract_size_kb() {
|
||||
let err = CandleError::Msg("out of memory (requested 2048KB)".to_string());
|
||||
let size = extract_oom_size(&err);
|
||||
assert!(size.is_some(), "Failed to extract KB size");
|
||||
assert_eq!(size.unwrap(), 2048 * 1024);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_extract_size_bytes() {
|
||||
let err = CandleError::Msg("allocate 1024 bytes failed".to_string());
|
||||
let size = extract_oom_size(&err);
|
||||
assert!(size.is_some(), "Failed to extract byte size");
|
||||
assert_eq!(size.unwrap(), 1024);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_extract_size_no_info() {
|
||||
// Error messages without size information
|
||||
let errors = vec![
|
||||
CandleError::Msg("out of memory".to_string()),
|
||||
CandleError::Msg("cuda error 2".to_string()),
|
||||
CandleError::Msg("allocation failed".to_string()),
|
||||
];
|
||||
|
||||
for err in errors {
|
||||
assert!(
|
||||
extract_oom_size(&err).is_none(),
|
||||
"Should return None for error without size: {:?}",
|
||||
err
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_extract_size_case_insensitive() {
|
||||
let errors = vec![
|
||||
CandleError::Msg("allocate 1GB".to_string()),
|
||||
CandleError::Msg("allocate 1gb".to_string()),
|
||||
CandleError::Msg("allocate 1Gb".to_string()),
|
||||
CandleError::Msg("allocate 1gB".to_string()),
|
||||
];
|
||||
|
||||
for err in errors {
|
||||
let size = extract_oom_size(&err);
|
||||
assert!(
|
||||
size.is_some(),
|
||||
"Failed case-insensitive extraction: {:?}",
|
||||
err
|
||||
);
|
||||
assert_eq!(size.unwrap(), 1024 * 1024 * 1024);
|
||||
}
|
||||
}
|
||||
|
||||
// ========================================================================
|
||||
// Edge Cases
|
||||
// ========================================================================
|
||||
|
||||
#[test]
|
||||
fn test_multiple_sizes_in_message() {
|
||||
// Should extract first size found
|
||||
let err = CandleError::Msg(
|
||||
"tried to allocate 2GB but only 1GB available".to_string()
|
||||
);
|
||||
let size = extract_oom_size(&err);
|
||||
assert!(size.is_some());
|
||||
assert_eq!(size.unwrap(), 2 * 1024 * 1024 * 1024);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_decimal_sizes() {
|
||||
let err = CandleError::Msg("allocate 1.5GB failed".to_string());
|
||||
let size = extract_oom_size(&err);
|
||||
assert!(size.is_some());
|
||||
|
||||
let expected = (1.5 * 1024.0 * 1024.0 * 1024.0) as usize;
|
||||
let actual = size.unwrap();
|
||||
let diff = if actual > expected {
|
||||
actual - expected
|
||||
} else {
|
||||
expected - actual
|
||||
};
|
||||
assert!(diff < expected / 100, "Decimal size mismatch");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_gib_vs_gb() {
|
||||
// Both GiB and GB should work (treated as binary)
|
||||
let err1 = CandleError::Msg("allocate 1GiB".to_string());
|
||||
let err2 = CandleError::Msg("allocate 1GB".to_string());
|
||||
|
||||
assert_eq!(
|
||||
extract_oom_size(&err1),
|
||||
extract_oom_size(&err2),
|
||||
"GiB and GB should extract same size"
|
||||
);
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user