# AutoBatchSizer VRAM Budget Fix — Implementation Plan

> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (`- [ ]`) syntax for tracking.

**Goal:** Fix AutoBatchSizer to account for ALL GPU memory consumers, preventing OOM on any GPU while maximizing batch size for performance.

**Architecture:** Replace the current `model_memory_mb`-only budget with a comprehensive VRAM budget that includes IQN (~1.1GB at batch=8192), attention (28MB), IQL (16MB), replay buffer (70% VRAM), experience collector, and CUDA graph overhead. The batch_size is computed LAST, using only the VRAM remaining after all fixed allocations.

**Tech Stack:** Rust, cudarc 0.17.3, CUDA driver API

---

## Problem Statement

The current `AutoBatchSizer::max_safe_batch_size()` only considers:
- Model parameters (DQN trunk: ~1.2MB)
- Optimizer states (Adam: 2x model)
- Gradients (1x model)
- Activations (1x model)
- Safety margin (20%)
- Per-sample batch data (80 features x bf16 = 160 bytes)

It IGNORES:
- **IQN head**: 1.1GB at batch=8192 (dominates VRAM on small GPUs)
- **Attention**: 28MB at batch=8192
- **IQL**: 16MB at batch=8192
- **HER**: ~1MB
- **Replay buffer**: 70% of VRAM (allocated BEFORE batch sizer runs)
- **Experience collector**: 40MB+ (cuBLAS handles, output buffers)
- **CUDA graphs**: ~10MB (captured kernel topology)

Result: On RTX 3050 (4GB, 1.7GB free), returns batch=8192 when safe max is ~128.

## Root Cause

1. `AutoBatchSizer::new()` probes `free_memory_mb` at constructor time
2. But the replay buffer (70% VRAM) is allocated AFTER the batch sizer runs
3. The `model_memory_mb` passed to `max_safe_batch_size()` is only the DQN trunk (~1.2MB), not the full training pipeline
4. IQN VRAM scales linearly with batch_size but isn't in the per-sample calculation

## VRAM Allocation Order (current)

```
1. AutoBatchSizer probes free VRAM     → sees 81GB (H100) or 1.7GB (RTX 3050)
2. batch_size = max_safe_batch(config)  → computes from DQN trunk only
3. AutoReplaySizer allocates buffer     → consumes 70% of free VRAM
4. FusedTrainingCtx::new() allocates:
   a. GpuDqnTrainer (batch buffers + optimizer states)
   b. GpuHer
   c. GpuIqlTrainer
   d. GpuIqnHead (HUGE: 1.1GB at batch=8192)
   e. GpuAttention (28MB at batch=8192)
5. Experience collector allocates       → 40MB+
```

The batch_size is decided at step 2 but the big allocations happen at steps 3-5.

## Fix: VRAM Budget Model

```
free_vram = probe_free_vram()

# Fixed allocations (batch-independent)
replay_vram    = free_vram * replay_buffer_vram_fraction  # 70% default
exp_collector  = n_episodes * timesteps * state_dim * 4   # ~40MB
cuda_overhead  = 200MB                                     # graphs, cuBLAS handles, kernels

remaining = free_vram - replay_vram - exp_collector - cuda_overhead

# Per-sample VRAM (scales with batch_size)
per_sample_mb = (
    state_dim * 2 * 2                           # states + next_states (bf16)
  + state_dim * 4 * 2                           # f32 shadows (states + next)
  + 4 + 4 + 4 + 4                               # actions(i32) + rewards(f32) + dones(f32) + weights(f32)
  + hidden_h1 * 2 + hidden_h2 * 2               # save_h_s1 + save_h_s2 (bf16)
  + value_h * 2 + adv_h * 2 * 3                 # save_h_v + save_h_b0..b2 (bf16)
  + num_atoms * 3 * 2 * 2                        # save_current_lp + save_projected (bf16)
  + 2 + 2                                        # per_sample_loss + td_errors (bf16)
  + total_actions * 2                             # q_out_buf (bf16)
  + iqn_per_sample                                # IQN: quantile embeddings, forward scratch
  + attn_per_sample                               # Attention: 4-head scratch
  + iql_per_sample                                # IQL: forward + loss scratch
) / (1024 * 1024)

max_batch = (remaining / per_sample_mb).floor()
batch_size = max_batch.clamp(min_batch, 16384)
```

## File Structure

### Modified Files

| File | Changes |
|------|---------|
| `crates/ml/src/trainers/dqn/trainer/constructor.rs` | Replace `AutoBatchSizer` call with comprehensive VRAM budget |
| `crates/ml-core/src/memory_optimization/auto_batch_size.rs` | Add `DqnVramBudget` struct with per-component VRAM math |

---

### Task 1: Add DqnVramBudget to auto_batch_size.rs

**Files:**
- Modify: `crates/ml-core/src/memory_optimization/auto_batch_size.rs`

- [ ] **Step 1: Add the DqnVramBudget struct and compute method**

