foxhunt/WAVE_4_AGENT_2_DQN_CUDA_FIX_GUIDE.md

DQN CUDA Device Mismatch - Quick Fix Guide

Issue: DQN networks on GPU, input tensors on CPU → device mismatch errors Impact: 0% GPU utilization, 21.6% test failures, no GPU training possible Fix Time: 4-6 hours (3 files, ~50 lines changed)

---

Root Cause

WorkingDQN (GPU) → forward(input_cpu) → ERROR: device mismatch in matmul
       ↓
Q-network weights: CUDA
Input tensor: CPU
       ↓
Candle cannot multiply CPU × GPU tensors

---

Fix 1: Add Device to WorkingDQN (CRITICAL)

File: /home/jgrusewski/Work/foxhunt/ml/src/dqn/dqn.rs

Change 1: Add device field (line ~259)

pub struct WorkingDQN {
    config: WorkingDQNConfig,
    q_network: Sequential,
    target_network: Sequential,
    memory: Arc<Mutex<ExperienceReplayBuffer>>,
    epsilon: f32,
    training_steps: u64,
    optimizer: Option<Adam>,
    device: Device,  // ✅ ADD THIS LINE
}

Change 2: Store device in new() (line ~279)

pub fn new(config: WorkingDQNConfig) -> Result<Self, MLError> {
    let device = Device::cuda_if_available(0)?;

    let q_network = Sequential::new(
        config.state_dim,
        &config.hidden_dims,
        config.num_actions,
        device.clone(),  // ✅ Clone for q_network
    )?;

    let mut target_network = Sequential::new(
        config.state_dim,
        &config.hidden_dims,
        config.num_actions,
        device.clone(),  // ✅ Clone for target_network
    )?;

    // ... existing code ...

    Ok(Self {
        config,
        q_network,
        target_network,
        memory: Arc::new(Mutex::new(replay_buffer)),
        epsilon: config.epsilon_start,
        training_steps: 0,
        optimizer: Some(optimizer),
        device,  // ✅ Store device
    })
}

Change 3: Add device getter (after line ~274)

/// Get the device this DQN is using (CPU or CUDA)
pub fn device(&self) -> &Device {
    &self.device
}

Change 4: Auto-convert inputs in forward() (line ~42)

pub fn forward(&self, state: &Tensor) -> Result<Tensor, MLError> {
    // Auto-convert input to correct device if needed
    let state = if state.device() != &self.device {
        state.to_device(&self.device).map_err(|e| {
            MLError::ModelError(format!("Failed to move tensor to device: {}", e))
        })?
    } else {
        state.clone()
    };

    self.q_network.forward(&state)
}

---

Fix 2: Update DQNTrainableAdapter (CRITICAL)

File: /home/jgrusewski/Work/foxhunt/ml/src/dqn/trainable_adapter.rs

Change 1: Add device field (line ~16)

pub struct DQNTrainableAdapter {
    dqn: WorkingDQN,
    config: WorkingDQNConfig,
    device: Device,  // ✅ ADD THIS LINE
    learning_rate: f64,
    latest_metrics: TrainingMetrics,
    current_step: usize,
    loss_history: Vec<f64>,
}

Change 2: Store device in new() (line ~33)

pub fn new(config: WorkingDQNConfig) -> Result<Self, MLError> {
    let learning_rate = config.learning_rate;
    let dqn = WorkingDQN::new(config.clone())?;
    let device = dqn.device().clone();  // ✅ Get device from DQN

    Ok(Self {
        dqn,
        config,
        device,  // ✅ Store device
        learning_rate,
        latest_metrics: TrainingMetrics::default(),
        current_step: 0,
        loss_history: Vec::new(),
    })
}

Change 3: Fix device() method (line ~88)

fn device(&self) -> &Device {
    &self.device  // ✅ Return stored device (not hardcoded CPU)
}

---

Fix 3: Update load_checkpoint (IMPORTANT)

File: /home/jgrusewski/Work/foxhunt/ml/src/dqn/trainable_adapter.rs

Change: Load tensors to correct device (line ~254)

fn load_checkpoint(&mut self, checkpoint_path: &str) -> Result<CheckpointMetadata, MLError> {
    // ... existing metadata loading ...

    let safetensors_path = format!("{}.safetensors", checkpoint_path);
    let tensors = candle_core::safetensors::load(
        &safetensors_path,
        &self.device  // ✅ Load to actual device (not CPU)
    ).map_err(|e| {
        MLError::CheckpointError(format!("Failed to load safetensors: {}", e))
    })?;

    // ... rest of checkpoint loading ...
}

---

Fix 4: Update Tests (MEDIUM PRIORITY)

Pattern for All Test Files

Every test that creates input tensors must use the model's device:

// ❌ OLD (creates CPU tensor)
let state = Tensor::zeros(&[1, state_dim], DType::F32, &Device::Cpu)?;

// ✅ NEW (creates tensor on model's device)
let device = Device::cuda_if_available(0)?;
let state = Tensor::zeros(&[1, state_dim], DType::F32, &device)?;

// OR (get device from model)
let device = dqn.device();
let state = Tensor::zeros(&[1, state_dim], DType::F32, device)?;

