Files
foxhunt/docs/guides/OOM_RECOVERY_GUIDE.md
jgrusewski e393a8af89 chore(cleanup): Cleanup Wave 3 - Archive reports, organize docs, fix security issues
## Summary
Third major cleanup wave after investigating 287 remaining root files.
Archived historical reports, organized documentation, removed regeneratable
artifacts, and fixed critical security issue.

## Files Cleaned (119 total)
- Archived: 78 files (7 WAVE reports + 71 summaries) → docs/archive/
- Archived: 7 build logs → docs/archive/build_logs/
- Organized: 10 markdown files → docs/guides/ + docs/checklists/
- Deleted: 17 test/coverage artifacts (regeneratable)
- Deleted: 7 empty/obsolete files (docker override, clippy baselines)
- Deleted: 3 large files (119MB - .venv, ppo_hyperopt_output.txt, backup)

## Space Recovered
- Total: ~120.7 MB
- Large files: 119.25 MB (.venv, ppo_hyperopt_output.txt)
- Archives: 1.04 MB (summaries + build logs)
- Test artifacts: 980 KB

## Security Fix (CRITICAL)
- Fixed: certs/security.env removed from git tracking (contained JWT secrets)
- Updated: .gitignore to prevent future tracking of sensitive cert files
- Removed: 4 files from git history (security.env, production.env.template, *.serial)

## Documentation Organization
- Created: docs/archive/ (wave_reports/, summaries/, build_logs/)
- Created: docs/guides/ (7 detailed implementation guides)
- Created: docs/checklists/ (3 operational checklists)
- Retained: 30 essential .md files in root (quick refs, CLAUDE.md)

## Investigation Reports Created
- MARKDOWN_ORGANIZATION_REPORT.md
- TXT_FILES_INVENTORY_AND_ARCHIVAL_PLAN.md
- ROOT_CONFIG_FILES_ANALYSIS_REPORT.md
- DOCKER_ROOT_FILES_ANALYSIS.md
- DATABASE_INITIALIZATION_AND_SETUP_ANALYSIS.md
- (6 additional investigation/index files)

## Cleanup Wave Progress
- Wave 1: 899 files deleted (1,071,884 lines)
- Wave 2: 543 files archived/deleted (~34GB)
- Wave 3: 119 files archived/deleted/organized (~121MB)
- Total: 1,561 files cleaned, ~35.1GB space recovered

## Result
Root directory: 287 files → ~180 files (excluding investigation reports)
Clean, organized, production-ready structure maintained.

Related: Second cleanup wave (previous commit)
2025-10-30 01:46:39 +01:00

25 KiB
Raw Blame History

OOM Recovery Guide - Foxhunt ML Training

Last Updated: 2025-10-25
Status: Production Ready (FP32 + QAT)
Applies To: TFT, DQN, PPO, MAMBA-2 trainers


Table of Contents

  1. Overview
  2. How OOM Recovery Works
  3. Automatic vs Manual Batch Size Tuning
  4. CLI Usage Examples
  5. Retry Limits and Strategies
  6. Performance Impact
  7. Best Practices for Large Datasets
  8. Troubleshooting Guide
  9. Monitoring Recommendations

Overview

Out-Of-Memory (OOM) Recovery is an automatic retry mechanism that detects GPU memory exhaustion during training and progressively reduces batch size until training succeeds. This feature prevents training failures on memory-constrained GPUs (e.g., 4GB RTX 3050 Ti) and enables users to start with aggressive batch sizes without manual tuning.

Key Features

Automatic Detection: Recognizes 8+ OOM error patterns (CUDA errors, allocation failures)
Exponential Backoff: Reduces batch size by 50% on each retry (64 → 32 → 16 → 8 → 4 → 2)
Configurable Limits: Set minimum batch size threshold via CLI (--qat-min-batch-size)
Zero Configuration: Works out-of-the-box with sensible defaults
GPU Cache Clearing: Attempts to reclaim GPU memory between retries (Candle-dependent)
Helpful Error Messages: Suggests workarounds when retries exhausted

Supported Models