```rust
/// Comprehensive VRAM budget for DQN training pipeline.
/// Accounts for ALL GPU memory consumers, not just the DQN trunk.
pub struct DqnVramBudget {
    pub state_dim: usize,
    pub hidden_h1: usize,
    pub hidden_h2: usize,
    pub value_h: usize,
    pub adv_h: usize,
    pub num_atoms: usize,
    pub total_params: usize,      // DQN trunk params
    pub iqn_params: usize,        // IQN head params (always active)
    pub attn_params: usize,       // Attention params (always active)
    pub iql_params: usize,        // IQL params (always active)
    pub replay_vram_fraction: f64, // fraction of free VRAM for replay buffer
    pub n_episodes: usize,        // experience collector episodes
    pub timesteps: usize,         // experience collector timesteps
}

impl DqnVramBudget {
    /// Compute maximum safe batch size from actual free VRAM.
    pub fn max_batch_size(&self, free_vram_mb: f64) -> usize {
        // 1. Fixed allocations (batch-independent)
        let replay_mb = free_vram_mb * self.replay_vram_fraction;
        let exp_collector_mb = (self.n_episodes * self.timesteps * self.state_dim * 4) as f64
            / (1024.0 * 1024.0);
        // cuBLAS handles, CUDA graphs, kernel modules, segment tree
        let cuda_overhead_mb = 200.0;
        // DQN optimizer: params_buf + target_params_buf + m_buf + v_buf + grad_buf = 5x f32
        let dqn_optim_mb = (self.total_params * 5 * 4) as f64 / (1024.0 * 1024.0);
        // IQN fixed: params + target + m + v + grad = 6x (mixed bf16/f32)
        let iqn_fixed_mb = (self.iqn_params * 6 * 2) as f64 / (1024.0 * 1024.0);
        // Attention fixed: params + m + v + d_params = 4x f32
        let attn_fixed_mb = (self.attn_params * 4 * 4) as f64 / (1024.0 * 1024.0);
        // IQL fixed: params + m + v + grad = 4x f32
        let iql_fixed_mb = (self.iql_params * 4 * 4) as f64 / (1024.0 * 1024.0);

        let fixed_mb = replay_mb + exp_collector_mb + cuda_overhead_mb
            + dqn_optim_mb + iqn_fixed_mb + attn_fixed_mb + iql_fixed_mb;

        let remaining_mb = (free_vram_mb - fixed_mb).max(0.0);

        // 2. Per-sample VRAM (scales with batch_size)
        let sd = self.state_dim;
        let h1 = self.hidden_h1;
        let h2 = self.hidden_h2;
        let vh = self.value_h;
        let ah = self.adv_h;
        let na = self.num_atoms;
        let branches = 3_usize;

        // DQN trainer per-sample bytes:
        let dqn_per_sample = sd * 2 * 2          // states + next_states (bf16, padded)
            + sd * 4 * 2                          // f32 states + next for experience upload
            + 4 + 4 + 4 + 4                       // actions(i32) + rewards(f32) + dones(f32) + is_weights(f32)
            + h1 * 2 + h2 * 2                     // save_h_s1 + save_h_s2 (bf16)
            + vh * 2 + ah * 2 * branches           // save_h_v + save_h_b0..b2 (bf16)
            + na * branches * 2 * 2                // save_current_lp + save_projected (bf16)
            + 2 + 2                                // per_sample_loss + td_errors (bf16)
            + (na * branches + vh + ah * branches) * 2; // forward output logits (bf16)

        // IQN per-sample: quantile embeddings + forward scratch (always active)
        let nq = 64_usize; // num_quantiles
        let iqn_h = h2;    // IQN hidden matches trunk h2
        let iqn_per_sample =
            nq * 64 * 2       // tau embeddings [B, nq, embed_dim] bf16
            + nq * iqn_h * 2  // quantile hidden [B, nq, hidden] bf16
            + nq * na * 2     // quantile Q-values [B, nq, num_atoms] bf16
            + h2 * 4;          // d_h_s2 gradient (f32)

        // Attention per-sample: 4-head scratch (always active)
        let attn_per_sample =
            h2 * 2             // attended output [B, h2] bf16
            + h2 * 4           // d_input scratch (f32)
            + h2 * 2;          // saved_input (bf16)

        // IQL per-sample: V(s) forward + loss (always active)
        let iql_per_sample =
            128 * 2            // IQL hidden (bf16)
            + 2 + 2;           // v_out + loss (bf16)

        let total_per_sample_bytes = dqn_per_sample + iqn_per_sample
            + attn_per_sample + iql_per_sample;
        let per_sample_mb = total_per_sample_bytes as f64 / (1024.0 * 1024.0);

        if per_sample_mb <= 0.0 || remaining_mb <= 0.0 {
            return 64; // absolute minimum
        }

        let max_batch = (remaining_mb / per_sample_mb).floor() as usize;
        max_batch.clamp(64, 16384)
    }
}
```

- [ ] **Step 2: Compile and verify**

```bash
SQLX_OFFLINE=true cargo check -p ml-core
```

Expected: PASS

- [ ] **Step 3: Commit**

```bash
git commit -m "feat: DqnVramBudget — comprehensive VRAM accounting for batch sizing"
```

---

