From f14440be5566c31988d850f87c4b8c55b08bfef4 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Wed, 15 Apr 2026 08:40:09 +0200 Subject: [PATCH] =?UTF-8?q?plan:=20add=20Tasks=2010-11=20=E2=80=94=20Q-mea?= =?UTF-8?q?n=20centering=20+=20atom=20entropy=20recovery=20(execute=20firs?= =?UTF-8?q?t)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Task 10: q_mean_reduce + q_mean_subtract two-phase kernel. Zero params. Fixes Q-mean drift 0→+0.47 from bootstrapping bias. Task 11: Adaptive entropy_coeff based on utilization_ema. Zero params. Fixes atom utilization collapse 100%→20%. Execution order: 10, 11, then 1-9. Co-Authored-By: Claude Opus 4.6 (1M context) --- ...-04-15-supervised-architecture-transfer.md | 213 +++++++++++++++++- 1 file changed, 211 insertions(+), 2 deletions(-) diff --git a/docs/superpowers/plans/2026-04-15-supervised-architecture-transfer.md b/docs/superpowers/plans/2026-04-15-supervised-architecture-transfer.md index 88e58e842..af613b77f 100644 --- a/docs/superpowers/plans/2026-04-15-supervised-architecture-transfer.md +++ b/docs/superpowers/plans/2026-04-15-supervised-architecture-transfer.md @@ -2,9 +2,9 @@ > **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:** Transfer 7 architectural concepts from the supervised model suite (KAN, xLSTM, Diffusion, TGNN, Liquid ODE, TFT Quantile, TLOB) into the DQN training pipeline to improve action selection quality, representational power, and training dynamics. +**Goal:** Transfer 9 architectural concepts (7 from supervised suite + 2 stability fixes from live H100 observations) into the DQN training pipeline. -**Architecture:** 7 independent components that compose through the existing forward/backward/training pipeline. Each borrows a specific concept from a supervised architecture already implemented in `crates/ml-supervised/`. Components are deployed incrementally: quantile Q-select first (zero params, highest OOS impact), then graph message passing, KAN gates, diffusion refinement, TLOB injection, xLSTM context, and RK4 ODE last. +**Architecture:** 9 components. Tasks 10-11 (Q-mean centering + atom entropy recovery) should be executed FIRST — they fix active instabilities observed in live H100 run train-vpb4w (Q-mean drift 0→+0.47, atom utilization 100%→20%). Then the 7 supervised transfers deploy incrementally. **Tech Stack:** Rust 1.85+, CUDA 12.4 (NVRTC precompiled cubins via build.rs), cuBLAS cublasLt TF32, cudarc driver bindings. All kernels in `experience_kernels.cu`, all Rust integration in the `crates/ml/src/cuda_pipeline/` and `crates/ml/src/trainers/dqn/` modules. @@ -1718,6 +1718,215 @@ Co-Authored-By: Claude Opus 4.6 (1M context) " --- +## Task 10: Q-Mean Drift Correction (Component 8) — EXECUTE FIRST + +**Priority:** Execute BEFORE Tasks 1-9. Fixes active instability in live H100 run. + +**Files:** +- `crates/ml/src/cuda_pipeline/experience_kernels.cu` +- `crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs` + +- [ ] **Step 1:** Add `q_mean_center` kernel to `experience_kernels.cu`: + +```cuda +/** + * Zero-mean Q-values: subtract global mean from all Q-values. + * Prevents bootstrapping drift from accumulating across epochs. + * Preserves action ranking (constant subtraction doesn't change argmax). + * + * Phase 1 (block 0): compute mean of Q_raw[B*total_actions] + * Phase 2 (all blocks): Q_centered[i] = Q_raw[i] - mean + * + * Grid: ceil(B*total_actions/256), Block: 256. + */ +extern "C" __global__ void q_mean_center( + float* __restrict__ q_values, /* [B, total_actions] in-place */ + float* __restrict__ mean_scratch, /* [1] scratch for global mean */ + int total /* B * total_actions */ +) { + /* Phase 1: block 0 computes mean */ + if (blockIdx.x == 0) { + __shared__ float sdata[256]; + float sum = 0.0f; + for (int i = threadIdx.x; i < total; i += blockDim.x) { + sum += q_values[i]; + } + sdata[threadIdx.x] = sum; + __syncthreads(); + for (int s = blockDim.x / 2; s > 0; s >>= 1) { + if (threadIdx.x < s) sdata[threadIdx.x] += sdata[threadIdx.x + s]; + __syncthreads(); + } + if (threadIdx.x == 0) { + mean_scratch[0] = sdata[0] / (float)total; + } + } + /* Global barrier via grid sync not available — use two-phase launch */ +} +``` + +Actually, this needs a two-phase approach (can't barrier across blocks). Use two kernel launches: + +```cuda +/** + * Phase 1: Compute mean of Q-values into a scratch scalar. + * Grid: (1, 1, 1), Block: (256, 1, 1). + */ +extern "C" __global__ void q_mean_reduce( + const float* __restrict__ q_values, /* [total] */ + float* __restrict__ mean_out, /* [1] output */ + int total) +{ + __shared__ float sdata[256]; + float sum = 0.0f; + for (int i = threadIdx.x; i < total; i += 256) { + sum += q_values[i]; + } + sdata[threadIdx.x] = sum; + __syncthreads(); + for (int s = 128; s > 0; s >>= 1) { + if (threadIdx.x < s) sdata[threadIdx.x] += sdata[threadIdx.x + s]; + __syncthreads(); + } + if (threadIdx.x == 0) mean_out[0] = sdata[0] / (float)total; +} + +/** + * Phase 2: Subtract mean from all Q-values. + * Grid: ceil(total/256), Block: 256. + */ +extern "C" __global__ void q_mean_subtract( + float* __restrict__ q_values, /* [total] in-place */ + const float* __restrict__ mean_val, /* [1] */ + int total) +{ + int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i >= total) return; + q_values[i] -= mean_val[0]; +} +``` + +- [ ] **Step 2:** Load kernels in `gpu_dqn_trainer.rs` constructor. Add fields `q_mean_reduce_kernel`, `q_mean_subtract_kernel`, `q_mean_scratch: CudaSlice` (1 float). + +- [ ] **Step 3:** Add `launch_q_mean_center` method: + +```rust +pub(crate) fn launch_q_mean_center(&self, batch_size: usize) -> Result<(), MLError> { + let total = (batch_size * self.total_actions()) as i32; + let q_ptr = self.q_out_buf.raw_ptr(); + let scratch_ptr = self.q_mean_scratch.raw_ptr(); + // Phase 1: reduce to mean + unsafe { + self.stream.launch_builder(&self.q_mean_reduce_kernel) + .arg(&q_ptr).arg(&scratch_ptr).arg(&total) + .launch(LaunchConfig { grid_dim: (1,1,1), block_dim: (256,1,1), shared_mem_bytes: 0 }) + .map_err(|e| MLError::ModelError(format!("q_mean_reduce: {e}")))?; + } + // Phase 2: subtract mean + let blocks = ((total as u32 + 255) / 256).max(1); + unsafe { + self.stream.launch_builder(&self.q_mean_subtract_kernel) + .arg(&q_ptr).arg(&scratch_ptr).arg(&total) + .launch(LaunchConfig { grid_dim: (blocks,1,1), block_dim: (256,1,1), shared_mem_bytes: 0 }) + .map_err(|e| MLError::ModelError(format!("q_mean_subtract: {e}")))?; + } + Ok(()) +} +``` + +- [ ] **Step 4:** Wire into Q-value pipeline. In `reduce_current_q_stats`, AFTER `compute_expected_q` and BEFORE `launch_q_attention`: + +```rust +self.launch_q_mean_center(batch_size)?; +``` + +- [ ] **Step 5:** Add passthrough in `fused_training.rs`. + +- [ ] **Step 6:** Verify build and commit: +```bash +SQLX_OFFLINE=true cargo check -p ml 2>&1 | tail -5 +git add crates/ml/src/cuda_pipeline/experience_kernels.cu \ + crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs \ + crates/ml/src/trainers/dqn/fused_training.rs +git commit -m "feat: Q-mean centering — prevents bootstrapping drift (Component 8) + +Two-phase kernel: q_mean_reduce computes global mean, q_mean_subtract +centers all Q-values. Runs after compute_expected_q, before Q-attention. +Fixes Q-mean drift 0→+0.47 observed in train-vpb4w. + +Co-Authored-By: Claude Opus 4.6 (1M context) " +``` + +--- + +## Task 11: Atom Utilization Recovery (Component 9) — EXECUTE SECOND + +**Priority:** Execute AFTER Task 10, BEFORE Tasks 1-9. Fixes atom collapse in live H100 run. + +**Files:** +- `crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs` +- `crates/ml/src/trainers/dqn/fused_training.rs` + +- [ ] **Step 1:** Add `utilization_ema` field to `GpuDqnTrainer`: + +```rust +/// EMA of atom utilization [0,1] for adaptive entropy regularization. +utilization_ema: f32, +``` + +Initialize to `1.0` in the constructor (fully utilized at start). + +- [ ] **Step 2:** In `reduce_current_q_stats`, after reading `atom_utilization` from the Q-stats readback, update the EMA: + +```rust +// Adaptive-rate EMA for atom utilization +let util = stats.atom_utilization; // already in QValueStatsResult +if self.utilization_ema > 0.99 { + self.utilization_ema = util; +} else { + let err = (util - self.utilization_ema).abs(); + let alpha = (err / (err + 0.1)).clamp(0.01, 0.3); + self.utilization_ema = (1.0 - alpha) * self.utilization_ema + alpha * util; +} +``` + +- [ ] **Step 3:** Compute adaptive entropy coefficient. In the method that launches `c51_grad_kernel` (or wherever `entropy_coeff` is prepared), replace the fixed hyperparameter: + +```rust +// Adaptive entropy: ramp up when atom utilization drops below 50% +let base_entropy = self.config.entropy_coefficient; +let adaptive_entropy = if self.utilization_ema < 0.5 { + base_entropy * (1.0 - self.utilization_ema / 0.5) // linear ramp: 0 at 50%, base at 0% +} else { + 0.0 // no entropy bonus when atoms are healthy +}; +let entropy_coeff = base_entropy + adaptive_entropy; +``` + +Find where `entropy_coeff` is passed to `launch_c51_grad` and replace the fixed value with the adaptive one. + +- [ ] **Step 4:** Add accessor in `fused_training.rs`: + +```rust +pub(crate) fn utilization_ema(&self) -> f32 { self.trainer.utilization_ema } +``` + +- [ ] **Step 5:** Verify build and commit: +```bash +SQLX_OFFLINE=true cargo check -p ml 2>&1 | tail -5 +git add crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs \ + crates/ml/src/trainers/dqn/fused_training.rs +git commit -m "feat: adaptive atom entropy — recovers utilization when C51 distribution collapses (Component 9) + +entropy_coeff ramps up when atom utilization EMA drops below 50%. +Zero at 50%+, linear ramp to base_entropy at 0%. Prevents distributional +collapse observed in train-vpb4w (100%→20% over 18 epochs). + +Co-Authored-By: Claude Opus 4.6 (1M context) " +``` + +--- + ## Task 9: Build Verification + Smoke Tests **Files:**