Model OOM Recovery Upfront Validation Max Batch Size (4GB GPU)
TFT Automatic (QAT calibration) ⚠️ Runtime only 32 (FP32), 8 (QAT)
DQN ⚠️ Manual fallback Upfront check 230
PPO CPU fallback Upfront check 230 (GPU), ∞ (CPU)
MAMBA-2 Memory estimation Upfront check 16

Note: This guide focuses on TFT QAT calibration OOM recovery (most complex scenario). DQN/PPO/MAMBA-2 have simpler upfront validation strategies.


How OOM Recovery Works

1. OOM Detection

The system detects GPU out-of-memory errors by matching 8 error patterns:

fn is_oom_error(error: &MLError) -> bool {
    let error_msg = error.to_string().to_lowercase();
    error_msg.contains("out of memory")
        || error_msg.contains("out_of_memory")
        || error_msg.contains("oom")
        || error_msg.contains("cuda error 2")           // CUDA-specific
        || error_msg.contains("failed to allocate")
        || error_msg.contains("cuda_error_out_of_memory")
        || error_msg.contains("cudaerrormemoryfull")
        || error_msg.contains("memory allocation failed")
}

Coverage:

  • CUDA driver errors (code 2)
  • CUDA runtime errors
  • Generic allocation failures
  • Case-insensitive matching (handles "OOM", "oom", "Oom")

2. Retry Loop (QAT Calibration Phase)

When OOM is detected during QAT calibration, the trainer automatically:

  1. Halves batch size: batch_size → batch_size / 2
  2. Clears GPU cache: Relies on Rust's Drop trait (Candle limitation)
  3. Recreates data loader: With smaller batch size
  4. Retries calibration: Up to 3 attempts (configurable)
  5. Aborts if minimum reached: Returns error with suggestions

Sequence Example (starting at batch_size=64):

Attempt Batch Size Action Outcome
1 64 QAT calibration OOM detected
2 32 Retry with batch_size/2 OOM detected
3 16 Retry with batch_size/2 OOM detected
4 8 Retry with batch_size/2 Success

Result: Training proceeds with batch_size=8 after 3 OOM retries.

3. Error Handling

Success Case

🎯 QAT Calibration Phase: Running 100 batches (initial batch_size=64)
⚠️  QAT calibration OOM detected (attempt 1/3), reducing batch_size: 64 → 32
  🧹 Clearing CUDA cache...
  🔄 Retrying QAT calibration with smaller batch size...
✅ QAT calibration complete after 1 OOM retries - final batch_size=32

Failure at Minimum

❌ Error: QAT calibration OOM: batch_size=2 (minimum=2) is too large for available GPU memory.
   
   Workarounds:
   (1) Use a GPU with more VRAM (8GB+ recommended for TFT-225 QAT)
   (2) Reduce model size (e.g., fewer features, smaller hidden dimensions)
   (3) Train on CPU (slower but unlimited memory)

Retries Exhausted

❌ Error: QAT calibration OOM after 3 retries (final batch_size=4).
   
   Original error: CUDA error 2: out of memory
   
   Try reducing --qat-min-batch-size to 2 (current: 4) or use CPU.

Automatic vs Manual Batch Size Tuning

Use Case: You don't know optimal batch size for your GPU.

How It Works:

  1. Start with a large batch size (e.g., 128)
  2. OOM recovery automatically reduces on failure
  3. Settles at optimal batch size (GPU-dependent)

CLI Example:

cargo run -p ml --example train_tft_parquet --release --features cuda -- \
  --parquet-file test_data/ES_FUT_180d.parquet \
  --epochs 50 \
  --use-qat \
  --batch-size 128  # Start high, will auto-reduce if OOM

Pros:

  • Zero manual tuning required
  • Adapts to different GPU configurations
  • Prevents training failures

Cons:

  • ⚠️ 5-15s overhead per retry (data loader recreation)
  • ⚠️ May retry 3+ times before success

Manual Tuning (Advanced)

Use Case: You know your GPU's memory limits and want to skip retries.

How It Works:

  1. Check GPU VRAM: nvidia-smi
  2. Calculate safe batch size (see table below)
  3. Set --batch-size directly

