From 36d4c39f2dbba8068114ee48fddf082aaf3dfe6c Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Sat, 4 Apr 2026 11:49:13 +0200 Subject: [PATCH] docs: spec + plan for C51 loss accumulator fix (separate MSE/C51 buffers) Co-Authored-By: Claude Opus 4.6 (1M context) --- .../2026-04-04-c51-loss-accumulator-fix.md | 167 ++++++++++++++++++ ...6-04-04-c51-loss-accumulator-fix-design.md | 58 ++++++ 2 files changed, 225 insertions(+) create mode 100644 docs/superpowers/plans/2026-04-04-c51-loss-accumulator-fix.md create mode 100644 docs/superpowers/specs/2026-04-04-c51-loss-accumulator-fix-design.md diff --git a/docs/superpowers/plans/2026-04-04-c51-loss-accumulator-fix.md b/docs/superpowers/plans/2026-04-04-c51-loss-accumulator-fix.md new file mode 100644 index 000000000..bc475dbcd --- /dev/null +++ b/docs/superpowers/plans/2026-04-04-c51-loss-accumulator-fix.md @@ -0,0 +1,167 @@ +# C51 Loss Accumulator Fix + +> **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 NaN at step 22 by separating MSE and C51 loss accumulators. Blend at readback time so C51 NaN during warmup (alpha=0) doesn't contaminate the loss scalar. + +**Architecture:** Add `mse_loss_buf: CudaSlice` alongside existing `total_loss_buf` (now C51-only). MSE kernel writes to `mse_loss_buf`. At readback, blend: `loss = (1-α)*mse + α*c51`. CUDA Graph compatible (no conditional kernel launches). + +**Tech Stack:** CUDA C, Rust (cudarc), f32 atomicAdd + +--- + +## File Map + +| File | Change | +|------|--------| +| `crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs` | Add `mse_loss_buf`, update MSE kernel launch, blend at readback | +| `crates/ml/src/cuda_pipeline/mse_loss_kernel.cu` | Accept separate output buffer (if currently hardcoded to total_loss_buf) | + +--- + +### Task 1: Add mse_loss_buf and separate the accumulators + +**Files:** +- Modify: `crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs` +- Possibly modify: `crates/ml/src/cuda_pipeline/mse_loss_kernel.cu` + +This is a single focused task — all changes are in the loss accumulation path. + +- [ ] **Step 1: Read the current code** + +Read these sections of `gpu_dqn_trainer.rs`: +1. `total_loss_buf` field declaration (line 558) — understand the type and allocation +2. `total_loss_buf` allocation (line 2231) — how it's allocated +3. `GpuDqnPtrs` struct (line 425) — where `total_loss_buf: u64` raw pointer is stored +4. `launch_mse_loss()` (line 4689) — find where it passes `total_loss_buf` to the MSE kernel +5. `launch_c51_loss()` (line 4523) — find where it passes `total_loss_buf` to the C51 kernel +6. The readback path (around line 3030-3040) — where `total_loss_buf` is read back to host +7. The memset-zero before each step — where `total_loss_buf` is zeroed + +Also read `mse_loss_kernel.cu` to understand if the loss buffer is passed as a kernel argument or if it uses the pointer from `GpuDqnPtrs`. + +- [ ] **Step 2: Add `mse_loss_buf` field** + +Add next to `total_loss_buf` (line 558): +```rust + pub(crate) mse_loss_buf: CudaSlice, // [1] MSE loss accumulator (separate from C51) +``` + +If `GpuDqnPtrs` is used to pass pointers to kernels, also add: +```rust + mse_loss_buf: u64, +``` + +- [ ] **Step 3: Allocate `mse_loss_buf`** + +At line 2231, after `total_loss_buf` allocation: +```rust + let mse_loss_buf = stream.alloc_zeros::(1) + .map_err(|e| MLError::ModelError(format!("alloc mse_loss_f32: {e}")))?; +``` + +Wire it into the struct construction. + +- [ ] **Step 4: Update MSE loss kernel launch to use `mse_loss_buf`** + +In `launch_mse_loss()`, find where `total_loss_buf` (or its raw pointer) is passed as an argument to the MSE kernel. Change it to `mse_loss_buf`. + +If the kernel reads the buffer pointer from `GpuDqnPtrs` (a struct uploaded to GPU), you need to set `ptrs.mse_loss_buf = mse_loss_buf.raw_ptr()` instead. + +If the kernel takes it as a direct launch argument, change the `.arg()` call. + +- [ ] **Step 5: Zero both buffers before each step** + +Find where `total_loss_buf` is zeroed (memset) before each training step. Add the same for `mse_loss_buf`: +```rust +self.stream.memset_zeros(&mut self.mse_loss_buf) + .map_err(|e| MLError::ModelError(format!("zero mse_loss: {e}")))?; +``` + +- [ ] **Step 6: Blend at readback** + +Find the loss readback (around line 3030-3040). Currently it reads `total_loss_buf[0]`. Change to read BOTH buffers and blend: + +```rust +// Read both loss accumulators +let mut c51_loss = [0.0_f32; 1]; +let mut mse_loss = [0.0_f32; 1]; +self.stream.memcpy_dtoh(&self.total_loss_buf, &mut c51_loss) + .map_err(|e| MLError::ModelError(format!("readback c51_loss: {e}")))?; +self.stream.memcpy_dtoh(&self.mse_loss_buf, &mut mse_loss) + .map_err(|e| MLError::ModelError(format!("readback mse_loss: {e}")))?; + +// Blend: (1-α)*MSE + α*C51 — when α=0 (warmup), only MSE contributes +let alpha = self.c51_alpha; +let blended_loss = (1.0 - alpha) * mse_loss[0] + alpha * c51_loss[0]; +``` + +Use `blended_loss` instead of the old `total_loss` value. + +**IMPORTANT:** If `alpha * c51_loss[0]` is `0.0 * NaN = NaN`, you need to guard: +```rust +let blended_loss = if alpha < 1e-6 { + mse_loss[0] +} else if alpha > 1.0 - 1e-6 { + c51_loss[0] +} else { + (1.0 - alpha) * mse_loss[0] + alpha * c51_loss[0] +}; +``` + +This prevents `0.0 * NaN = NaN` contamination. + +- [ ] **Step 7: Verify compilation** + +```bash +SQLX_OFFLINE=true cargo check -p ml +``` + +- [ ] **Step 8: Quick NaN test** + +```bash +SQLX_OFFLINE=true FOXHUNT_TEST_DATA=test_data/futures-baseline cargo test -p ml --release --test dqn_training_pipeline_test -- test_dqn_trains_on_es_fut --nocapture --exact 2>&1 | tail -10 +``` + +If step 22 NaN is gone, the fix works. Expected: test runs past step 22 (may still fail for other reasons, but no NaN). + +- [ ] **Step 9: Commit** + +```bash +git add crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs crates/ml/src/cuda_pipeline/mse_loss_kernel.cu +git commit -m "fix: separate MSE/C51 loss accumulators (prevents NaN during warmup)" +``` + +--- + +### Task 2: Run integration tests + +- [ ] **Step 1: Pipeline test (fast NaN check)** + +```bash +SQLX_OFFLINE=true FOXHUNT_TEST_DATA=test_data/futures-baseline cargo test -p ml --release --test dqn_training_pipeline_test -- test_dqn_trains_on_es_fut --nocapture --exact +``` + +Expected: no NaN, completes or fails for non-NaN reason. + +- [ ] **Step 2: Early stopping tests** + +```bash +SQLX_OFFLINE=true cargo test -p ml --release --test dqn_early_stopping_termination_test -- --nocapture +``` + +Expected: all 4 pass, <2 min. + +- [ ] **Step 3: Full test suite** + +```bash +SQLX_OFFLINE=true cargo test -p ml-core --lib -- gpu::profile +SQLX_OFFLINE=true cargo test -p ml --lib -- fxcache feature_cache training_profile +SQLX_OFFLINE=true cargo test -p ml --release --test dqn_action_collapse_fix_test +``` + +- [ ] **Step 4: Push** + +```bash +git push origin main +``` diff --git a/docs/superpowers/specs/2026-04-04-c51-loss-accumulator-fix-design.md b/docs/superpowers/specs/2026-04-04-c51-loss-accumulator-fix-design.md new file mode 100644 index 000000000..871d8f0c3 --- /dev/null +++ b/docs/superpowers/specs/2026-04-04-c51-loss-accumulator-fix-design.md @@ -0,0 +1,58 @@ +# C51 Loss Accumulator Fix + +**Goal:** Fix NaN at step 22 caused by C51 loss kernel contaminating the shared loss accumulator during MSE warmup. Separate loss accumulators for MSE and C51, blend at readback time. + +**Root cause:** Both MSE and C51 loss kernels always run (CUDA Graph compatibility) and both write to the same `total_loss_buf` via `atomicAdd`. During warmup (`c51_alpha=0`), the C51 loss on random/extreme logits produces NaN which contaminates the shared buffer. The gradient blending is correct (C51 grad × 0 = 0), but the loss scalar is irreversibly NaN. + +--- + +## The Fix + +### Separate loss accumulators + +Add a second loss buffer `mse_loss_buf: CudaSlice` to `GpuDqnTrainer`. The existing `total_loss_buf` becomes the C51 accumulator. + +| Buffer | Written by | Purpose | +|--------|-----------|---------| +| `total_loss_buf` | C51 loss kernel | C51 cross-entropy loss | +| `mse_loss_buf` | MSE loss kernel | MSE warmup loss | + +### Blended readback + +At the end of `submit_forward_ops()` (or when the training guard reads the loss), compute: + +``` +blended_loss = (1 - alpha) * mse_loss + alpha * c51_loss +``` + +When `alpha=0` (warmup): `blended_loss = mse_loss` — C51 NaN is ignored. +When `alpha=1` (converged): `blended_loss = c51_loss` — MSE is ignored. + +This keeps both kernels running unconditionally (CUDA Graph compatible) but prevents NaN contamination. + +### Changes + +1. **Add `mse_loss_buf`** to `GpuDqnTrainer` struct (allocate alongside `total_loss_buf`) +2. **Update MSE loss kernel launch** to write to `mse_loss_buf` instead of `total_loss_buf` +3. **Update loss readback** (in training guard or `process_epoch_boundary`) to blend the two buffers using `alpha` +4. **Zero both buffers** before each training step + +### CUDA Graph compatibility + +Both kernels keep running on every step. The conditional logic is ONLY at the scalar readback (host-side blend after GPU→CPU copy). No kernel launches change. Graph stays fixed. + +--- + +## Files touched + +| File | Change | +|------|--------| +| `crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs` | Add `mse_loss_buf`, update MSE kernel launch arg, blend at readback | +| `crates/ml/src/cuda_pipeline/mse_loss_kernel.cu` | May need to accept separate output buffer (check if it already takes a loss_buf arg) | + +## What this does NOT change + +- C51 loss kernel (still runs, still writes to `total_loss_buf`) +- Gradient blending (already correct with scale kernel) +- CUDA Graph structure (no conditional launches) +- The SAXPY/scale kernel fix (separate bug, already fixed)