gitadmin/foxhunt
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix(critical): mean-reduce gradients + fix ExposureLevel 4-branch mismatch

Browse Source 
Three root causes found and fixed:

1. SUM-reduced gradients without 1/N: all CUDA loss gradient kernels
   (C51, MSE, IQN backward, CQL) accumulated per-sample gradients as
   raw SUM. At batch=16384 (H100) the raw norm was 282x larger than
   batch=58, causing gradient clipping to destroy signal-to-noise ratio
   and collapse training at epoch 2-3. Now all kernels multiply by
   1/batch_size, making gradient scale batch-invariant.

2. ExposureLevel::target_exposure() used a flat 9-level scale that did
   not match the 4-branch dir*mag encoding. The Rust backtest evaluator
   computed wrong position sizes (e.g. 4x oversize for Short+Small).
   Now uses dir x mag formula. Also fixed is_buy/is_sell/is_hold and
   from_trading_action for 4-branch semantics.

3. Rust epsilon-greedy only explored 5/9 exposure combos (0..5 instead
   of dir*3+mag), ignored the magnitude branch on greedy, and used wrong
   indices for order/urgency (get(1)/get(2) instead of get(2)/get(3)).

LR recalibrated: old gradient_clip_norm was accidentally a batch-size-
dependent LR reducer (~60x at batch=58, ~16000x at batch=16384). With
mean-reduced gradients the clip rarely fires, so LR is now the sole
training speed control. Smoketest 1e-4 -> 2e-6, production 1e-4 -> 1e-5,
hyperopt range [1e-5,3e-4] -> [1e-7,1e-4].

Diagnostics: FOXHUNT_GRAD_DIAG=1 enables per-stage gradient norm logging.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-04-09 14:18:36 +02:00
parent 5bdeb6c6d6
commit 30f46c04c1
14 changed files with 252 additions and 90 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
config/training/dqn-hyperopt.toml
View File

@@ -1,10 +1,14 @@
# DQN Hyperopt Profile — PSO search space definition
# All bounds are [min, max] ranges. The adapter reads these at runtime.
# To constrain the search space, narrow the ranges here — no code changes needed.
#
# Mean-reduced gradients (2026-04-09): LR range shifted down because gradient
# clipping no longer acts as a hidden LR reducer. Old range [1e-5, 3e-4] had
# effective LR of [~5e-10, ~1.5e-8] due to SUM-reduced clipping at batch=16384.
[search_space]
# Base parameters
learning_rate = [0.00001, 0.0003] # log scale in adapter
learning_rate = [1e-7, 1e-4] # log scale in adapter — mean-reduced gradients
batch_size = [4096, 16384]
gamma = [0.90, 0.99] # wider range — v_range computed dynamically from gamma
buffer_size = [50000, 100000] # log scale in adapter

11
config/training/dqn-production.toml
View File

@@ -4,12 +4,17 @@
# Mixed-precision (NVIDIA AMP pattern):
# - f32 master weights, Adam moments, gradients, loss accumulators
# - bf16 shadow copies for cuBLAS GemmEx tensor core GEMM
# - Cosine LR decay: 1e-4 → 1e-5 over 100 epochs
# - Cosine LR decay: 1e-5 → 1e-6 over 200 epochs
#
# Mean-reduced gradients (2026-04-09): all loss kernels divide by batch_size.
# gradient_clip_norm=1.0 is a safety net (norm ≈ 1.0, clip rarely fires).
# LR is the sole training speed control. Old lr=1e-4 was masked by SUM-reduced
# clipping (effective LR was ~6e-9). Start at 1e-5 and tune from there.
[training]
epochs = 200
batch_size = 16384
learning_rate = 0.0001
learning_rate = 1e-5
gamma = 0.99
weight_decay = 0.0001
adam_epsilon = 1e-8
@@ -18,7 +23,7 @@ hidden_dim_base = 256
reward_scale = 1.0
huber_delta = 1.0
lr_decay_type = 2
lr_min = 0.00001
lr_min = 1e-6
max_bars = 0
# Data source: "ohlcv" (1-min candles) or "mbp10" (imbalance bars from MBP-10 order book)
data_source = "mbp10"

2
config/training/dqn-smoketest.toml
View File

@@ -12,7 +12,7 @@
[training]
epochs = 3
batch_size = 64
learning_rate = 0.0001
learning_rate = 2e-6
gamma = 0.95
weight_decay = 0.0001
adam_epsilon = 1e-8