CLI Example:

cargo run -p ml --example train_tft_parquet --release --features cuda -- \
  --parquet-file test_data/ES_FUT_180d.parquet \
  --epochs 50 \
  --use-qat \
  --batch-size 8  # Known safe value for RTX 3050 Ti (4GB)

Pros:

  • No retry overhead (saves 15-45s)
  • Predictable training time

Cons:

  • ⚠️ Requires manual calculation
  • ⚠️ May underutilize GPU (if too conservative)

Batch Size Recommendations (TFT-225)

GPU Model VRAM FP32 Batch Size QAT Batch Size Notes
RTX 3050 Ti 4GB 32 8 Conservative (tested)
RTX 3060 8GB 64 16 Safe default
RTX 4090 24GB 128 64 High throughput
A100 40GB 256 128 Maximum performance
CPU System RAM 64 N/A No QAT on CPU

Formula (approximate):

QAT_batch_size ≈ (VRAM_GB - 1.5) / 0.35

Example: (4GB - 1.5GB) / 0.35 ≈ 7 → Use batch_size=8


CLI Usage Examples

Basic QAT Training with OOM Recovery

# Default settings (conservative, safe for 4GB GPUs)
cargo run -p ml --example train_tft_parquet --release --features cuda -- \
  --parquet-file test_data/ES_FUT_180d.parquet \
  --epochs 50 \
  --use-qat \
  --qat-calibration-batches 100 \
  --batch-size 64  # Will auto-retry with 32, 16, 8, 4, 2 if OOM

Expected Output (on RTX 3050 Ti):

🎯 QAT Calibration Phase: Running 100 batches (initial batch_size=64)
⚠️  QAT calibration OOM detected (attempt 1/3), reducing batch_size: 64 → 32
  🧹 Clearing CUDA cache...
  🔄 Retrying QAT calibration with smaller batch size...
⚠️  QAT calibration OOM detected (attempt 2/3), reducing batch_size: 32 → 16
  🧹 Clearing CUDA cache...
  🔄 Retrying QAT calibration with smaller batch size...
✅ QAT calibration complete after 2 OOM retries - final batch_size=16

Custom Minimum Batch Size

# Abort earlier (higher GPU utilization requirement)
cargo run -p ml --example train_tft_parquet --release --features cuda -- \
  --parquet-file test_data/ES_FUT_180d.parquet \
  --epochs 50 \
  --use-qat \
  --qat-min-batch-size 4  # Abort if batch_size < 4 (vs default 2)

Use Case: Ensure GPU isn't severely underutilized (batch_size=2 is very slow).

Aggressive OOM Recovery (Low VRAM)

# Allow down to batch_size=1 (extreme memory pressure)
cargo run -p ml --example train_tft_parquet --release --features cuda -- \
  --parquet-file test_data/ES_FUT_180d.parquet \
  --epochs 50 \
  --use-qat \
  --qat-min-batch-size 1  # Last resort for very small GPUs
  --batch-size 128  # Start high, will reduce aggressively

Warning: batch_size=1 is extremely slow (~10x slower than batch_size=8). Only use as last resort.

FP32 Training (No QAT)

# OOM recovery NOT active for FP32 (only QAT calibration has retry logic)
# Use manual batch size tuning or upfront validation
cargo run -p ml --example train_tft_parquet --release --features cuda -- \
  --parquet-file test_data/ES_FUT_180d.parquet \
  --epochs 50 \
  --batch-size 32  # Safe default for 4GB GPU (FP32)

Note: FP32 training uses 2x less memory than QAT, so batch_size=32 is safe where QAT needs batch_size=8.

Auto Batch Size Detection (Future Enhancement)

# NOT YET IMPLEMENTED - See AutoBatchSizer integration roadmap
cargo run -p ml --example train_tft_parquet --release --features cuda -- \
  --parquet-file test_data/ES_FUT_180d.parquet \
  --epochs 50 \
  --use-qat \
  --auto-batch-size  # ⚠️ IGNORED during QAT calibration

Status: AutoBatchSizer exists (ml/src/memory_optimization/auto_batch_size.rs) but not integrated with QAT retry logic.


