fix: restore batch_size to GPU profiles — profile defaults, not auto-scaling

Auto-scaling returned 8192 on 4GB RTX 3050 (should be 64), causing
17s/epoch instead of 0.28s. The VRAM math didn't account for IQN (1.1GB),
attention, IQL, replay buffer (70% VRAM).

Profile-tested values:
- RTX 3050: 64 (4GB, minimal)
- A100: 2048 (40-80GB)
- H100: 8192 (80GB, production)
- Default: 256 (conservative)

Auto-scaling remains as fallback for batch_size=0 (unknown GPU).

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

config/gpu/a100.toml

@@ -1,6 +1,6 @@
 # A100 (40-80GB VRAM)
 [training]
-batch_size = 0  # 0 = auto-compute from VRAM
+batch_size = 2048
 num_atoms = 51
 buffer_size = 0  # 0 = auto from VRAM
 hidden_dim_base = 256

config/gpu/default.toml

@@ -1,6 +1,6 @@
 # Default GPU profile -- conservative settings for unknown GPUs
 [training]
-batch_size = 0  # 0 = auto-compute from VRAM
+batch_size = 256
 num_atoms = 21
 buffer_size = 0  # 0 = auto from VRAM
 hidden_dim_base = 256

config/gpu/h100.toml

@@ -1,6 +1,6 @@
 # H100 PCIe/SXM (80GB VRAM) -- full production
 [training]
-batch_size = 0  # 0 = auto-compute from VRAM via AutoBatchSizer
+batch_size = 8192
 num_atoms = 51
 buffer_size = 0  # 0 = auto from VRAM
 hidden_dim_base = 256

config/gpu/rtx3050.toml

@@ -1,6 +1,6 @@
 # RTX 3050 Ti (4GB VRAM) -- minimal footprint
 [training]
-batch_size = 0  # 0 = auto-compute from VRAM
+batch_size = 64
 num_atoms = 11
 buffer_size = 0  # 0 = auto from VRAM
 hidden_dim_base = 256

crates/ml/examples/train_baseline_rl.rs

@@ -506,7 +506,7 @@ fn train_dqn_fold(
 
     let mut hyperparams = DQNHyperparameters {
         learning_rate: hp_f64(hp, "learning_rate").unwrap_or(args.learning_rate),
-        batch_size: 0, // auto from VRAM
+        batch_size: gpu_profile.training.batch_size,
         gamma: hp_f64(hp, "gamma").unwrap_or(0.95),
         epsilon_start,
         epsilon_end: hp_f64(hp, "epsilon_end").unwrap_or(0.01),
@@ -689,7 +689,7 @@ fn train_ppo_fold(
         learning_rate: hp_f64(hp, "learning_rate").unwrap_or(args.learning_rate),
         actor_learning_rate: Some(hp_f64(hp, "policy_learning_rate").unwrap_or(args.learning_rate)),
         critic_learning_rate: Some(hp_f64(hp, "value_learning_rate").unwrap_or(args.learning_rate * 3.0)),
-        batch_size: 0, // auto from VRAM
+        batch_size: ml_core::gpu::profile::GpuProfile::load().training.batch_size,
         gamma: hp_f64(hp, "gamma").unwrap_or(0.99),
         clip_epsilon: hp_f64(hp, "clip_epsilon").unwrap_or(0.2) as f32,
         vf_coef: hp_f64(hp, "value_loss_coeff").unwrap_or(0.5) as f32,

crates/ml/src/trainers/dqn/trainer/constructor.rs

@@ -95,7 +95,7 @@ impl DQNTrainer {
 
         // batch_size == 0 -> auto-compute from VRAM (no artificial cap)
         if hyperparams.batch_size == 0 {
-            hyperparams.batch_size = max_safe_batch.min(8192);
+            hyperparams.batch_size = max_safe_batch.min(8192); // fallback only — profiles should set batch_size
             info!("AutoBatchSizer: batch_size auto-computed to {}", hyperparams.batch_size);
         }
 
@@ -106,7 +106,7 @@ impl DQNTrainer {
                 "DQN batch_size capped from {} -> {} (VRAM ceiling)",
                 hyperparams.batch_size, max_safe_batch
             );
-            hyperparams.batch_size = max_safe_batch.min(8192);
+            hyperparams.batch_size = max_safe_batch.min(8192); // fallback only — profiles should set batch_size
         }
 
         // Use override device if provided (hyperopt shares one CUDA context),

docs/superpowers/plans/2026-04-02-autobatchsizer-vram-budget.md

@@ -0,0 +1,344 @@
+# 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 |