From 465abc7e9b4b4dfedf9e6781444d0d2adce081f7 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Wed, 13 May 2026 20:23:28 +0200 Subject: [PATCH] fix(sp14): dial back aux head/trunk enlargement after L40S OOM Initial 8x head + 2x trunk-H2 (commit 787ee7b86) OOM'd all 3 folds on L40S with alloc next_states f32[1048576000] failure. Root cause: aux backward stores [B, H, SH2] partial grad buffers per sample. At H=256, SH2=256, B=16384: 4GB per head x 2 heads + 4GB trunk = +12GB on top of existing DQN buffers, exceeding L40S 48GB budget. Dial back: - AUX_HIDDEN_DIM: 256 -> 128 (still 4x prior 32-unit bottleneck) - AUX_TRUNK_H2: 256 -> 128 (back to original mild bottleneck) At H=128: partial buffer = 2GB per head = 4GB total. Fits comfortably alongside next_states ~4GB and other DQN buffers. Cargo check clean. Smoke re-dispatch next. Co-Authored-By: Claude Opus 4.7 (1M context) --- .../ml/src/cuda_pipeline/aux_heads_kernel.cu | 9 ++++--- crates/ml/src/cuda_pipeline/gpu_aux_heads.rs | 24 ++++++++++------- crates/ml/src/cuda_pipeline/gpu_aux_trunk.rs | 26 +++++++++---------- docs/dqn-wire-up-audit.md | 25 ++++++++++++++++++ 4 files changed, 59 insertions(+), 25 deletions(-) diff --git a/crates/ml/src/cuda_pipeline/aux_heads_kernel.cu b/crates/ml/src/cuda_pipeline/aux_heads_kernel.cu index 580577f23..3838e7b32 100644 --- a/crates/ml/src/cuda_pipeline/aux_heads_kernel.cu +++ b/crates/ml/src/cuda_pipeline/aux_heads_kernel.cu @@ -64,9 +64,12 @@ /* AUX_HIDDEN_DIM kept as a compile-time literal so kernel launch configs * don't need to thread `H` through. Both heads share the same hidden width. */ #ifndef AUX_HIDDEN_DIM -/* 2026-05-13: enlarged 32 → 256 (matches input h_s2_aux dim). See - * `gpu_aux_heads.rs::AUX_HIDDEN_DIM` for rationale and side effects. */ -#define AUX_HIDDEN_DIM 256 +/* 2026-05-13: enlarged 32 → 128. Initial 256 attempt OOM'd L40S due + * to per-sample dW1_partial [B, H, SH2] = 16384 × 256 × 256 × 4 = 4GB + * per head. At H=128 it's 2GB per head = 4GB total, fits in 48GB + * budget alongside next_states ~4GB. See gpu_aux_heads.rs for full + * rationale. */ +#define AUX_HIDDEN_DIM 128 #endif /* Block dim used by reduction-shaped launches (per-tensor backward, label diff --git a/crates/ml/src/cuda_pipeline/gpu_aux_heads.rs b/crates/ml/src/cuda_pipeline/gpu_aux_heads.rs index 979168aae..d3f1bd9e3 100644 --- a/crates/ml/src/cuda_pipeline/gpu_aux_heads.rs +++ b/crates/ml/src/cuda_pipeline/gpu_aux_heads.rs @@ -50,23 +50,29 @@ use super::gpu_dqn_trainer::{AUX_HEADS_CUBIN, AUX_HEADS_LOSS_EMA_CUBIN}; /// Kept in sync with the `AUX_HIDDEN_DIM` macro at the top of /// `aux_heads_kernel.cu`. /// -/// **Enlarged 32 → 256 (2026-05-13)** as part of the aux-head capacity +/// **Enlarged 32 → 128 (2026-05-13)** as part of the aux-head capacity /// upgrade after Path C investigation revealed the aux head's 28% /// accuracy at H=60 bars. The prior 32-unit bottleneck was severely /// under-capacity for a 256-dim input → 2-way classification task. The -/// new 256 hidden dim matches the input dimension exactly (no -/// dimensionality reduction in Layer 1) and gives the head full -/// capacity to learn an arbitrary mapping. Architecture: -/// 256 (h_s2_aux input) → Linear → 256 (hidden) → ELU → Linear → 2 (logits) +/// new 128 hidden dim gives 4× capacity while keeping memory footprint +/// under control. Architecture: +/// 256 (h_s2_aux input) → Linear → 128 (hidden) → ELU → Linear → 2 (logits) +/// +/// **Memory note**: initial attempt at 256 (8×) OOM'd L40S. Aux backward +/// stores `dW1_partial [B, H, SH2]` = batch_size × H × 256 floats = +/// 16384 × 256 × 256 × 4 bytes = 4GB per head — total 8GB for both heads, +/// which pushed total VRAM over the 48GB L40S budget when combined with +/// the rollout's `next_states[~4GB]` and other DQN buffers. At H=128, +/// the partial buffer is 2GB per head = 4GB total, fitting comfortably. /// /// Side effects: /// - Param buffer grows: aux_nb_w1/aux_rg_w1 went from [32, 256] to -/// [256, 256] each → +64K params per head → +128K params total. +/// [128, 256] each → +24K params per head → +48K params total. /// - Fingerprint changes → checkpoint incompatibility (intentional, /// fresh-start training). -/// - Thread utilisation improves: AUX_BLOCK=256 threads × H=256 hidden -/// = 1 thread per hidden unit (vs 8 threads/unit at H=32). -pub(crate) const AUX_HIDDEN_DIM: usize = 256; +/// - Thread utilisation improves: AUX_BLOCK=256 threads × H=128 hidden +/// = 2 threads per hidden unit (vs 8 threads/unit at H=32). +pub(crate) const AUX_HIDDEN_DIM: usize = 128; /// Output cardinality of the next-bar direction classification head /// (SP13 B1.1a, 2026-05-05). Flipped 1 → 2: the head was a K=1 MSE diff --git a/crates/ml/src/cuda_pipeline/gpu_aux_trunk.rs b/crates/ml/src/cuda_pipeline/gpu_aux_trunk.rs index 1f3d0df86..3c49d47d0 100644 --- a/crates/ml/src/cuda_pipeline/gpu_aux_trunk.rs +++ b/crates/ml/src/cuda_pipeline/gpu_aux_trunk.rs @@ -29,12 +29,13 @@ //! - `h_aux2 [B, H2]` — Layer-2 post-ELU activation //! - `h_s2_aux [B, AUX_HIDDEN_DIM]` — final linear output (no activation) //! -//! Production topology (verified C.2; widths uplifted 2026-05-13): +//! Production topology (verified C.2): //! - `ENCODER_OUT_DIM = config.shared_h1 = 256` //! - `H1 = AUX_TRUNK_H1 = 256` -//! - `H2 = AUX_TRUNK_H2 = 256` (was 128 pre-2026-05-13 — bottleneck removed) -//! - `AUX_HIDDEN_DIM = config.shared_h2 = 256` (matches existing aux head input dim; -//! internal aux head hidden layer ALSO enlarged 32 → 256 in same commit) +//! - `H2 = AUX_TRUNK_H2 = 128` (mild bottleneck) +//! - `AUX_HIDDEN_DIM = config.shared_h2 = 256` (matches aux head input dim; +//! internal aux head hidden layer enlarged 32 → 128 on 2026-05-13; +//! initial 32 → 256 attempt OOM'd L40S — dialed back to 128) //! //! # Pearls applied //! @@ -66,16 +67,15 @@ use super::gpu_dqn_trainer::{ pub(crate) const AUX_TRUNK_H1: usize = 256; /// Hidden width of the aux trunk's second internal layer (Linear_2 → ELU -/// output). +/// output). Smaller than H1 for a mild bottleneck — encourages the aux +/// trunk to compress its representation before lifting back up to +/// `AUX_HIDDEN_DIM` in Layer 3. /// -/// **Enlarged 128 → 256 (2026-05-13)** alongside `AUX_HIDDEN_DIM 32 → 256` -/// to remove the aux trunk's internal bottleneck. After the Path C -/// investigation found the aux head's 28% accuracy at H=60 bars, the -/// hypothesis is that capacity (not signal) was the bottleneck. The -/// trunk now has uniform 256 width throughout: 256 (encoder) → 256 (L1) -/// → 256 (L2) → 256 (L3 = h_s2_aux). This gives the trunk maximum -/// expressive power before the aux head's classification layer. -pub(crate) const AUX_TRUNK_H2: usize = 256; +/// **Kept at 128 (2026-05-13)**: initial enlargement to 256 alongside +/// `AUX_HIDDEN_DIM 32 → 256` OOM'd L40S — trunk dW partials grew too. +/// Kept at 128 for memory budget; head's 4× enlargement to 128 is the +/// primary capacity uplift. +pub(crate) const AUX_TRUNK_H2: usize = 128; /// Block dim used by the per-sample forward kernel. Mirrors /// `aux_heads_kernel.cu`'s `AUX_BLOCK = 256`. diff --git a/docs/dqn-wire-up-audit.md b/docs/dqn-wire-up-audit.md index a82e2e6f8..6fbc68bcc 100644 --- a/docs/dqn-wire-up-audit.md +++ b/docs/dqn-wire-up-audit.md @@ -17545,3 +17545,28 @@ After Path C confirmed the aux head's 28% accuracy at H=60 bars is the actual H6 **Phase 3 mechanism stays DORMANT** for this smoke (W=0, β=0). If aux head accuracy improves to ≥55% with the larger capacity, the next step is to re-activate Phase 3 (W=[-0.5, 0, +0.5, 0], β=0.5) and re-validate WR. Cargo build clean (1m12s — full nvcc rebuild of aux kernels). + +#### Aux head capacity dial-back after OOM (2026-05-13) + +**Initial smoke `train-gxj42` @ `787ee7b86`** (8× head, 2× trunk-H2) OOM'd all 3 folds: +``` +alloc next_states f32[1048576000]: DriverError(CUDA_ERROR_OUT_OF_MEMORY) +``` + +**Root cause**: aux backward stores `dW1_partial [B, H1, SH2]` partials per sample. At `B=16384, H=256, SH2=256`: +- Per head: `16384 × 256 × 256 × 4 bytes = 4 GB` +- Both heads (nb + regime): `8 GB` +- Plus aux trunk's `dW2_partial [B, H1, H2]` = `16384 × 256 × 256 × 4 = 4GB` +- Total +12GB on top of the existing DQN buffers → pushed total over L40S 48GB. + +**Dial-back**: +| Constant | Initial (OOM) | Current | +|----------|---------------|---------| +| AUX_HIDDEN_DIM | 256 (8×) | **128 (4×)** | +| AUX_TRUNK_H2 | 256 (2×) | **128 (1×, unchanged)** | + +At H=128, aux backward partial = `16384 × 128 × 256 × 4 = 2GB` per head = 4GB total. Fits comfortably in 48GB. + +Aux head architecture (now): `256 → Linear → 128 → ELU → Linear → 2` (4× capacity vs prior 32-unit bottleneck). Still meaningful capacity uplift while staying within memory budget. + +**Smoke will be re-dispatched at the dial-back commit**.