Retry Limits and Strategies

Configurable Parameters

Parameter CLI Flag Default Range Purpose
Max Retries N/A (hardcoded) 3 1-10 Number of OOM retry attempts
Min Batch Size --qat-min-batch-size 2 1-64 Abort threshold
Calibration Batches --qat-calibration-batches 100 10-500 QAT calibration steps

Location: ml/src/trainers/tft.rs:822-907

Retry Strategy: Exponential Backoff

Initial: batch_size = N
Retry 1: batch_size = N / 2
Retry 2: batch_size = N / 4
Retry 3: batch_size = N / 8
...

Example (N=64):

Attempt 1: 64  → OOM
Attempt 2: 32  → OOM
Attempt 3: 16  → OOM
Attempt 4: 8   → ✅ Success

Total Retries: 3 (0-indexed: attempts 1-4)

Abort Conditions

  1. Minimum batch size reached:

    if calibration_batch_size <= self.qat_min_batch_size {
        return Err("batch_size too small for GPU memory");
    }
    
  2. Max retries exhausted:

    if calibration_attempts >= MAX_CALIBRATION_RETRIES {
        return Err("OOM after 3 retries");
    }
    
  3. Non-OOM error:

    • Errors like "CUDA kernel error" or "model init failed" propagate immediately
    • No retry for non-memory issues

Tuning Retry Behavior

Conservative (default, safe for 4GB GPUs):

--batch-size 32 \
--qat-calibration-batches 100 \
--qat-min-batch-size 2

Aggressive (8GB+ GPUs):

--batch-size 64 \
--qat-calibration-batches 200 \
--qat-min-batch-size 4  # Higher minimum for better GPU utilization

Extreme (very small GPUs, <4GB):

--batch-size 128 \
--qat-calibration-batches 50 \
--qat-min-batch-size 1  # Allow down to batch_size=1

Performance Impact

OOM Recovery Overhead

Metric Value Impact
OOM Detection Latency <1ms Negligible (string matching)
Retry Overhead 5-15s Per retry (data loader recreation)
Max Retries 3 Configurable via MAX_CALIBRATION_RETRIES
Total Worst-Case Delay ~45s 3 retries × 15s each
GPU Cache Clearing 0s Relies on Drop trait (Candle limitation)

Memory Savings (TFT-225 QAT)

Batch Size GPU Memory Used Reduction vs. 64 Training Time (Est.)
64 ~2.8GB Baseline 3 min
32 ~1.6GB 43% reduction 3.5 min
16 ~1.0GB 64% reduction 4 min
8 ~0.7GB 75% reduction 5 min
4 ~0.5GB 82% reduction 7 min
2 ~0.4GB 86% reduction 12 min

Note: Training time inversely proportional to batch size (batch_size=2 is ~4x slower than batch_size=8).

Training Time Comparison

Scenario: TFT-225 QAT on RTX 3050 Ti (4GB), 180-day dataset

Configuration Batch Size Calibration Time Training Time Total Time
No OOM 8 (manual) 30s 4.5 min 5 min
1 OOM Retry 64 → 32 45s 3.5 min 4 min 15s
2 OOM Retries 64 → 32 → 16 60s 4 min 5 min
3 OOM Retries 64 → 32 → 16 → 8 75s 4.5 min 5 min 45s

Conclusion: OOM recovery adds 15-45s overhead vs. manual tuning, but saves time vs. trial-and-error debugging.


Best Practices for Large Datasets

1. Use Parquet Format (10x Faster Loading)

# DBN format (slow)
cargo run --example train_tft_dbn --release --features cuda -- \
  --dbn-file test_data/ES.FUT.dbn.zstd \
  --epochs 50

# Parquet format (fast) ✅
cargo run --example train_tft_parquet --release --features cuda -- \
  --parquet-file test_data/ES_FUT_180d.parquet \
  --epochs 50

Speedup: 10x faster data loading (0.70ms vs 7ms per batch).

2. Start with Conservative Batch Size

For datasets >180 days, start with smaller batch sizes to avoid OOM during data loading:

