From 0de7f694186f33d80ed01bc5a2ad104467ffe71a Mon Sep 17 00:00:00 2001
From: jgrusewski <jgrusewski@localhost>
Date: Thu, 9 Apr 2026 00:09:17 +0200
Subject: [PATCH] fix(training): gradient cliff at C51 warmup + stale val_loss
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Two issues found in the H100 production run:

1. Primary gradient budget dropped from 1.0 to 0.10 when C51 alpha
   ramped to 1.0 (c51_frac=0.10 with IQN at 60%). Branch heads
   depend entirely on primary gradient — starved → grad_norm=0.000
   at epoch 7. Fix: floor c51_frac at 0.30 so branch heads always
   get meaningful gradient.

2. val_loss identical at epochs 1-2 (6.228316) because the backtest
   evaluator's CUDA Graph was never invalidated between epochs.
   Fix: call invalidate_dqn_graph() before each evaluation.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
---
 crates/ml/src/trainers/dqn/fused_training.rs  | 13 +++++++++----
 crates/ml/src/trainers/dqn/trainer/metrics.rs |  6 ++++++
 2 files changed, 15 insertions(+), 4 deletions(-)

diff --git a/crates/ml/src/trainers/dqn/fused_training.rs b/crates/ml/src/trainers/dqn/fused_training.rs
index 64bcad381..5e155fc8d 100644
--- a/crates/ml/src/trainers/dqn/fused_training.rs
+++ b/crates/ml/src/trainers/dqn/fused_training.rs
@@ -907,15 +907,20 @@ impl FusedTrainingCtx {
             ).map_err(|e| anyhow::anyhow!("HER in-place relabel kernel: {e}"))?;
         }
 
-        // Primary gradient budget clip: blends from full max_grad_norm during
-        // MSE warmup to the C51 fraction as c51_alpha ramps up. Prevents
-        // artificially suppressing MSE gradients when C51 isn't contributing.
+        // Primary gradient budget clip: C51+MSE share of max_grad_norm.
+        // With IQN as primary (60%), IQN contributes only via trunk gradient (SAXPY).
+        // Branch heads depend ENTIRELY on the primary (C51+MSE) gradient — if the
+        // primary budget is too small, branch heads starve and gradients collapse.
+        //
+        // Floor at max(c51_frac, 0.30) to ensure branch heads always get meaningful
+        // gradient even after C51 warmup. The alpha ramp prevents MSE suppression
+        // during warmup, but MUST NOT drop below 30% post-warmup.
         {
             let alpha = self.trainer.c51_alpha();
             let cql_frac = if self.trainer.has_cql() { CQL_GRAD_BUDGET } else { 0.0 };
             let iqn_frac = if self.gpu_iqn.is_some() { IQN_GRAD_BUDGET } else { 0.0 };
             let ens_frac = if !self.ensemble_extra_heads.is_empty() { ENS_GRAD_BUDGET } else { 0.0 };
-            let c51_frac = 1.0 - cql_frac - iqn_frac - ens_frac;
+            let c51_frac = (1.0 - cql_frac - iqn_frac - ens_frac).max(0.30);
             let mgn = self.trainer.config().max_grad_norm;
             let primary_budget = mgn * (1.0 - alpha) + mgn * c51_frac * alpha;
             self.trainer.clip_grad_buf_inplace(primary_budget)
diff --git a/crates/ml/src/trainers/dqn/trainer/metrics.rs b/crates/ml/src/trainers/dqn/trainer/metrics.rs
index cfb5d6bd3..3e550e55a 100644
--- a/crates/ml/src/trainers/dqn/trainer/metrics.rs
+++ b/crates/ml/src/trainers/dqn/trainer/metrics.rs
@@ -502,6 +502,12 @@ impl DQNTrainer {
         let evaluator = self.gpu_evaluator.as_mut()
             .ok_or_else(|| anyhow::anyhow!("gpu_evaluator must be initialized"))?;
 
+        // Invalidate CUDA Graph — weights changed since last epoch.
+        // The graph captures DtoD copies from DuelingWeightSet pointers, which are
+        // stable (same CudaSlice), but the portfolio state and step returns also need
+        // fresh computation. Without invalidation, the graph replays stale results.
+        evaluator.invalidate_dqn_graph();
+
         let hp = &self.hyperparams;
         #[allow(clippy::cast_possible_truncation)]
         let dqn_cfg = DqnBacktestConfig {