fix(training): gradient cliff at C51 warmup + stale val_loss

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>
This commit is contained in:
jgrusewski
2026-04-09 00:09:17 +02:00
parent aa611a938a
commit 0de7f69418
2 changed files with 15 additions and 4 deletions

View File

@@ -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)

View File

@@ -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 {