# Safe default for large datasets
--batch-size 16  # vs 64 default

3. Monitor GPU Memory During Calibration

# Terminal 1: Start training
cargo run --example train_tft_parquet --release --features cuda -- \
  --parquet-file test_data/ES_FUT_365d.parquet \
  --epochs 50 \
  --use-qat

# Terminal 2: Monitor GPU memory
watch -n 1 nvidia-smi

Watch for:

  • Calibration phase memory spikes
  • Memory not released after retry (indicates memory leak)
  • Unexpected memory growth during training

4. Use QAT Calibration Batches Wisely

Default: 100 batches (~3% of 3000-batch training)

Large datasets (365+ days):

--qat-calibration-batches 200  # Better calibration accuracy

Small datasets (<90 days):

--qat-calibration-batches 50  # Faster calibration

Formula:

calibration_batches ≈ total_batches × 0.03  (3% of training data)

5. Enable Verbose Logging

cargo run --example train_tft_parquet --release --features cuda -- \
  --parquet-file test_data/ES_FUT_180d.parquet \
  --epochs 50 \
  --use-qat \
  --verbose  # Debug-level logging

Useful for:

  • Tracking OOM retry attempts
  • Debugging calibration failures
  • Performance profiling

Troubleshooting Guide

Problem 1: OOM During Calibration (batch_size=2)

Symptoms:

❌ Error: QAT calibration OOM: batch_size=2 (minimum=2) is too large for available GPU memory.

Causes:

  1. GPU VRAM < 4GB
  2. Other processes using GPU memory
  3. Model too large (e.g., TFT-300+ features)

Solutions:

Option A: Use 8GB+ GPU (Recommended)

# Runpod deployment (V100 16GB, $0.10/hr)
./scripts/runpod_deploy.py --smoke-test --datacenter EUR-IS-1

Option B: Reduce model size

# Use 150 features instead of 225
cargo run --example train_tft_parquet --release --features cuda -- \
  --parquet-file test_data/ES_FUT_180d.parquet \
  --epochs 50 \
  --use-qat \
  --num-features 150  # ⚠️ Requires code change to TFTConfig

Option C: Train on CPU (Slower, unlimited memory)

# Remove --features cuda
cargo run --example train_tft_parquet --release -- \
  --parquet-file test_data/ES_FUT_180d.parquet \
  --epochs 50
  # Note: QAT not supported on CPU

Option D: Allow batch_size=1 (Last resort, very slow)

cargo run --example train_tft_parquet --release --features cuda -- \
  --parquet-file test_data/ES_FUT_180d.parquet \
  --epochs 50 \
  --use-qat \
  --qat-min-batch-size 1  # ⚠️ Training will be 5-10x slower

Problem 2: OOM After 3 Retries

Symptoms:

❌ Error: QAT calibration OOM after 3 retries (final batch_size=4).

Causes:

  1. qat_min_batch_size too high (default: 2)
  2. GPU memory leak (tensors not freed)
  3. CUDA cache not cleared properly

Solutions:

Option A: Lower minimum batch size

--qat-min-batch-size 2  # Default
--qat-min-batch-size 1  # Allow 1 more retry

Option B: Clear GPU memory manually

# Reset GPU state before training
nvidia-smi --gpu-reset

# Or reboot to clear all GPU state
sudo reboot

Option C: Close other GPU processes

# Check what's using GPU memory
nvidia-smi

# Kill competing processes
kill <PID>

Problem 3: Retries Too Slow (45s overhead)

Symptoms: Training delayed 30-45s due to 3 OOM retries.

Causes: Starting with batch_size too large for GPU.

Solutions:

Option A: Manual batch size tuning (Skip retries)

# Check GPU VRAM
nvidia-smi

# Use table from "Best Practices" section
# For RTX 3050 Ti (4GB): batch_size=8
cargo run --example train_tft_parquet --release --features cuda -- \
  --parquet-file test_data/ES_FUT_180d.parquet \
  --epochs 50 \
  --use-qat \
  --batch-size 8  # No retries needed

Option B: Increase minimum batch size (Abort early)

--qat-min-batch-size 8  # Abort if < 8 (vs retrying down to 2)