### Task 2: Wire DqnVramBudget into constructor.rs

**Files:**
- Modify: `crates/ml/src/trainers/dqn/trainer/constructor.rs` (use `sed` — Unicode-safe)

Replace the `AutoBatchSizer` call with `DqnVramBudget::max_batch_size()`. The budget struct is populated from the hyperparams and network config that are already available in the constructor.

- [ ] **Step 1: Replace batch sizing logic**

In `constructor.rs`, replace the `AutoBatchSizer::new()` block (lines ~76-99) with:

```rust
// Comprehensive VRAM budget: accounts for ALL GPU consumers.
let max_safe_batch = {
    use ml_core::memory_optimization::auto_batch_size::DqnVramBudget;
    use ml_core::memory_optimization::detect_gpu_hardware;

    match detect_gpu_hardware() {
        Ok(hw) => {
            let budget = DqnVramBudget {
                state_dim: input_dim,
                hidden_h1: hidden_dims[0],
                hidden_h2: *hidden_dims.last().unwrap_or(&128),
                value_h: hidden_dims.last().copied().unwrap_or(64),
                adv_h: hidden_dims.last().copied().unwrap_or(64),
                num_atoms: hyperparams.num_atoms,
                total_params: param_count,
                iqn_params: 20_495,   // always active (from GpuIqnHead init logs)
                attn_params: 263_680, // always active (from GpuAttention init logs)
                iql_params: 27_009,   // always active (from GpuIqlTrainer init logs)
                replay_vram_fraction: hyperparams.replay_buffer_vram_fraction,
                n_episodes: 4096, // worst-case auto-scaled
                timesteps: hyperparams.gpu_timesteps_per_episode,
            };
            let safe = budget.max_batch_size(hw.free_memory_mb);
            info!(
                "DqnVramBudget: batch_size={} (free={:.0}MB, replay={:.0}%, per_sample={:.1}KB)",
                safe, hw.free_memory_mb,
                hyperparams.replay_buffer_vram_fraction * 100.0,
                budget.per_sample_kb(),
            );
            safe
        }
        Err(e) => {
            info!("GPU detection failed ({e}), using batch_size=256");
            256
        }
    }
};

if hyperparams.batch_size == 0 {
    hyperparams.batch_size = max_safe_batch;
    info!("batch_size auto-computed to {}", hyperparams.batch_size);
}
```

Remove the `.min(8192)` hard cap — the budget handles the ceiling.

- [ ] **Step 2: Compile**

```bash
SQLX_OFFLINE=true cargo check -p ml
```

- [ ] **Step 3: Run smoketest locally**

```bash
SQLX_OFFLINE=true FOXHUNT_TEST_DATA=test_data/futures-baseline cargo test -p ml --lib -- smoke_tests::training_stability::test_gpu_collector_auto_initializes --ignored --nocapture 2>&1 | grep -E "batch_size|DqnVramBudget|per-step|Epoch complete|ok|FAIL"
```

Expected:
- RTX 3050 (4GB): batch_size ~128-256 (not 8192)
- Per-step ~7.5ms
- Epoch ~1.8s (not 17s)

- [ ] **Step 4: Commit**

```bash
git commit -m "fix: wire DqnVramBudget into constructor — accounts for IQN/attention/IQL/replay VRAM"
```

---

### Task 3: Validate on H100

- [ ] **Step 1: Submit H100 run**

```bash
./scripts/argo-train.sh dqn --epochs 5 --trials 0 --gpu-pool ci-training-h100
```

Expected:
- H100 (80GB): batch_size ~8192-16384
- n_episodes ~4096
- Per-step <10ms
- Epoch <5s
- No OOM

- [ ] **Step 2: Verify epoch metrics**

Check logs for:
- `DqnVramBudget: batch_size=...` — should be 8192-16384
- `Phase 2: done in ...ms` — experience collection
- `Phase 3: done in ...ms` — training steps
- `Epoch complete` — full epoch timing
- No `CUDA_ERROR_OUT_OF_MEMORY`

---

## Expected Results

| GPU | Free VRAM | Replay (70%) | Remaining | Per-Sample | Max Batch |
|-----|-----------|-------------|-----------|------------|-----------|
| RTX 3050 | 1,700 MB | 1,190 MB | ~310 MB | ~2.5 KB | ~128 |
| A100 40GB | 35,000 MB | 24,500 MB | ~10,300 MB | ~2.5 KB | ~4,096 |
| H100 80GB | 75,000 MB | 52,500 MB | ~22,300 MB | ~2.5 KB | ~8,900 |

Note: these are estimates — actual values depend on IQN/attention config and experience collector scaling.

## Risks

| Risk | Impact | Mitigation |
|------|--------|-----------|
| Per-sample estimate too conservative | Low batch → slow training | Log actual VRAM usage vs estimate, tune |
| Per-sample estimate too aggressive | OOM at training start | Safety margin (20%) + cudarc OOM detection |
| IQN VRAM varies with num_quantiles | Incorrect estimate | Use hyperparams.num_quantiles in budget |
| Replay fraction varies by profile | Budget miscalculation | Read from hyperparams, not hardcoded |
