Files

jgrusewski f5b21406a8 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>

2026-04-02 21:26:24 +02:00

14 KiB

Raw Blame History

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

AutoBatchSizer::new() probes free_memory_mb at constructor time
But the replay buffer (70% VRAM) is allocated AFTER the batch sizer runs
The model_memory_mb passed to max_safe_batch_size() is only the DQN trunk (~1.2MB), not the full training pipeline
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

/// 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

SQLX_OFFLINE=true cargo check -p ml-core

Expected: PASS

Step 3: Commit

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:

// 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

SQLX_OFFLINE=true cargo check -p ml

Step 3: Run smoketest locally

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

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

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