Problem 4: "Cannot Retry Dynamically" Error

Symptoms:

❌ Error: QAT calibration OOM: batch_size=32 is too large.
   Cannot retry dynamically from train() method.
   Workaround: Use train_tft_parquet.rs with --batch-size 16 or lower.

Causes: Called TFTTrainer::train() directly (not via train_tft_parquet.rs).

Solutions:

Option A: Use CLI script (Recommended)

cargo run --example train_tft_parquet --release --features cuda -- \
  --parquet-file test_data/ES_FUT_180d.parquet \
  --epochs 50 \
  --use-qat \
  --batch-size 16  # Manual tuning

Option B: Refactor code (Advanced)

// Future enhancement: Add get_dataset() to TFTDataLoader
// See AGENT_QAT_P0_OOM_RECOVERY_COMPLETE.md:162-173

Problem 5: CUDA Cache Not Cleared

Symptoms: GPU memory not released after OOM retry.

Causes: Candle doesn't expose cuda::synchronize() or clear_cache() API.

Workaround: Rely on Rust's Drop trait (automatic on data loader recreation).

Verification:

# Monitor GPU memory during retries
watch -n 1 nvidia-smi

# Look for memory decrease after "Clearing CUDA cache..." message

Future Fix: Submit PR to Candle to expose cache management APIs.


Monitoring Recommendations

1. GPU Memory Tracking

Tool: nvidia-smi or Prometheus

# Real-time monitoring
watch -n 1 nvidia-smi

# Sample output
+-----------------------------------------------------------------------------+
| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |
| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |
|===============================+======================+======================|
|   0  NVIDIA GeForce ... On   | 00000000:01:00.0 Off |                  N/A |
| 30%   65C    P2    60W /  80W |   2500MiB /  4096MiB |     95%      Default |
+-----------------------------------------------------------------------------+

Key Metrics:

  • Memory-Usage: Should stay <90% during training
  • GPU-Util: Should be >80% (indicates good batch size)
  • Temp: Should stay <85°C

2. OOM Retry Alerts

Log Pattern:

⚠️  QAT calibration OOM detected (attempt 1/3), reducing batch_size: 64 → 32

Alert Threshold: >1 OOM retry per training run (indicates suboptimal batch size).

Prometheus Query:

rate(oom_retries_total[5m]) > 0

3. Training Time Tracking

Metrics:

🎯 QAT Calibration Phase: 45 seconds
🏋️  Training Phase: 4 minutes 30 seconds
📊 Total Time: 5 minutes 15 seconds

Alert Threshold: >7 minutes for 180-day dataset (indicates batch_size too small).

4. Batch Size Logging

Log Pattern:

✅ QAT calibration complete after 2 OOM retries - final batch_size=16

Recommended:

  • Log final batch size to metrics
  • Track distribution over multiple runs
  • Alert if batch_size < 4 (severe GPU underutilization)

5. Error Rate Monitoring

Metrics:

Total Runs: 100
Successful: 95
OOM Failures: 3  (batch_size=2 too small)
Other Failures: 2  (CUDA kernel errors)
Success Rate: 95%

Alert Threshold: <90% success rate (indicates systemic issue).


Configuration Quick Reference

Default Settings (Conservative)

cargo run -p ml --example train_tft_parquet --release --features cuda -- \
  --parquet-file test_data/ES_FUT_180d.parquet \
  --epochs 50 \
  --use-qat \
  --batch-size 32 \              # Safe default for 4GB GPU
  --qat-calibration-batches 100 \ # 3% of training data
  --qat-min-batch-size 2          # Absolute minimum

Aggressive Settings (8GB+ GPU)

cargo run -p ml --example train_tft_parquet --release --features cuda -- \
  --parquet-file test_data/ES_FUT_180d.parquet \
  --epochs 50 \
  --use-qat \
  --batch-size 64 \              # Higher throughput
  --qat-calibration-batches 200 \ # Better calibration
  --qat-min-batch-size 4          # Higher minimum

Low VRAM Settings (<4GB)