Files to Update

1. /home/jgrusewski/Work/foxhunt/ml/tests/dqn_tests.rs

   • Lines with &Device::Cpu tensor creation
   • 8 real-data tests

2. /home/jgrusewski/Work/foxhunt/ml/tests/dqn_edge_cases_test.rs

   • Input tensor creation for edge cases

3. /home/jgrusewski/Work/foxhunt/ml/tests/dqn_rainbow_test.rs

   • Rainbow agent real market data tests

4. /home/jgrusewski/Work/foxhunt/ml/src/dqn/agent_new_tests.rs

   • Agent validation tests with real features

---

Validation Checklist

After applying fixes, run these checks:

1. Compilation

cargo build -p ml --release
# Should compile without errors

2. Basic Tests

cargo test -p ml --test dqn_tests --release -- --nocapture
# Expected: 37/37 passing (100%)

3. Device Verification

cargo test -p ml --test verify_dqn_cuda --release -- --nocapture
# Expected: test_dqn_uses_cuda_device PASS
# Output should show: "✅ DQN is using CUDA GPU acceleration"

4. GPU Memory Usage

# In terminal 1:
watch -n 1 nvidia-smi

# In terminal 2:
cargo test -p ml --test dqn_tests --release

# Expected during test: 50-150MB GPU memory usage

5. Real Data Tests

cargo test -p ml --test dqn_tests test_dqn_training_with_real_market_data --release -- --nocapture
# Expected: PASS with GPU acceleration

---

Expected Outcomes After Fix

Test Results

• Pass Rate: 37/37 (100%) ← was 29/37 (78.4%)
• Device Mismatch Errors: 0 ← was 8
• Real Data Tests: All passing ← 6 failing

GPU Utilization

• Memory Usage: 50-150MB ← was 3MB (idle)
• Temperature: 60-70°C ← was 46°C (idle)
• Utilization: 20-40% ← was 0%

Performance

• Q-network Forward Pass: <100μs (GPU) ← N/A (failed)
• Training Step: 10-50x faster than CPU
• Experience Replay: GPU tensor operations functional

---

Testing Methodology

Sequential Validation

1. Apply Fix 1 (WorkingDQN) → compile → test device getter
2. Apply Fix 2 (Adapter) → compile → test adapter.device()
3. Apply Fix 3 (checkpoint) → compile → test save/load
4. Apply Fix 4 (tests) → run full test suite

Rollback Plan

If fix breaks other tests:

git diff ml/src/dqn/dqn.rs > /tmp/dqn_fix.patch
git checkout ml/src/dqn/dqn.rs  # Rollback
# Analyze issue, adjust fix, re-apply

---

Risk Assessment

Low Risk Changes

• Adding device field to WorkingDQN (backward compatible)
• Adding device() getter (new method, no conflicts)
• Storing device in adapter (internal state)

Medium Risk Changes

• Auto-converting inputs in forward() (performance impact: +1-2μs per call)
• Changing checkpoint load device (could break existing checkpoints on CPU)

Mitigation

• Test with both CPU and GPU checkpoints
• Add device validation in checkpoint load
• Document device conversion overhead

---

Performance Impact

Expected Improvements

• Training Speed: 10-50x faster (GPU vs CPU)
• Inference Latency: 100μs → <10μs (GPU acceleration)
• Batch Processing: 1000 experiences in ~5ms (was ~200ms on CPU)

Negligible Overhead

• Device check in forward(): ~0.1μs (branch prediction)
• Auto-conversion (if needed): ~2μs per tensor (rarely triggered)

---

Common Pitfalls

Pitfall 1: Forgetting to Clone Device

// ❌ WRONG - moves device
let q_network = Sequential::new(..., device)?;
let target_network = Sequential::new(..., device)?;  // ERROR: device moved

// ✅ CORRECT - clone device
let q_network = Sequential::new(..., device.clone())?;
let target_network = Sequential::new(..., device.clone())?;

Pitfall 2: Not Updating All Test Files

• Must update ALL files that create input tensors
• Use rg "Device::Cpu" ml/tests/ to find all occurrences

Pitfall 3: Checkpoint Device Incompatibility

• CPU checkpoints loaded on GPU device → works (auto-converts)
• GPU checkpoints loaded on CPU device → works but slower
• Document device in checkpoint metadata

---

Summary

Total Changes: 3 files, ~50 lines Risk Level: Low-Medium (well-isolated changes) Test Coverage: 100% (all existing tests validate fix) Estimated Time: 4-6 hours (including testing)

Priority: CRITICAL - Blocks Wave 4 Agent 3 (PPO) and Agent 4 (TFT)

Success Metric:

• Test pass rate: 78.4% → 100%
• GPU memory: 3MB → 50-150MB
• GPU utilization: 0% → 20-40%

Validation Command:

cargo test -p ml dqn --release && \
  watch -n 1 "nvidia-smi --query-gpu=memory.used --format=csv,noheader,nounits"

Expected: Tests pass + GPU memory rises to 50-150MB during execution

---

Document Version: 1.0 Last Updated: 2025-10-15 Agent: Wave 4 Agent 2 Status: Ready for implementation