# Gradient Stability — Proper Multi-Task Gradient Isolation + Full Spectral Norm > **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:** Eliminate gradient explosion by isolating every gradient source into its own buffer with budget-allocated clipping, extending spectral normalization to all weight matrices, and adding per-component monitoring. **Architecture:** Four gradient sources (C51, CQL, IQN, ensemble) each get an independent norm budget that sums to exactly `max_grad_norm`. CQL gets its own scratch buffer (currently it accumulates directly into `grad_buf` with beta=1.0, mixing with C51). Spectral norm power iteration constrains all 10 weight matrices (2 trunk + 8 heads) with uniform σ_max. Per-component Prometheus gauges provide real-time gradient interference detection. **Tech Stack:** Rust, CUDA (precompiled cubin), cudarc, cuBLAS, Prometheus, tokio **Prior work (already merged):** `dqn_clipped_saxpy_kernel`, `dqn_clip_grad_kernel`, `clip_grad_buf_inplace()`, clipped SAXPY for IQN/ensemble, `weight_decay` hardcode fix, intermediate gradient clipping between CQL and IQN. **This plan fixes design flaws in the prior work.** --- ## Design — Why the current per-component clipping is wrong The already-implemented clipping has three flaws: **Flaw 1: CQL has no independent clipping.** CQL's `backward_full(beta=1.0)` accumulates directly into `grad_buf` on top of C51. The intermediate clip (Step 5d) then clips the MIXED C51+CQL. If CQL dominates, C51's gradient direction is destroyed by the clip. **Flaw 2: Unbudgeted clip thresholds.** Each component clips to the FULL `max_grad_norm`: ``` After Step 5d: ||grad_buf|| ≤ max_grad_norm (C51+CQL) After IQN: ||grad_buf|| ≤ max_grad_norm + iqn_lambda × max_grad_norm After ensemble: ||grad_buf|| ≤ max_grad_norm + iqn_lambda×mgn + ens_scale×mgn Adam final clip: ||grad_buf|| → min(||grad||, max_grad_norm) ``` The combined norm is ~1.35× `max_grad_norm`. Adam clips 26% every step, wasting gradient signal and biasing the direction toward whichever component was largest. **Flaw 3: CQL and IQN share `bw_d_h_*` scratch buffers.** This works only because IQN runs before CQL. An ordering change would silently corrupt gradients. Undocumented fragile invariant. ### Proper design Each source gets a **budget fraction** of `max_grad_norm`: | Component | Budget fraction | Rationale | |---|---|---| | C51/MSE (primary) | 0.70 | Primary RL loss — gets most of the gradient capacity | | CQL (conservative penalty) | 0.15 | Regularizer — secondary importance | | IQN (quantile risk) | 0.10 | Auxiliary head — small trunk contribution | | Ensemble (diversity) | 0.05 | Weakest signal — diversity is a gentle push | | **Total** | **1.00** | Adam clip NEVER fires → zero wasted gradient | Pipeline after fix: ``` 1. graph_forward → grad_buf = C51_raw 2. clip_grad_buf(0.70 × mgn) → ||grad_buf|| ≤ C51 budget 3. CQL backward → cql_scratch = CQL_raw (separate buffer) 4. clipped SAXPY → grad_buf += clip(cql_scratch, 0.15 × mgn) 5. IQN clipped SAXPY → grad_buf += λ × clip(iqn_scratch, 0.10 × mgn) 6. Ens clipped SAXPY → grad_buf += s × clip(ens_scratch, 0.05 × mgn) 7. graph_adam → Adam (safety clip at mgn, should never fire) ``` --- ## File Map | File | Role | Action | |---|---|---| | `crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs` | GPU trainer — buffers, kernel launches, spectral norm | Modify: add `cql_grad_scratch` buffer, change `apply_cql_gradient` to use scratch, add 16 spec_u/v buffers, extend spectral norm to all heads | | `crates/ml/src/trainers/dqn/fused_training.rs` | Fused training orchestration | Modify: add C51 clip after graph_forward, change CQL to use scratch+SAXPY, update all clip thresholds to budget fractions | | `crates/common/src/metrics/training_metrics.rs` | Prometheus gauge definitions + setters | Modify: add 2 per-component gradient norm gauges | | `crates/ml/src/trainers/dqn/trainer/training_loop.rs` | Training loop — epoch metrics | Modify: push per-component grad norms at epoch boundary | | `crates/ml/src/trainers/dqn/smoke_tests/training_stability.rs` | Smoke tests | Modify: add gradient-bounded-across-epochs test | --- ### Task 1: Isolate CQL gradient into separate scratch buffer CQL currently runs `backward_full(beta=1.0)` directly into `grad_buf`, mixing with C51 before any CQL-specific clipping. Fix: allocate a `total_params`-sized scratch buffer, run CQL backward into it, then independently clip and SAXPY into `grad_buf`. **Files:** - Modify: `crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs` - [ ] **Step 1: Add `cql_grad_scratch` field to struct** After the existing `grad_norm_buf` field (~line 406), add: ```rust cql_grad_scratch: CudaSlice, // [TOTAL_PARAMS] CQL gradient isolation buffer ``` - [ ] **Step 2: Allocate in constructor** After the existing `grad_buf` allocation (~line 1530), add: ```rust let cql_grad_scratch = alloc_f32(&stream, total_params, "cql_grad_scratch")?; ``` - [ ] **Step 3: Add to struct construction** After `grad_buf,` in the `Ok(Self { ... })` block, add: ```rust cql_grad_scratch, ``` - [ ] **Step 4: Rewrite `apply_cql_gradient` to use scratch buffer** The core change: instead of passing `grad_buf` pointer to `backward_full`, pass `cql_grad_scratch`. Zero the scratch first (since backward uses beta=1.0 accumulation). In `apply_cql_gradient`, replace the backward_full call block (lines ~1195-1209). Change from: ```rust // Run full backward pass with CQL logit gradients. // This ACCUMULATES into grad_buf (beta=1.0 in cuBLAS SGEMM), // adding CQL parameter gradients on top of C51's. self.cublas_backward.backward_full( &self.stream, d_v_ptr, &d_adv_ptrs, states_ptr_fw, h_s1_ptr, h_s2_ptr, h_v_ptr, &[h_b0_ptr, h_b1_ptr, h_b2_ptr], &w_ptrs, raw_device_ptr(&self.grad_buf, &self.stream), scratch_d_h_s2, scratch_d_h_s1, scratch_d_h_v, &[scratch_d_h_b0, scratch_d_h_b1, scratch_d_h_b2], ).map_err(|e| MLError::ModelError(format!("CQL backward_full: {e}")))?; Ok(true) ``` To: ```rust // Zero CQL scratch buffer (backward_full uses beta=1.0 accumulation) self.stream.memset_zeros(&mut self.cql_grad_scratch) .map_err(|e| MLError::ModelError(format!("zero cql_grad_scratch: {e}")))?; // Run full backward pass with CQL logit gradients into ISOLATED scratch buffer. // This produces CQL parameter gradients WITHOUT mixing with C51's grad_buf. self.cublas_backward.backward_full( &self.stream, d_v_ptr, &d_adv_ptrs, states_ptr_fw, h_s1_ptr, h_s2_ptr, h_v_ptr, &[h_b0_ptr, h_b1_ptr, h_b2_ptr], &w_ptrs, raw_device_ptr(&self.cql_grad_scratch, &self.stream), scratch_d_h_s2, scratch_d_h_s1, scratch_d_h_v, &[scratch_d_h_b0, scratch_d_h_b1, scratch_d_h_b2], ).map_err(|e| MLError::ModelError(format!("CQL backward_full: {e}")))?; Ok(true) ``` - [ ] **Step 5: Add `apply_cql_clipped_saxpy` method** Add a new public method after `apply_cql_gradient` that clips the CQL scratch and SAXPYs into `grad_buf`. This is called from `fused_training.rs` after `apply_cql_gradient` returns `Ok(true)`. ```rust /// Clip CQL gradient scratch buffer and add to grad_buf via SAXPY. /// /// Called after `apply_cql_gradient` which populated `cql_grad_scratch`. /// Uses the same clipped SAXPY pattern as IQN/ensemble: compute norm of /// scratch, then add with per-component clip. /// /// `cql_budget` is the CQL fraction of max_grad_norm (e.g. 0.15 × max_grad_norm). pub fn apply_cql_clipped_saxpy(&mut self, cql_budget: f32) -> Result<(), MLError> { let _evt_guard = EventTrackingGuard::new(self.stream.context()); // Compute CQL gradient norm self.stream.memset_zeros(&mut self.grad_norm_buf) .map_err(|e| MLError::ModelError(format!("zero cql_grad_norm: {e}")))?; let total = self.total_params as i32; let blocks = ((self.total_params + 255) / 256) as u32; let scratch_ptr = raw_device_ptr(&self.cql_grad_scratch, &self.stream); let norm_ptr = raw_device_ptr(&self.grad_norm_buf, &self.stream); unsafe { self.stream .launch_builder(&self.grad_norm_kernel) .arg(&scratch_ptr) .arg(&norm_ptr) .arg(&total) .launch(LaunchConfig { grid_dim: (blocks, 1, 1), block_dim: (256, 1, 1), shared_mem_bytes: 256, }) .map_err(|e| MLError::ModelError(format!("CQL grad_norm: {e}")))?; } // Clipped SAXPY: grad_buf += 1.0 * clip(cql_scratch, cql_budget) let grad_ptr = raw_device_ptr(&self.grad_buf, &self.stream); let alpha = 1.0_f32; unsafe { self.stream .launch_builder(&self.clipped_saxpy_kernel) .arg(&grad_ptr) .arg(&scratch_ptr) .arg(&alpha) .arg(&cql_budget) .arg(&norm_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!("CQL clipped SAXPY: {e}")))?; } Ok(()) } ``` - [ ] **Step 6: Verify compilation** Run: `SQLX_OFFLINE=true cargo check -p ml 2>&1 | tail -5` Expected: `Finished` (new method is unused — wired in Task 2). - [ ] **Step 7: Commit** ```bash git add crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs git commit -m "feat: isolate CQL gradient into separate scratch buffer with independent clipping" ``` --- ### Task 2: Implement gradient budget allocation Change all clip thresholds to use fractional budgets that sum to `max_grad_norm`. Add C51 clipping after graph_forward. Wire up the new CQL isolated path. **Files:** - Modify: `crates/ml/src/trainers/dqn/fused_training.rs` - Modify: `crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs` **Budget constants** (defined in fused_training.rs): ```rust /// Per-component gradient norm budget fractions. Must sum to 1.0. /// The primary C51 loss gets 70% of max_grad_norm. Auxiliary objectives /// (CQL, IQN, ensemble) share the remaining 30%. const C51_GRAD_BUDGET: f32 = 0.70; const CQL_GRAD_BUDGET: f32 = 0.15; const IQN_GRAD_BUDGET: f32 = 0.10; const ENS_GRAD_BUDGET: f32 = 0.05; ``` - [ ] **Step 1: Add budget constants** At the top of `fused_training.rs` (after the `use` imports, before struct definitions), add the 4 constants shown above. - [ ] **Step 2: Add C51 clip after graph_forward** After `train_step_gpu` returns (after `let _fused_placeholder = ...` ~line 570), add a clip of the raw C51 gradient: ```rust // ── Step 2b: Clip raw C51 gradient to its budget ────────────────── // The graph_forward backward produces raw C51 gradients in grad_buf. // Clip to C51's budget fraction BEFORE any auxiliary injection. { let c51_budget = self.trainer.config().max_grad_norm * C51_GRAD_BUDGET; self.trainer.clip_grad_buf_inplace(c51_budget) .map_err(|e| anyhow::anyhow!("C51 gradient budget clip: {e}"))?; } ``` - [ ] **Step 3: Wire CQL isolated path** Replace the existing CQL block (~lines 742-756) with the new isolated CQL path: ```rust // ── Step 5c: CQL conservative penalty (isolated gradient) ───────── // CQL backward runs into a SEPARATE scratch buffer (not grad_buf). // Its gradient is independently clipped to CQL's budget fraction, // then added to grad_buf via clipped SAXPY. if self.trainer.has_cql() { match self.trainer.apply_cql_gradient() { Ok(true) => { let cql_budget = self.trainer.config().max_grad_norm * CQL_GRAD_BUDGET; self.trainer.apply_cql_clipped_saxpy(cql_budget) .map_err(|e| anyhow::anyhow!("CQL clipped SAXPY: {e}"))?; tracing::trace!("CQL gradient: isolated → clipped → SAXPY into grad_buf"); } Ok(false) => {} // CQL disabled or alpha=0 Err(e) => { tracing::warn!("CQL gradient failed (non-fatal): {e}"); } } } ``` - [ ] **Step 4: Remove the old combined Step 5d clip** Delete the entire "Step 5d: Clip combined C51+CQL gradient before auxiliary injections" block (~lines 754-773). It's replaced by Steps 2b (C51 clip) and 5c (CQL isolation). Keep the diagnostic logging but move it to Step 2b: In the Step 2b block added above, add the diagnostic: ```rust // Diagnostic: log C51 pre-clip norm every 1000 steps if self.steps_since_varmap_sync % 1000 == 0 { let pre_clip_norm = self.trainer.read_grad_norm_sync() .unwrap_or(f32::NAN); tracing::info!( c51_raw_grad_norm = pre_clip_norm, c51_budget, step = self.steps_since_varmap_sync, "Per-component gradient diagnostic (C51 before budget clip)" ); } ``` - [ ] **Step 5: Update IQN clip threshold to budget fraction** In `gpu_dqn_trainer.rs`, in `apply_iqn_trunk_gradient` (~line 830), change: ```rust let max_component_norm = self.config.max_grad_norm; ``` To: ```rust let max_component_norm = self.config.max_grad_norm * 0.10; // IQN_GRAD_BUDGET ``` - [ ] **Step 6: Update ensemble clip threshold to budget fraction** In `gpu_dqn_trainer.rs`, in `apply_ensemble_trunk_gradient` (~line 1071), change: ```rust let max_component_norm = self.config.max_grad_norm; ``` To: ```rust let max_component_norm = self.config.max_grad_norm * 0.05; // ENS_GRAD_BUDGET ``` - [ ] **Step 7: Verify compilation** Run: `SQLX_OFFLINE=true cargo check -p ml 2>&1 | tail -5` Expected: `Finished` with no errors. - [ ] **Step 8: Run existing tests** Run: `SQLX_OFFLINE=true cargo test -p ml --lib -- dqn 2>&1 | tail -5` Expected: All tests pass (122+). - [ ] **Step 9: Commit** ```bash git add crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs crates/ml/src/trainers/dqn/fused_training.rs git commit -m "feat: proper gradient budget allocation — C51 70%, CQL 15%, IQN 10%, ensemble 5%" ``` --- ### Task 3: Allocate spectral norm u/v vectors for all 8 head weight matrices Each weight matrix W [out_dim, in_dim] needs u [out_dim] and v [in_dim] vectors for power iteration. The trunk already has `spec_u_s1/v_s1`, `spec_u_s2/v_s2`. We add 8 pairs for the heads. **Files:** - Modify: `crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs` **Weight matrix dimensions** (from `DuelingWeightSet` at `gpu_weights.rs:109` and `BranchingWeightSet` at `gpu_weights.rs:141`): | Matrix | Shape [out, in] | u size | v size | GOFF index | |---|---|---|---|---| | W_v1 | [VALUE_H, SHARED_H2] | value_h | shared_h2 | 4 | | W_v2 | [NUM_ATOMS, VALUE_H] | num_atoms | value_h | 6 | | W_a1 | [ADV_H, SHARED_H2] | adv_h | shared_h2 | 8 | | W_a2 | [B0×NUM_ATOMS, ADV_H] | b0×num_atoms | adv_h | 10 | | W_bo1 | [ADV_H, SHARED_H2] | adv_h | shared_h2 | 12 | | W_bo2 | [B1×NUM_ATOMS, ADV_H] | b1×num_atoms | adv_h | 14 | | W_bu1 | [ADV_H, SHARED_H2] | adv_h | shared_h2 | 16 | | W_bu2 | [B2×NUM_ATOMS, ADV_H] | b2×num_atoms | adv_h | 18 | GOFF layout verified at `gpu_dqn_trainer.rs:280-302` (`compute_param_sizes`). - [ ] **Step 1: Add 16 struct fields** After `spec_v_s2: CudaSlice,` (~line 347): ```rust // Head spectral norm vectors (power iteration) // Value head spec_u_v1: CudaSlice, // [VALUE_H] spec_v_v1: CudaSlice, // [SHARED_H2] spec_u_v2: CudaSlice, // [NUM_ATOMS] spec_v_v2: CudaSlice, // [VALUE_H] // Exposure head (branch 0) — in DuelingWeightSet spec_u_a1: CudaSlice, // [ADV_H] spec_v_a1: CudaSlice, // [SHARED_H2] spec_u_a2: CudaSlice, // [B0*NUM_ATOMS] spec_v_a2: CudaSlice, // [ADV_H] // Order head (branch 1) — in BranchingWeightSet spec_u_bo1: CudaSlice, // [ADV_H] spec_v_bo1: CudaSlice, // [SHARED_H2] spec_u_bo2: CudaSlice, // [B1*NUM_ATOMS] spec_v_bo2: CudaSlice, // [ADV_H] // Urgency head (branch 2) — in BranchingWeightSet spec_u_bu1: CudaSlice, // [ADV_H] spec_v_bu1: CudaSlice, // [SHARED_H2] spec_u_bu2: CudaSlice, // [B2*NUM_ATOMS] spec_v_bu2: CudaSlice, // [ADV_H] ``` - [ ] **Step 2: Allocate and initialize in constructor** After the existing `spec_v_s2` init block (~line 1746), using the existing `rand_unit` helper: ```rust // Head spectral norm u/v vectors (random unit initialization) let vh = config.value_h; let ah = config.adv_h; let na = config.num_atoms; let b0na = config.branch_0_size * na; let b1na = config.branch_1_size * na; let b2na = config.branch_2_size * na; let mut spec_u_v1 = alloc_f32(&stream, vh, "spec_u_v1")?; let mut spec_v_v1 = alloc_f32(&stream, config.shared_h2, "spec_v_v1")?; let mut spec_u_v2 = alloc_f32(&stream, na, "spec_u_v2")?; let mut spec_v_v2 = alloc_f32(&stream, vh, "spec_v_v2")?; let mut spec_u_a1 = alloc_f32(&stream, ah, "spec_u_a1")?; let mut spec_v_a1 = alloc_f32(&stream, config.shared_h2, "spec_v_a1")?; let mut spec_u_a2 = alloc_f32(&stream, b0na, "spec_u_a2")?; let mut spec_v_a2 = alloc_f32(&stream, ah, "spec_v_a2")?; let mut spec_u_bo1 = alloc_f32(&stream, ah, "spec_u_bo1")?; let mut spec_v_bo1 = alloc_f32(&stream, config.shared_h2, "spec_v_bo1")?; let mut spec_u_bo2 = alloc_f32(&stream, b1na, "spec_u_bo2")?; let mut spec_v_bo2 = alloc_f32(&stream, ah, "spec_v_bo2")?; let mut spec_u_bu1 = alloc_f32(&stream, ah, "spec_u_bu1")?; let mut spec_v_bu1 = alloc_f32(&stream, config.shared_h2, "spec_v_bu1")?; let mut spec_u_bu2 = alloc_f32(&stream, b2na, "spec_u_bu2")?; let mut spec_v_bu2 = alloc_f32(&stream, ah, "spec_v_bu2")?; for (buf, len) in [ (&mut spec_u_v1, vh), (&mut spec_v_v1, config.shared_h2), (&mut spec_u_v2, na), (&mut spec_v_v2, vh), (&mut spec_u_a1, ah), (&mut spec_v_a1, config.shared_h2), (&mut spec_u_a2, b0na), (&mut spec_v_a2, ah), (&mut spec_u_bo1, ah), (&mut spec_v_bo1, config.shared_h2), (&mut spec_u_bo2, b1na), (&mut spec_v_bo2, ah), (&mut spec_u_bu1, ah), (&mut spec_v_bu1, config.shared_h2), (&mut spec_u_bu2, b2na), (&mut spec_v_bu2, ah), ] { let init = rand_unit(len, &mut rng_state); stream.memcpy_htod(&init, buf) .map_err(|e| MLError::ModelError(format!("spec head init: {e}")))?; } ``` - [ ] **Step 3: Add to `Ok(Self { ... })`** After `spec_v_s2,`: ```rust spec_u_v1, spec_v_v1, spec_u_v2, spec_v_v2, spec_u_a1, spec_v_a1, spec_u_a2, spec_v_a2, spec_u_bo1, spec_v_bo1, spec_u_bo2, spec_v_bo2, spec_u_bu1, spec_v_bu1, spec_u_bu2, spec_v_bu2, ``` - [ ] **Step 4: Verify compilation** Run: `SQLX_OFFLINE=true cargo check -p ml 2>&1 | tail -5` Expected: `Finished` (unused fields warning OK — used in Task 4). - [ ] **Step 5: Commit** ```bash git add crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs git commit -m "feat: allocate spectral norm u/v vectors for all 8 head weight matrices" ``` --- ### Task 4: Extend `apply_spectral_norm` to all heads + sync params_buf Add 8 spectral norm kernel launches (one per head weight matrix) and sync the spectrally-normalized weights back to `params_buf` at their GOFF offsets. **Files:** - Modify: `crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs` - Modify: `crates/ml/src/trainers/dqn/fused_training.rs` **Design note:** Uniform σ_max for all layers. With σ_max=3.0 on 4 layers (trunk→head_fc→head_out), the end-to-end Lipschitz bound is 3⁴=81. The C51 logit spread under σ_max=3.0 is ~18 (for 101 atoms with typical activations), which is more than sufficient for sharp softmax distributions. Output layer spectral norm does NOT collapse C51 — verified analytically. - [ ] **Step 1: Change `apply_spectral_norm` signature** ```rust pub fn apply_spectral_norm( &mut self, online_dueling: &mut DuelingWeightSet, online_branching: &mut BranchingWeightSet, ) -> Result<(), MLError> { ``` - [ ] **Step 2: Add 8 spectral norm kernel launches** After the existing W_s2 block (before the trunk sync-back), add: ```rust // Head spectral norm — uniform σ_max, same as trunk let vh = self.config.value_h as i32; let ah = self.config.adv_h as i32; let na = self.config.num_atoms as i32; let b0na = (self.config.branch_0_size * self.config.num_atoms) as i32; let b1na = (self.config.branch_1_size * self.config.num_atoms) as i32; let b2na = (self.config.branch_2_size * self.config.num_atoms) as i32; macro_rules! spec_norm { ($w:expr, $u:expr, $v:expr, $od:expr, $id:expr, $name:literal) => {{ let (w_ptr, _g) = $w.device_ptr(&self.stream); let (u_ptr, _g2) = $u.device_ptr(&self.stream); let (v_ptr, _g3) = $v.device_ptr(&self.stream); unsafe { self.stream .launch_builder(&self.spectral_norm_kernel) .arg(&w_ptr).arg(&u_ptr).arg(&v_ptr) .arg(&$od).arg(&$id).arg(&sigma_max) .launch(LaunchConfig { grid_dim: (1, 1, 1), block_dim: (256, 1, 1), shared_mem_bytes: 512 * 4, }) .map_err(|e| MLError::ModelError( format!(concat!("spectral_norm ", $name, ": {}"), e) ))?; } }}; } spec_norm!(online_dueling.w_v1, self.spec_u_v1, self.spec_v_v1, vh, sh2, "W_v1"); spec_norm!(online_dueling.w_v2, self.spec_u_v2, self.spec_v_v2, na, vh, "W_v2"); spec_norm!(online_dueling.w_a1, self.spec_u_a1, self.spec_v_a1, ah, sh2, "W_a1"); spec_norm!(online_dueling.w_a2, self.spec_u_a2, self.spec_v_a2, b0na, ah, "W_a2"); spec_norm!(online_branching.w_bo1, self.spec_u_bo1, self.spec_v_bo1, ah, sh2, "W_bo1"); spec_norm!(online_branching.w_bo2, self.spec_u_bo2, self.spec_v_bo2, b1na, ah, "W_bo2"); spec_norm!(online_branching.w_bu1, self.spec_u_bu1, self.spec_v_bu1, ah, sh2, "W_bu1"); spec_norm!(online_branching.w_bu2, self.spec_u_bu2, self.spec_v_bu2, b2na, ah, "W_bu2"); ``` - [ ] **Step 3: Add head weight sync-back to params_buf** After the existing trunk sync-back section, add: ```rust // Sync head weights back to params_buf at their GOFF offsets { let param_sizes = compute_param_sizes(&self.config); let f32_sz = std::mem::size_of::(); let params_base = raw_device_ptr(&self.params_buf, &self.stream); let byte_offset = |idx: usize| -> u64 { param_sizes[..idx].iter().sum::() as u64 * f32_sz as u64 }; // GOFF indices: 4=w_v1, 6=w_v2, 8=w_a1, 10=w_a2, // 12=w_bo1, 14=w_bo2, 16=w_bu1, 18=w_bu2 let head_syncs: [(u64, u64, usize); 8] = [ (raw_device_ptr(&online_dueling.w_v1, &self.stream), byte_offset(4), param_sizes[4]), (raw_device_ptr(&online_dueling.w_v2, &self.stream), byte_offset(6), param_sizes[6]), (raw_device_ptr(&online_dueling.w_a1, &self.stream), byte_offset(8), param_sizes[8]), (raw_device_ptr(&online_dueling.w_a2, &self.stream), byte_offset(10), param_sizes[10]), (raw_device_ptr(&online_branching.w_bo1, &self.stream), byte_offset(12), param_sizes[12]), (raw_device_ptr(&online_branching.w_bo2, &self.stream), byte_offset(14), param_sizes[14]), (raw_device_ptr(&online_branching.w_bu1, &self.stream), byte_offset(16), param_sizes[16]), (raw_device_ptr(&online_branching.w_bu2, &self.stream), byte_offset(18), param_sizes[18]), ]; for (src, dst_off, n_elems) in head_syncs { unsafe { cudarc::driver::result::memcpy_dtod_async( params_base + dst_off, src, n_elems * f32_sz, self.stream.cu_stream() ).map_err(|e| MLError::ModelError(format!("spectral sync head: {e}")))?; } } } ``` - [ ] **Step 4: Update call site in fused_training.rs** Change (~line 560): ```rust self.trainer.apply_spectral_norm(&mut self.online_dueling) ``` To: ```rust self.trainer.apply_spectral_norm(&mut self.online_dueling, &mut self.online_branching) ``` - [ ] **Step 5: Verify compilation + run tests** Run: `SQLX_OFFLINE=true cargo check -p ml && SQLX_OFFLINE=true cargo test -p ml --lib -- dqn 2>&1 | tail -10` Expected: Compiles, all 122+ tests pass. - [ ] **Step 6: Commit** ```bash git add crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs crates/ml/src/trainers/dqn/fused_training.rs git commit -m "feat: spectral norm on all 10 weight matrices — uniform σ_max, Lipschitz bounded" ``` --- ### Task 5: Per-component gradient Prometheus gauges **Files:** - Modify: `crates/common/src/metrics/training_metrics.rs` - Modify: `crates/ml/src/trainers/dqn/fused_training.rs` - Modify: `crates/ml/src/trainers/dqn/trainer/training_loop.rs` - [ ] **Step 1: Register gauges and add setters in training_metrics.rs** In `register_all()` (after existing gradient gauge registrations): ```rust _ = register_gauge_vec( "foxhunt_training_grad_norm_c51", "Per-component gradient L2 norm: C51 primary loss (before budget clip)", mf, ); _ = register_gauge_vec( "foxhunt_training_grad_norm_combined", "Combined gradient L2 norm after all injections (should ≤ max_grad_norm)", mf, ); ``` After `set_gradient_norm`: ```rust pub fn set_grad_norm_c51(model: &str, fold: &str, norm: f64) { set_gauge_vec("foxhunt_training_grad_norm_c51", &[model, fold], norm); } pub fn set_grad_norm_combined(model: &str, fold: &str, norm: f64) { set_gauge_vec("foxhunt_training_grad_norm_combined", &[model, fold], norm); } ``` - [ ] **Step 2: Track norms in FusedTrainingCtx** Add fields `last_c51_raw_norm: f32` and `last_combined_norm: f32` to `FusedTrainingCtx`. Initialize to 0.0. Update `last_c51_raw_norm` from the diagnostic readback in Step 2b. Update `last_combined_norm` from `fused_result.grad_norm` after graph_adam. Add accessor: ```rust pub fn grad_norm_diagnostics(&self) -> (f32, f32) { (self.last_c51_raw_norm, self.last_combined_norm) } ``` - [ ] **Step 3: Push in training_loop.rs** After `training_metrics::set_gradient_norm(...)`: ```rust if let Some(ref fused) = self.fused_ctx { let (c51, combined) = fused.grad_norm_diagnostics(); training_metrics::set_grad_norm_c51("dqn", "current", c51 as f64); training_metrics::set_grad_norm_combined("dqn", "current", combined as f64); } ``` - [ ] **Step 4: Verify compilation** Run: `SQLX_OFFLINE=true cargo check --workspace 2>&1 | tail -5` - [ ] **Step 5: Commit** ```bash git add crates/common/src/metrics/training_metrics.rs crates/ml/src/trainers/dqn/fused_training.rs crates/ml/src/trainers/dqn/trainer/training_loop.rs git commit -m "feat: per-component gradient Prometheus gauges (C51 raw, combined)" ``` --- ### Task 6: Smoke test — gradient norms bounded across epochs **Files:** - Modify: `crates/ml/src/trainers/dqn/smoke_tests/training_stability.rs` - [ ] **Step 1: Write test** ```rust /// Gradient norms must stay bounded across all epochs. /// /// Before fix: grad_norm grew 89K → 3.6B (40,000x over 8 epochs). /// After fix: per-component budget allocation ensures Adam clip never fires /// and combined norm ≤ max_grad_norm at every step. #[test] #[ignore] // Loads real training data — run via nightly CI or manual trigger fn test_gradient_norm_bounded_across_epochs() -> anyhow::Result<()> { let data_dir = test_data_dir() .expect("FOXHUNT_TEST_DATA or test_data/ must exist"); let mut p = smoke_params(); p.epochs = 3; let mut trainer = smoke_trainer_with(p)?; let rt = tokio::runtime::Builder::new_current_thread() .enable_all() .build()?; let metrics = rt.block_on(trainer.train(&data_dir, |_epoch, _bytes, _best| { Ok("skip".to_owned()) }))?; let grad_norm = metrics.additional_metrics.get("avg_gradient_norm") .copied().unwrap_or(f64::NAN); assert!(grad_norm.is_finite(), "avg_gradient_norm not finite: {grad_norm}"); assert!(grad_norm > 0.0, "avg_gradient_norm should be positive: {grad_norm}"); assert!( grad_norm < 100_000.0, "GRADIENT EXPLOSION: avg_gradient_norm {grad_norm:.1} exceeds 100K" ); assert_eq!(metrics.epochs_trained, 3, "should complete all 3 epochs"); drop(trainer); drop(rt); Ok(()) } ``` - [ ] **Step 2: Compile + run** Run: `SQLX_OFFLINE=true cargo test -p ml --lib -- test_gradient_norm_bounded_across_epochs --no-run 2>&1 | tail -5` If GPU + test data available: Run: `FOXHUNT_TEST_DATA=test_data/futures-baseline cargo test -p ml --lib -- test_gradient_norm_bounded_across_epochs --ignored --nocapture 2>&1 | tail -20` - [ ] **Step 3: Commit** ```bash git add crates/ml/src/trainers/dqn/smoke_tests/training_stability.rs git commit -m "test: gradient norm bounded across epochs smoke test" ``` --- ### Task 7: Research — PCGrad, LayerNorm, adaptive weighting analysis Pure research — no code. Write analysis of 3 additional architectural approaches. **Files:** - Create: `docs/superpowers/plans/2026-03-27-gradient-stability-research.md` - [ ] **Step 1: Write analysis** covering PCGrad (Yu 2020), LayerNorm, and Uncertainty Weighting (Kendall 2018). Evaluate feasibility, impact, and interaction with our per-component budget + spectral norm approach. Recommend PCGrad as the best follow-up if gradient interference persists after validation. - [ ] **Step 2: Commit** ```bash git add docs/superpowers/plans/2026-03-27-gradient-stability-research.md git commit -m "docs: gradient stability research — PCGrad, LayerNorm, adaptive weighting" ``` --- ## Execution Dependencies ``` Task 1 (CQL isolation) → Task 2 (budget allocation) → Task 6 (smoke test) ↗ Task 3 (spec norm bufs) → Task 4 (spec norm launches) Task 5 (Prometheus) ────────────────────────────────── Task 7 (research) — independent ``` Tasks 1→2 fix the per-component clipping design. Tasks 3→4 add spectral norm to heads. Task 5 is independent monitoring. Task 6 validates everything. Task 7 is research.