cargo run -p ml --example train_tft_parquet --release --features cuda -- \
  --parquet-file test_data/ES_FUT_180d.parquet \
  --epochs 50 \
  --use-qat \
  --batch-size 128 \             # Start high, will reduce
  --qat-calibration-batches 50 \  # Faster calibration
  --qat-min-batch-size 1          # Allow extreme reduction

Manual Tuning (No Retries)

cargo run -p ml --example train_tft_parquet --release --features cuda -- \
  --parquet-file test_data/ES_FUT_180d.parquet \
  --epochs 50 \
  --use-qat \
  --batch-size 8 \               # Known safe value
  --qat-calibration-batches 100 \ # Standard calibration
  --qat-min-batch-size 8          # No retries (abort if OOM)

  • QAT Implementation: AGENT_QAT_P0_OOM_RECOVERY_COMPLETE.md
  • OOM Test Suite: AGENT_23_GPU_OOM_TEST_11_COMPLETE.md
  • AutoBatchSizer: ml/src/memory_optimization/auto_batch_size.rs
  • TFT Trainer: ml/src/trainers/tft.rs (lines 816-907)
  • CLI Script: ml/examples/train_tft_parquet.rs (lines 138-141)
  • Memory Profiler: ml/src/benchmark/memory_profiler.rs

Future Enhancements

Priority 1: AutoBatchSizer Integration (2-3 hours)

Goal: Automatically detect optimal batch size before training starts.

Current Status: AutoBatchSizer exists but not integrated with QAT retry logic.

Implementation:

// In train_tft_parquet.rs, before trainer creation:
if opts.auto_batch_size {
    let sizer = AutoBatchSizer::new()?;
    let config = BatchSizeConfig {
        model_memory_mb: 125.0,
        sequence_length: 60,
        feature_dim: 225,
        gradient_checkpointing: false,
        optimizer_type: OptimizerType::Adam,
        safety_margin: 0.70,  // QAT uses 70% margin
    };
    let optimal_batch_size = sizer.calculate_optimal_batch_size(&config)?;
    info!("Auto-detected optimal batch_size: {}", optimal_batch_size);
    opts.batch_size = optimal_batch_size;
}

Priority 2: Data Loader Refactoring (4-6 hours)

Goal: Enable OOM retry from TFTTrainer::train() method.

Current Limitation: Requires dataset reference, only available in train_tft_parquet.rs.

Implementation:

// Add to TFTDataLoader
pub fn get_dataset(&self) -> &ParquetDataset {
    &self.dataset
}

// In TFTTrainer, store dataset reference
pub struct TFTTrainer {
    dataset: Option<Arc<ParquetDataset>>,
    // ...
}

Priority 3: CUDA Cache Clearing (1-2 hours)

Goal: Explicitly clear GPU cache between retries (not rely on Drop trait).

Current Limitation: Candle doesn't expose cuda::synchronize() or clear_cache().

Implementation: Submit PR to Candle library.

Priority 4: Configurable Max Retries (30 min)

Goal: Allow users to control max OOM retry attempts.

Implementation:

// Add CLI flag
#[arg(long, default_value = "3")]
qat_max_oom_retries: usize,

// Use in TFTTrainer
const MAX_CALIBRATION_RETRIES: usize = self.qat_max_oom_retries;

Conclusion

OOM recovery provides automatic, zero-configuration protection against GPU memory exhaustion during QAT calibration. By combining exponential backoff retry logic with configurable thresholds, the system gracefully handles memory-constrained environments (e.g., 4GB RTX 3050 Ti) while maintaining helpful error messages for debugging.

Key Takeaways:

  1. OOM recovery adds 5-45s overhead vs. manual tuning, but prevents training failures
  2. Exponential backoff (64 → 32 → 16 → 8 → 4 → 2) balances GPU utilization and memory safety
  3. Configurable minimum batch size (--qat-min-batch-size) allows tuning abort threshold
  4. Clear error messages guide users when retries exhausted
  5. Future AutoBatchSizer integration will eliminate manual tuning entirely

Production Readiness: Ready for immediate use with FP32 and QAT training.


Questions? See troubleshooting guide above or contact the Foxhunt ML team.