32
crates/ml-core/src/action_space.rs
View File

@@ -77,11 +77,18 @@ mod tests {
#[test]
fn test_target_exposure_values() {
assert_eq!(ExposureLevel::Short100.target_exposure(), -1.0);
assert_eq!(ExposureLevel::Short50.target_exposure(), -0.5);
assert_eq!(ExposureLevel::Flat.target_exposure(), 0.0);
assert_eq!(ExposureLevel::Long50.target_exposure(), 0.5);
assert_eq!(ExposureLevel::Long100.target_exposure(), 1.0);
// 4-branch: dir × mag. Index = dir*3 + mag.
// dir: 0=Short(-1), 1=Flat(0), 2=Long(+1)
// mag: 0=Small(0.25), 1=Half(0.50), 2=Full(1.00)
assert_eq!(ExposureLevel::Short100.target_exposure(), -0.25); // Short×Small
assert_eq!(ExposureLevel::Short75.target_exposure(), -0.50); // Short×Half
assert_eq!(ExposureLevel::Short50.target_exposure(), -1.0); // Short×Full
assert_eq!(ExposureLevel::Short25.target_exposure(), 0.0); // Flat×Small
assert_eq!(ExposureLevel::Flat.target_exposure(), 0.0); // Flat×Half
assert_eq!(ExposureLevel::Long25.target_exposure(), 0.0); // Flat×Full
assert_eq!(ExposureLevel::Long50.target_exposure(), 0.25); // Long×Small
assert_eq!(ExposureLevel::Long75.target_exposure(), 0.50); // Long×Half
assert_eq!(ExposureLevel::Long100.target_exposure(), 1.0); // Long×Full
}
#[test]
@@ -169,26 +176,27 @@ mod tests {
#[test]
fn test_extreme_actions() {
// Test Short100 aggressive market (index 0 * 9 + 0 * 3 + 2 = 2)
// 4-branch: Short100(=0) is Short×Small = -0.25 (NOT -1.0)
// Full short is Short50(=2) = Short×Full = -1.0
let short_extreme = FactoredAction::new(
ExposureLevel::Short100,
ExposureLevel::Short50, // Short×Full = -1.0
OrderType::Market,
Urgency::Aggressive,
);
assert_eq!(short_extreme.to_index(), 2);
assert_eq!(short_extreme.to_index(), 2 * 9 + 0 * 3 + 2); // 20
assert_eq!(short_extreme.target_exposure(), -1.0);
assert_eq!(short_extreme.transaction_cost(), 0.0015); // Wave 2.5 calibration
assert_eq!(short_extreme.transaction_cost(), 0.0015);
assert_eq!(short_extreme.urgency_weight(), 1.5);
// Test Long100 aggressive market (index 8 * 9 + 0 * 3 + 2 = 74)
// Long100(=8) is Long×Full = +1.0 (still correct)
let long_extreme = FactoredAction::new(
ExposureLevel::Long100,
OrderType::Market,
Urgency::Aggressive,
);
assert_eq!(long_extreme.to_index(), 74);
assert_eq!(long_extreme.to_index(), 8 * 9 + 0 * 3 + 2); // 74
assert_eq!(long_extreme.target_exposure(), 1.0);
assert_eq!(long_extreme.transaction_cost(), 0.0015); // Wave 2.5 calibration
assert_eq!(long_extreme.transaction_cost(), 0.0015);
assert_eq!(long_extreme.urgency_weight(), 1.5);
}

116
crates/ml-core/src/common/action.rs
View File

@@ -66,34 +66,48 @@ impl fmt::Display for OrderType {
/// Exposure level for position sizing (-100% to +100%)
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub enum ExposureLevel {
Short100 = 0, // -100% of max position
Short75 = 1, // -75%
Short50 = 2, // -50%
Short25 = 3, // -25%
Flat = 4, // 0% (neutral)
Long25 = 5, // +25%
Long50 = 6, // +50%
Long75 = 7, // +75%
Long100 = 8, // +100%
/// 4-branch factored DQN: discriminant = dir * 3 + mag
/// dir: 0=Short(-1), 1=Flat(0), 2=Long(+1)
/// mag: 0=Small(0.25), 1=Half(0.50), 2=Full(1.00)
///
/// LEGACY NAMES: variant names reflect the old flat 9-level model.
/// Actual position = direction × magnitude (see target_exposure()).
Short100 = 0, // dir=0(Short), mag=0(Small) → -0.25
Short75 = 1, // dir=0(Short), mag=1(Half) → -0.50
Short50 = 2, // dir=0(Short), mag=2(Full) → -1.00
Short25 = 3, // dir=1(Flat), mag=0(Small) → 0
Flat = 4, // dir=1(Flat), mag=1(Half) → 0
Long25 = 5, // dir=1(Flat), mag=2(Full) → 0
Long50 = 6, // dir=2(Long), mag=0(Small) → +0.25
Long75 = 7, // dir=2(Long), mag=1(Half) → +0.50
Long100 = 8, // dir=2(Long), mag=2(Full) → +1.00
}
impl ExposureLevel {
/// Get target portfolio value percentage (-1.0 to +1.0 in 0.25 steps)
/// Target exposure for the 4-branch factored DQN.
///
/// Decodes from composite index `dir * 3 + mag`:
/// direction: [-1.0, 0.0, +1.0]
/// magnitude: [0.25, 0.50, 1.00]
/// result: direction × magnitude
pub fn target_exposure(&self) -> f64 {
match self {
ExposureLevel::Short100 => -1.0,
ExposureLevel::Short75 => -0.75,
ExposureLevel::Short50 => -0.5,
ExposureLevel::Short25 => -0.25,
ExposureLevel::Flat => 0.0,
ExposureLevel::Long25 => 0.25,
ExposureLevel::Long50 => 0.5,
ExposureLevel::Long75 => 0.75,
ExposureLevel::Long100 => 1.0,
}
let idx = *self as usize;
let dir = idx / 3;
let mag = idx % 3;
let direction: f64 = match dir {
0 => -1.0,
2 => 1.0,
_ => 0.0,
};
let magnitude: f64 = match mag {
0 => 0.25,
1 => 0.50,
_ => 1.00,
};
direction * magnitude
}
/// Convert from index (0-8)
/// Convert from composite index (0-8): dir * 3 + mag
pub fn from_index(idx: usize) -> Result<Self, MLError> {
match idx {
0 => Ok(ExposureLevel::Short100),
@@ -115,16 +129,17 @@ impl ExposureLevel {
impl fmt::Display for ExposureLevel {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
// Display as 4-branch dir×mag combo for clarity
match self {
ExposureLevel::Short100 => write!(f, "Short100"),
ExposureLevel::Short75 => write!(f, "Short75"),
ExposureLevel::Short50 => write!(f, "Short50"),
ExposureLevel::Short25 => write!(f, "Short25"),
ExposureLevel::Flat => write!(f, "Flat"),
ExposureLevel::Long25 => write!(f, "Long25"),
ExposureLevel::Long50 => write!(f, "Long50"),
ExposureLevel::Long75 => write!(f, "Long75"),
ExposureLevel::Long100 => write!(f, "Long100"),
ExposureLevel::Short100 => write!(f, "S25"), // Short×Small
ExposureLevel::Short75 => write!(f, "S50"), // Short×Half
ExposureLevel::Short50 => write!(f, "S100"), // Short×Full
ExposureLevel::Short25 => write!(f, "F25"), // Flat×Small
ExposureLevel::Flat => write!(f, "F50"), // Flat×Half
ExposureLevel::Long25 => write!(f, "F100"), // Flat×Full
ExposureLevel::Long50 => write!(f, "L25"), // Long×Small
ExposureLevel::Long75 => write!(f, "L50"), // Long×Half
ExposureLevel::Long100 => write!(f, "L100"), // Long×Full
}
}
}
@@ -245,18 +260,21 @@ impl FactoredAction {
}
/// Create a FactoredAction from a TradingAction.
/// Maps Buy -> Long100/Market/Normal, Sell -> Short100/Market/Normal, Hold -> Flat/Market/Normal.
/// Buy → Long+Full(+1.0), Sell → Short+Full(-1.0), Hold → Flat+Half(0).
///
/// In the 4-branch encoding: Long+Full = Long100(=8), Short+Full = Short50(=2),
/// Flat+Half = Flat(=4). Legacy variant names don't match 4-branch semantics.
pub fn from_trading_action(action: crate::trading_action::TradingAction) -> Self {
use crate::trading_action::TradingAction;
match action {
TradingAction::Buy => {
Self::new(ExposureLevel::Long100, OrderType::Market, Urgency::Normal)
Self::new(ExposureLevel::Long100, OrderType::Market, Urgency::Normal) // Long+Full = +1.0
}
TradingAction::Sell => {
Self::new(ExposureLevel::Short100, OrderType::Market, Urgency::Normal)
Self::new(ExposureLevel::Short50, OrderType::Market, Urgency::Normal) // Short+Full = -1.0
}
TradingAction::Hold => {
Self::new(ExposureLevel::Flat, OrderType::Market, Urgency::Normal)
Self::new(ExposureLevel::Flat, OrderType::Market, Urgency::Normal) // Flat+Half = 0
}
}
}
@@ -292,25 +310,28 @@ impl FactoredAction {
trade_value * self.transaction_cost()
}
/// Check if this action is a buy (long exposure)
/// Check if this action is a buy (long direction, dir=2, indices 6-8)
pub fn is_buy(&self) -> bool {
matches!(
self.exposure,
ExposureLevel::Long100 | ExposureLevel::Long50
ExposureLevel::Long50 | ExposureLevel::Long75 | ExposureLevel::Long100
)
}
/// Check if this action is a sell (short exposure)
/// Check if this action is a sell (short direction, dir=0, indices 0-2)
pub fn is_sell(&self) -> bool {
matches!(
self.exposure,
ExposureLevel::Short100 | ExposureLevel::Short75 | ExposureLevel::Short50 | ExposureLevel::Short25
ExposureLevel::Short100 | ExposureLevel::Short75 | ExposureLevel::Short50
)
}
/// Check if this action is neutral (flat exposure)
/// Check if this action is neutral (flat direction, dir=1, indices 3-5)
pub fn is_hold(&self) -> bool {
matches!(self.exposure, ExposureLevel::Flat)
matches!(
self.exposure,
ExposureLevel::Short25 | ExposureLevel::Flat | ExposureLevel::Long25
)
}
/// Convert action to position delta
@@ -411,11 +432,12 @@ mod tests {
#[test]
fn test_target_exposure() {
assert_eq!(ExposureLevel::Short100.target_exposure(), -1.0);
assert_eq!(ExposureLevel::Short50.target_exposure(), -0.5);
assert_eq!(ExposureLevel::Flat.target_exposure(), 0.0);
assert_eq!(ExposureLevel::Long50.target_exposure(), 0.5);
assert_eq!(ExposureLevel::Long100.target_exposure(), 1.0);
// 4-branch: dir × mag formula
assert_eq!(ExposureLevel::Short100.target_exposure(), -0.25); // Short×Small
assert_eq!(ExposureLevel::Short50.target_exposure(), -1.0); // Short×Full
assert_eq!(ExposureLevel::Flat.target_exposure(), 0.0); // Flat×Half
assert_eq!(ExposureLevel::Long50.target_exposure(), 0.25); // Long×Small
assert_eq!(ExposureLevel::Long100.target_exposure(), 1.0); // Long×Full
}
#[test]
@@ -465,6 +487,6 @@ mod tests {
let action =
FactoredAction::new(ExposureLevel::Long50, OrderType::LimitMaker, Urgency::Patient);
let s = format!("{}", action);
assert_eq!(s, "Long50+LimitMaker+Patient");
assert_eq!(s, "L25+LimitMaker+Patient"); // Long50 displays as L25 (Long×Small)
}
}

21
crates/ml-dqn/src/dqn.rs
View File

@@ -1520,8 +1520,11 @@ impl DQN {
let urg_random = in_warmup || rng.gen::<f32>() < effective_epsilon;
if exp_random && ord_random && urg_random {
// All branches random — skip forward pass
let exposure = ExposureLevel::from_index(rng.gen_range(0..5_usize))?;
// All branches random — skip forward pass.
// 4-branch: exposure = dir*3 + mag (9 combos), order(3), urgency(3).
let dir = rng.gen_range(0..3_usize);
let mag = rng.gen_range(0..3_usize);
let exposure = ExposureLevel::from_index(dir * 3 + mag)?;
let order = OrderType::from_index(rng.gen_range(0..3_usize))?;
let urgency = Urgency::from_index(rng.gen_range(0..3_usize))?;
FactoredAction { exposure, order, urgency }
@@ -1536,20 +1539,26 @@ impl DQN {
let output = branching_net.forward_branches_eval(&state_tensor)?;
let greedy = super::branching::BranchingDuelingQNetwork::greedy_branch_actions(&output, &self.stream)?;
// 4-branch: greedy[0]=dir(0-2), greedy[1]=mag(0-2), composite = dir*3+mag
let exposure = if exp_random {
ExposureLevel::from_index(rng.gen_range(0..5_usize))?
let dir = rng.gen_range(0..3_usize);
let mag = rng.gen_range(0..3_usize);
ExposureLevel::from_index(dir * 3 + mag)?
} else {
ExposureLevel::from_index(greedy.first().copied().unwrap_or(2) as usize)?
let dir = greedy.first().copied().unwrap_or(1) as usize;
let mag = greedy.get(1).copied().unwrap_or(1) as usize;
ExposureLevel::from_index(dir * 3 + mag)?
};
// 4-branch greedy: [dir, mag, order, urgency] → order at [2], urgency at [3]
let order = if ord_random {
OrderType::from_index(rng.gen_range(0..3_usize))?
} else {
OrderType::from_index(greedy.get(1).copied().unwrap_or(0) as usize)?
OrderType::from_index(greedy.get(2).copied().unwrap_or(0) as usize)?
};
let urgency = if urg_random {
Urgency::from_index(rng.gen_range(0..3_usize))?
} else {
Urgency::from_index(greedy.get(2).copied().unwrap_or(1) as usize)?
Urgency::from_index(greedy.get(3).copied().unwrap_or(1) as usize)?
};
FactoredAction { exposure, order, urgency }
}

4
crates/ml/src/cuda_pipeline/backward_kernels.cu
View File

@@ -68,5 +68,9 @@ extern "C" __global__ void bias_grad_reduce_kernel(
for (int b = 0; b < batch_size; b++) {
sum += (float)dy[b * out_dim + j];
}
/* No 1/N here: the upstream dY is already mean-reduced (1/N applied in
* c51_grad_kernel / mse_grad_kernel). Chain rule propagates it through
* cuBLAS dX = dY @ W, so dY at every layer is already 1/N-scaled.
* Dividing again would be 1/N². */
atomicAdd(&db[j], sum);
}

10
crates/ml/src/cuda_pipeline/c51_grad_kernel.cu
View File

@@ -38,8 +38,14 @@ extern "C" __global__ void c51_grad_kernel(
float proj = (float)projected[tid];
/* Cross-entropy gradient: d/d_logits(-Sigma proj * lp) = exp(lp) - proj
* Float exp() handles full range — no bf16 overflow */
float d_combined = isw * (expf(lp) - proj);
* Float exp() handles full range — no bf16 overflow.
*
* MEAN-reduce: divide by batch_size so gradient scale is invariant to
* batch size. Without this, batch=16384 (H100) produces a 282× larger
* raw gradient sum than batch=58 (smoke test), causing the budget clip
* to destroy gradient SNR on large batches → epoch 2-3 collapse. */
float inv_batch = 1.0f / (float)batch_size;
float d_combined = inv_batch * isw * (expf(lp) - proj);
/* Entropy regularization — magnitude branch (d==1) gets 5× boost to prevent
* atom distribution collapse. Standardized advantages prevent Q-value scale

7
crates/ml/src/cuda_pipeline/cql_grad_kernel.cu
View File

@@ -28,6 +28,11 @@ extern "C" __global__ void cql_logit_grad_kernel(
int i = blockIdx.x * blockDim.x + threadIdx.x;
if (i >= N) return;
/* MEAN-reduce: scale all CQL gradient outputs by 1/N so downstream cuBLAS
* backward produces mean-reduced weight gradients. Keeps gradient scale
* invariant to batch size — consistent with C51/MSE/IQN kernels. */
float inv_batch = 1.0f / (float)N;
int total_actions = b0_size + b1_size + b2_size + b3_size;
float dz = (num_atoms > 1) ? (v_max - v_min) / (float)(num_atoms - 1) : 0.0f;
@@ -181,7 +186,7 @@ extern "C" __global__ void cql_logit_grad_kernel(
a_mean_j /= (float)bd;
float p = expf(val[j] + adv[j] - a_mean_j - log_sum);
float z = v_min + (float)j * dz;
float d_combined = d_cql_dq[a] * p * (z - eq);
float d_combined = inv_batch * d_cql_dq[a] * p * (z - eq);
d_adv[j] = d_combined;
if (j < 256) d_val_accum[j] += d_combined;

28
crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs
View File

@@ -2108,6 +2108,34 @@ impl GpuDqnTrainer {
Ok(norm_bf16[0].to_f32())
}
/// Compute and return the current grad_buf L2 norm using the STANDALONE
/// kernel handle (safe after mega-graph capture).
///
/// `read_grad_norm_sync()` uses the graph-captured kernel handle which
/// corrupts after `graph_mega` capture. This variant uses
/// `grad_norm_standalone` — a separately loaded handle of the same kernel.
///
/// Costs one stream sync + 4-byte DtoH. For diagnostic use only.
pub fn read_grad_norm_standalone_sync(&mut self) -> Result<f32, MLError> {
let _evt_guard = EventTrackingGuard::new(self.stream.context());
self.launch_grad_norm_standalone()?;
self.launch_grad_norm_finalize()?;
unsafe { cudarc::driver::sys::cuStreamSynchronize(self.stream.cu_stream()); }
let mut norm_f32 = [0.0_f32; 1];
unsafe {
cudarc::driver::sys::cuMemcpyDtoH_v2(
norm_f32.as_mut_ptr().cast(),
self.grad_norm_f32_buf.raw_ptr(),
std::mem::size_of::<f32>(),
);
}
// grad_norm_f32_buf holds sum-of-squares; return L2 norm
Ok(norm_f32[0].sqrt())
}
/// Apply GPU multi-head feature attention to `save_h_s2` (post-graph).
///
/// Runs 4-head self-attention over the trunk output `h_s2 [B, SHARED_H2]`

31
crates/ml/src/cuda_pipeline/iqn_dual_head_kernel.cu
View File

@@ -562,6 +562,11 @@ void iqn_backward_kernel(
/* Process each quantile τ_i */
float kappa = IQN_KAPPA;
float inv_n_sq = 1.0f / ((float)IQN_NUM_QUANTILES * (float)IQN_NUM_QUANTILES);
/* MEAN-reduce: all atomicAdd contributions are scaled by 1/batch_size so
* gradient magnitude is invariant to batch size. Without this, batch=16384
* produces ~282× larger raw gradients than batch=58, causing the budget
* clip system to destroy gradient SNR on H100. */
float inv_batch = 1.0f / (float)batch_size;
for (int ti = 0; ti < IQN_NUM_QUANTILES; ti++) {
float tau_i = my_taus[ti];
@@ -590,33 +595,33 @@ void iqn_backward_kernel(
for (int h = tid; h < hidden_dim; h += IQN_BLOCK_SIZE)
dL_dcomb[h / IQN_BLOCK_SIZE] = 0.0f;
/* Branch 0: accumulate gradient */
/* Branch 0: accumulate gradient (mean-reduced by inv_batch) */
for (int h = tid; h < hidden_dim; h += IQN_BLOCK_SIZE) {
dL_dcomb[h / IQN_BLOCK_SIZE] += w_b0[a0 * hidden_dim + h] * dL_dq;
/* dL/dW_b0[a0, h] += dL/dq × combined[h] */
/* dL/dW_b0[a0, h] += inv_batch × dL/dq × combined[h] */
atomicAdd(&grad_buf[off[2] + a0 * hidden_dim + h],
dL_dq * comb_dist[h / IQN_BLOCK_SIZE]);
inv_batch * dL_dq * comb_dist[h / IQN_BLOCK_SIZE]);
}
if (tid == 0)
atomicAdd(&grad_buf[off[3] + a0], dL_dq);
atomicAdd(&grad_buf[off[3] + a0], inv_batch * dL_dq);
/* Branch 1 */
for (int h = tid; h < hidden_dim; h += IQN_BLOCK_SIZE) {
dL_dcomb[h / IQN_BLOCK_SIZE] += w_b1[a1 * hidden_dim + h] * dL_dq;
atomicAdd(&grad_buf[off[4] + a1 * hidden_dim + h],
dL_dq * comb_dist[h / IQN_BLOCK_SIZE]);
inv_batch * dL_dq * comb_dist[h / IQN_BLOCK_SIZE]);
}
if (tid == 0)
atomicAdd(&grad_buf[off[5] + a1], dL_dq);
atomicAdd(&grad_buf[off[5] + a1], inv_batch * dL_dq);
/* Branch 2 */
for (int h = tid; h < hidden_dim; h += IQN_BLOCK_SIZE) {
dL_dcomb[h / IQN_BLOCK_SIZE] += w_b2[a2 * hidden_dim + h] * dL_dq;
atomicAdd(&grad_buf[off[6] + a2 * hidden_dim + h],
dL_dq * comb_dist[h / IQN_BLOCK_SIZE]);
inv_batch * dL_dq * comb_dist[h / IQN_BLOCK_SIZE]);
}
if (tid == 0)
atomicAdd(&grad_buf[off[7] + a2], dL_dq);
atomicAdd(&grad_buf[off[7] + a2], inv_batch * dL_dq);
/* ── Gradient through element-wise product ── */
/* combined = h_s2 ⊙ embed
@@ -625,9 +630,9 @@ void iqn_backward_kernel(
float dL_dembed[IQN_DIST_MAX];
for (int h = tid; h < hidden_dim; h += IQN_BLOCK_SIZE) {
dL_dembed[h / IQN_BLOCK_SIZE] = dL_dcomb[h / IQN_BLOCK_SIZE] * h_dist[h / IQN_BLOCK_SIZE];
/* Accumulate dL/d(h_s2) across all quantiles.
/* Accumulate dL/d(h_s2) across all quantiles (mean-reduced).
* This flows IQN's bounded Huber gradient to the shared trunk. */
float d_trunk = dL_dcomb[h / IQN_BLOCK_SIZE] * embed_dist[h / IQN_BLOCK_SIZE];
float d_trunk = inv_batch * dL_dcomb[h / IQN_BLOCK_SIZE] * embed_dist[h / IQN_BLOCK_SIZE];
if (isfinite(d_trunk))
atomicAdd(&d_h_s2_out[sample * hidden_dim + h], d_trunk);
}
@@ -647,12 +652,12 @@ void iqn_backward_kernel(
for (int h = tid; h < hidden_dim; h += IQN_BLOCK_SIZE) {
float dL_dpre = dL_dembed[h / IQN_BLOCK_SIZE];
if (!isfinite(dL_dpre)) continue;
/* Bias gradient */
atomicAdd(&grad_buf[off[1] + h], dL_dpre);
/* Bias gradient (mean-reduced) */
atomicAdd(&grad_buf[off[1] + h], inv_batch * dL_dpre);
/* Weight gradient: outer product with precomputed cosine features */
for (int d = 0; d < embed_dim; d++) {
atomicAdd(&grad_buf[off[0] + h * embed_dim + d],
dL_dpre * my_cos[d]);
inv_batch * dL_dpre * my_cos[d]);
}
}
}

9
crates/ml/src/cuda_pipeline/mse_grad_kernel.cu
View File

@@ -42,8 +42,13 @@ extern "C" __global__ void mse_grad_kernel(
float delta_z = (num_atoms > 1) ? (v_max - v_min) / (float)(num_atoms - 1) : 0.0f;
float z_j = v_min + (float)j * delta_z;
/* Gradient of MSE loss through softmax expectation (float arithmetic) */
float d_combined = isw * td_error * p_j * (z_j - e_q);
/* Gradient of MSE loss through softmax expectation (float arithmetic).
*
* MEAN-reduce: divide by batch_size so gradient scale is invariant to
* batch size. Without this, the budget clip is 282× more aggressive on
* H100 (batch=16384) than on smoke test (batch=58). */
float inv_batch = 1.0f / (float)batch_size;
float d_combined = inv_batch * isw * td_error * p_j * (z_j - e_q);
/* Magnitude entropy boost: prevent atom distribution collapse for d==1.
* MSE path has no entropy param — hardcoded 0.005 = 5× base (0.001). */

2
crates/ml/src/trainers/dqn/config.rs
View File

@@ -1589,7 +1589,7 @@ impl DQNHyperparameters {
time_decay_rate: 0.0005,
q_gap_threshold: 0.05, // Tier 2 default: mild conviction gating (hyperopt searches [0.0, 0.5])
huber_delta: 100.0, // BUG #12 FIX: Scale delta 100x for gradient explosion fix (was 1.0)
gradient_clip_norm: Some(10.0), // C51 101-atoms × 3 branches → gradient naturally 300x larger than DQN
gradient_clip_norm: Some(1.0), // Safety net — mean-reduced gradients have norm ≈ 1.0
hold_penalty_weight: 0.01, // Default: 1% penalty weight
movement_threshold: 0.02, // Default: 2% price movement threshold
preprocessing_window: 50, // Default: 50-bar rolling window

63
crates/ml/src/trainers/dqn/fused_training.rs
View File

@@ -217,6 +217,10 @@ pub(crate) struct FusedTrainingCtx {
pub(crate) hindsight_fraction: f64,
/// v8: Curriculum learning enabled (false by default).
pub(crate) curriculum_enabled: bool,
/// When true, disables mega-graph/graph_aux capture and logs per-stage
/// gradient norms (raw → post-primary-clip → post-IQN → post-CQL).
/// Costs ~4 stream syncs per step. For debugging H100 gradient collapse.
pub(crate) gradient_stage_diagnostics: bool,
}
impl Drop for FusedTrainingCtx {
@@ -610,6 +614,7 @@ impl FusedTrainingCtx {
popart_enabled: hyperparams.popart_enabled,
hindsight_fraction: hyperparams.hindsight_fraction,
curriculum_enabled: hyperparams.curriculum_enabled,
gradient_stage_diagnostics: std::env::var("FOXHUNT_GRAD_DIAG").is_ok(),
})
}
@@ -852,7 +857,8 @@ impl FusedTrainingCtx {
// Capture mega-graph at step 2 (all sub-trainers initialized).
// Mega-graph fuses spectral + forward + backward + aux → 1 launch.
if self.graph_mega.is_none() && self.steps_since_varmap_sync == 2 {
// Skip when gradient_stage_diagnostics is on — need ungraphed path for per-stage reads.
if self.graph_mega.is_none() && self.steps_since_varmap_sync == 2 && !self.gradient_stage_diagnostics {
if let Err(e) = self.capture_graph_mega(agent, gpu_batch) {
tracing::warn!(
"graph_mega capture failed (non-fatal, falling back to individual graphs): {e}"
@@ -916,6 +922,19 @@ impl FusedTrainingCtx {
// gradient even after C51 warmup. The alpha ramp prevents MSE suppression
// during warmup, but MUST NOT drop below 30% post-warmup.
{
// ── DIAG: raw C51+MSE gradient norm (before any clip) ─────────
if self.gradient_stage_diagnostics {
let raw_norm = self.trainer.read_grad_norm_standalone_sync()
.unwrap_or(f32::NAN);
self.last_c51_raw_norm = raw_norm;
tracing::warn!(
batch_size = self.batch_size,
raw_norm,
stage = "pre_primary_clip",
"GRAD_DIAG: raw C51+MSE gradient norm BEFORE budget clip"
);
}
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 };
@@ -925,6 +944,20 @@ impl FusedTrainingCtx {
let primary_budget = mgn * (1.0 - alpha) + mgn * c51_frac * alpha;
self.trainer.clip_grad_buf_inplace(primary_budget)
.map_err(|e| anyhow::anyhow!("Gradient budget clip: {e}"))?;
// ── DIAG: post-primary-clip gradient norm ─────────────────────
if self.gradient_stage_diagnostics {
let post_clip_norm = self.trainer.read_grad_norm_standalone_sync()
.unwrap_or(f32::NAN);
tracing::warn!(
batch_size = self.batch_size,
post_clip_norm,
primary_budget,
c51_alpha = alpha,
stage = "post_primary_clip",
"GRAD_DIAG: gradient norm AFTER primary budget clip"
);
}
}
// EMA target update.
@@ -991,6 +1024,18 @@ impl FusedTrainingCtx {
);
iqn.target_ema_update(tau as f32)
.map_err(|e| anyhow::anyhow!("IQN EMA update: {e}"))?;
// ── DIAG: post-IQN trunk SAXPY gradient norm ──────────
if self.gradient_stage_diagnostics {
let post_iqn_norm = self.trainer.read_grad_norm_standalone_sync()
.unwrap_or(f32::NAN);
tracing::warn!(
batch_size = self.batch_size,
post_iqn_norm,
stage = "post_iqn_saxpy",
"GRAD_DIAG: gradient norm AFTER IQN trunk clipped SAXPY"
);
}
}
Err(e) => {
tracing::warn!("IQN step failed (non-fatal): {e}");
@@ -1021,6 +1066,22 @@ impl FusedTrainingCtx {
}
}
// ── DIAG: final gradient norm (after all component injections) ────
if self.gradient_stage_diagnostics {
let final_norm = self.trainer.read_grad_norm_standalone_sync()
.unwrap_or(f32::NAN);
tracing::warn!(
batch_size = self.batch_size,
final_norm,
raw_norm = self.last_c51_raw_norm,
clip_ratio = if self.last_c51_raw_norm > 0.0 {
final_norm / self.last_c51_raw_norm
} else { 0.0 },
stage = "final_pre_adam",
"GRAD_DIAG: final gradient norm before Adam (all components injected)"
);
}
// Regime-adaptive PER scaling.
self.trainer.regime_scale_td_errors()
.map_err(|e| anyhow::anyhow!("Regime PER scaling: {e}"))?;