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

fix: validation reads GPU-resident weights instead of stale Candle VarMap

Browse Source 
The GPU evaluator was reading weights from agent.get_q_network_vars()
(Candle VarMap) which is NOT synced during fused CUDA training. Weights
stayed at epoch-0 values → val_loss=0.804688 every epoch (constant).

Fix: read weights directly from fused_ctx.online_dueling_ref() and
fused_ctx.online_branching_ref() — the GPU-resident buffers updated
by Adam each training step.

Added online_dueling_ref() and online_branching_ref() accessor methods
to FusedTrainingCtx.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-04-06 23:57:58 +02:00
parent 74f63f80b2
commit 7a008e71b0
2 changed files with 22 additions and 28 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -1477,6 +1477,14 @@ impl FusedTrainingCtx {
/// The training guard reads this pointer for NaN detection, so we must return
/// `mse_loss_buf` during warmup to avoid false NaN halts.
/// Once C51 is active (alpha > 0), returns `total_loss_buf` (C51 accumulator).
pub(crate) fn online_dueling_ref(&self) -> &DuelingWeightSet {
&self.online_dueling
}
pub(crate) fn online_branching_ref(&self) -> &BranchingWeightSet {
&self.online_branching
}
pub(crate) fn loss_gpu_buf(&self) -> &cudarc::driver::CudaSlice<f32> {
if self.trainer.c51_alpha() < 1e-6 {
&self.trainer.mse_loss_buf

42
crates/ml/src/trainers/dqn/trainer/metrics.rs
View File

@@ -472,36 +472,25 @@ impl DQNTrainer {
self.gpu_evaluator = Some(evaluator);
}
// ── Extract current weights from agent ──
// ── Extract current weights from fused CUDA trainer (GPU-resident) ──
// The fused_ctx owns the up-to-date weights (updated by Adam each step).
// Do NOT use agent.get_q_network_vars() — that reads the Candle VarMap
// which is NOT synced during fused CUDA training (stale epoch-0 weights).
let fused = self.fused_ctx.as_ref()
.ok_or_else(|| anyhow::anyhow!("fused_ctx required for GPU validation"))?;
let online_weights = fused.online_dueling_ref();
let branching_weights = fused.online_branching_ref();
let agent = self.agent.read().await;
let is_branching = agent.is_using_branching();
let network_dims = agent.network_dims();
let vars = agent.get_q_network_vars();
let (b0, b1, b2) = agent.branch_sizes();
let is_branching = agent.is_using_branching();
drop(agent);
let evaluator = self.gpu_evaluator.as_mut()
.ok_or_else(|| anyhow::anyhow!("gpu_evaluator must be initialized"))?;
let eval_stream = evaluator.stream();
let online_weights = if is_branching {
crate::cuda_pipeline::gpu_weights::extract_dueling_weights_branching(
vars, eval_stream,
).map_err(|e| anyhow::anyhow!("extract dueling weights (branching): {e}"))?
} else {
crate::cuda_pipeline::gpu_weights::extract_dueling_weights(
vars, eval_stream,
).map_err(|e| anyhow::anyhow!("extract dueling weights: {e}"))?
};
let branching_weights = if is_branching {
Some(crate::cuda_pipeline::gpu_weights::extract_branching_weights(
vars, eval_stream,
).map_err(|e| anyhow::anyhow!("extract branching weights: {e}"))?)
} else {
None
};
let hp = &self.hyperparams;
let (b0, b1, b2) = agent.branch_sizes();
#[allow(clippy::cast_possible_truncation)]
let dqn_cfg = DqnBacktestConfig {
shared_h1: network_dims.0,
@@ -516,16 +505,13 @@ impl DQNTrainer {
v_max: hp.v_max as f32,
};
// Drop agent read guard before mutable evaluator call
drop(agent);
// Weights changed since last epoch — invalidate the cached CUDA graph
// so evaluate_dqn_graphed re-captures with current weight pointers.
evaluator.invalidate_dqn_graph();
let metrics = evaluator.evaluate_dqn_graphed(
&online_weights,
branching_weights.as_ref(),
online_weights,
Some(branching_weights),
&dqn_cfg,
).map_err(|e| anyhow::anyhow!("GPU backtest evaluate_dqn_graphed: {e}"))?;