diff --git a/config/gpu/a100.toml b/config/gpu/a100.toml index 1cd527d16..cf56cf27e 100644 --- a/config/gpu/a100.toml +++ b/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 diff --git a/config/gpu/default.toml b/config/gpu/default.toml index cac4d8e09..755f04379 100644 --- a/config/gpu/default.toml +++ b/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 diff --git a/config/gpu/h100.toml b/config/gpu/h100.toml index b83339af3..bc4f4dcea 100644 --- a/config/gpu/h100.toml +++ b/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 diff --git a/config/gpu/rtx3050.toml b/config/gpu/rtx3050.toml index 88d88c663..1c810abef 100644 --- a/config/gpu/rtx3050.toml +++ b/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 diff --git a/crates/ml/examples/train_baseline_rl.rs b/crates/ml/examples/train_baseline_rl.rs index d48a8db9d..abd6358b3 100644 --- a/crates/ml/examples/train_baseline_rl.rs +++ b/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, diff --git a/crates/ml/src/trainers/dqn/trainer/constructor.rs b/crates/ml/src/trainers/dqn/trainer/constructor.rs index 2b035187f..bf4da312a 100644 --- a/crates/ml/src/trainers/dqn/trainer/constructor.rs +++ b/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), diff --git a/docs/superpowers/plans/2026-04-02-autobatchsizer-vram-budget.md b/docs/superpowers/plans/2026-04-02-autobatchsizer-vram-budget.md new file mode 100644 index 000000000..e95c5187a --- /dev/null +++ b/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 |