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

perf: GPU-native validation via GpuBacktestEvaluator (2s → ~50ms)

Browse Source 
Replaced the CPU validation path (compute_validation_loss) with the
existing GpuBacktestEvaluator. The CPU path downloaded Q-values to host,
did argmax + PnL + Sharpe on CPU — 2s/epoch from the GPU→CPU sync.

The GPU evaluator does everything on-device: cuBLAS forward, argmax,
PnL simulation, Sharpe/drawdown/PF computation. Lazy-initialized once
per fold from val_data, called each epoch with current weights.

Deleted:
- val_features_gpu, val_closes_gpu, val_ofi_gpu fields (CPU caches)
- CPU state tensor assembly (80-line loop)
- clone_htod_f32_to_bf16 upload
- dtoh_bf16_to_f32 download + GPU→CPU sync
- CPU argmax loop
- CPU Sharpe computation
- Epsilon save/restore

Net -61 lines of dead CPU code.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-04-06 22:57:38 +02:00
parent 3d0824eb9b
commit 6edb7b91fd
3 changed files with 129 additions and 190 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
crates/ml/src/trainers/dqn/trainer/constructor.rs
View File

@@ -752,9 +752,7 @@ impl DQNTrainer {
// OFI features: populated during data loading when MBP-10 data is available
ofi_features: None,
ofi_val_offset: 0,
val_features_gpu: None,
val_closes_gpu: None,
val_ofi_gpu: None,
gpu_evaluator: None,
// Phase C: Fill simulation and smart order routing
fill_simulator: FillSimulator::default(),

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

@@ -387,203 +387,152 @@ impl DQNTrainer {
None
}
/// Compute validation loss (negative Sharpe proxy) using cuBLAS forward — zero GpuTensor.
/// Compute validation loss (negative Sharpe) using the GPU backtest evaluator.
///
/// Builds the state tensor on CPU (features + portfolio broadcast + OFI + alignment
/// padding), uploads it once as a raw `CudaSlice<half::bf16>`, runs cuBLAS forward via
/// `fused_ctx`, then downloads Q-values and performs all action selection and Sharpe
/// computation on CPU. This avoids GpuTensor elementwise kernels (add, broadcast_mul,
/// cat, etc.) that cache `CudaFunction` handles on the shared CUDA primary context,
/// which causes sequential test failures.
/// Runs the full backtest step loop on GPU: gather kernel -> cuBLAS forward ->
/// env kernel -> metrics reduction. Only the final metrics (1 window x 14 floats)
/// are downloaded to CPU. Replaces the old CPU path (GPU->CPU Q-value download,
/// CPU argmax, CPU Sharpe) which took ~2s/epoch vs ~50ms here.
///
/// Epoch-boundary cold path — simplicity over GPU utilisation.
/// The evaluator is lazy-initialized on the first call and cached across epochs.
/// Weight pointers change each epoch (training updates weights), so we call
/// `invalidate_dqn_graph()` to force CUDA Graph re-capture.
pub(crate) async fn compute_validation_loss(&mut self) -> Result<f64> {
use crate::cuda_pipeline::gpu_backtest_evaluator::{
DqnBacktestConfig, GpuBacktestConfig, GpuBacktestEvaluator,
};
if self.val_data.is_empty() {
return Ok(0.0);
}
// Cap sample_size at fused_ctx batch_size to avoid exceeding cuBLAS buffer
let fused_batch = self.fused_ctx.as_ref().map(|f| f.batch_size()).unwrap_or(0);
if fused_batch == 0 {
return Ok(0.0); // fused_ctx not initialized — skip validation
}
// Save current epsilon and force to 0 for deterministic evaluation
let original_epsilon = self.get_epsilon().await?;
self.set_epsilon(0.0).await?; // Pure greedy selection
let sample_size = self.val_data.len().min(1000).min(fused_batch);
let aligned_state_dim = {
let agent = self.agent.read().await;
agent.get_state_dim()
let stream = match self.cuda_stream.as_ref() {
Some(s) => s,
None => return Ok(0.0), // CPU device — skip GPU validation
};
let stream = self.cuda_stream.as_ref()
.ok_or_else(|| anyhow::anyhow!("cuda_stream required for validation"))?;
// ── Lazy-init the GPU evaluator (once per fold, reused across epochs) ──
if self.gpu_evaluator.is_none() {
let ofi_enabled = self.ofi_features.is_some();
let market_dim: usize = 42;
let feature_dim: usize = if ofi_enabled { 50 } else { market_dim };
// ── CPU-cached validation data (lazy init, built once per training run) ──
// val_features_gpu stores flat features as Vec<f32> [sample_size * 42] on CPU.
// val_closes_gpu stores (current_closes, next_closes) as Vec<f32> on CPU.
// val_ofi_gpu stores flat OFI features as Vec<f32> [sample_size * 8] on CPU.
if self.val_features_gpu.is_none() {
let mut flat_features: Vec<f32> = Vec::with_capacity(sample_size * 42);
let mut current_closes: Vec<f32> = Vec::with_capacity(sample_size);
let mut next_closes: Vec<f32> = Vec::with_capacity(sample_size);
// Build a single window from all val_data
let mut prices: Vec<[f32; 4]> = Vec::with_capacity(self.val_data.len());
let mut features: Vec<Vec<f32>> = Vec::with_capacity(self.val_data.len());
for (feature_vec, target) in self.val_data.iter().take(sample_size) {
for &v in feature_vec.iter() {
flat_features.push(v as f32);
}
let cur = if target.len() >= 2 { target[0] } else { feature_vec[3] };
let nxt = if target.len() >= 2 { target[1] } else { cur };
current_closes.push(cur as f32);
next_closes.push(nxt as f32);
}
for (i, (fv, target)) in self.val_data.iter().enumerate() {
// Prices: [close, close, close, close] — same convention as hyperopt adapter
let close = if target.len() >= 2 { target[0] as f32 } else { fv[3] as f32 };
prices.push([close, close, close, close]);
self.val_features_gpu = Some(flat_features);
self.val_closes_gpu = Some((current_closes, next_closes));
if let Some(ref ofi) = self.ofi_features {
let mut flat_ofi: Vec<f32> = Vec::with_capacity(sample_size * 8);
for i in 0..sample_size {
let idx = self.ofi_val_offset + i;
if let Some(row) = ofi.get(idx) {
for &v in row.iter() {
flat_ofi.push(v as f32);
// Features: 42 market features, plus 8 OFI if enabled
let fv_slice = &fv[..market_dim.min(fv.len())];
let mut fv_f32: Vec<f32> = fv_slice.iter().map(|&v| v as f32).collect();
if ofi_enabled {
let ofi_idx = self.ofi_val_offset + i;
if let Some(ofi_row) = self.ofi_features.as_ref().and_then(|o| o.get(ofi_idx)) {
for &v in ofi_row.iter() {
fv_f32.push(v as f32);
}
} else {
flat_ofi.extend_from_slice(&[0.0_f32; 8]);
fv_f32.resize(feature_dim, 0.0);
}
}
self.val_ofi_gpu = Some(flat_ofi);
}
}
let flat_features = self.val_features_gpu.as_ref()
.ok_or_else(|| anyhow::anyhow!("val_features_gpu must be initialized"))?;
let (current_closes, next_closes) = self.val_closes_gpu.as_ref()
.ok_or_else(|| anyhow::anyhow!("val_closes_gpu must be initialized"))?;
// ── Build complete state tensor on CPU per validation call ──
// Portfolio features change every epoch (position, P&L), so we rebuild each call.
// Feature layout per row: [42 market features | 3 portfolio | 8 OFI (opt) | padding]
let val_price = self.val_data.first()
.map(|(fv, tgt)| if tgt.len() >= 2 { tgt[0] as f32 } else { fv[3] as f32 })
.unwrap_or(0.0);
let portfolio_f = self.portfolio_tracker.get_portfolio_features(val_price);
// State layout: [42 market | 8 portfolio | 16 MTF | 8 OFI | padding]
// We only have 42 features + 3 portfolio + 8 OFI = 53 values.
// Pad the rest with zeros to match state_dim_padded (pad128) per row.
// cuBLAS GemmEx reads with ldb=state_dim_padded (CUTLASS K-tile=128).
let state_dim_padded = (aligned_state_dim + 127) & !127;
let ofi_cols: usize = if self.val_ofi_gpu.is_some() { 8 } else { 0 };
let content_cols: usize = 42 + 3 + ofi_cols; // what we actually write
let pad_cols = state_dim_padded.saturating_sub(content_cols);
let mut state_flat: Vec<f32> = Vec::with_capacity(sample_size * state_dim_padded);
let feat_cols = flat_features.len() / sample_size; // 42
for i in 0..sample_size {
// Market features
let feat_start = i * feat_cols;
let feat_end = feat_start + feat_cols;
let feat_slice = flat_features.get(feat_start..feat_end)
.ok_or_else(|| anyhow::anyhow!("val features OOB at row {i}"))?;
state_flat.extend_from_slice(feat_slice);
// Portfolio features (same for all rows — broadcast)
state_flat.extend_from_slice(&portfolio_f[..3]);
// OFI features (optional, 8 floats per row)
if let Some(ref flat_ofi) = self.val_ofi_gpu {
let ofi_start = i * 8;
let ofi_end = ofi_start + 8;
let ofi_slice = flat_ofi.get(ofi_start..ofi_end)
.ok_or_else(|| anyhow::anyhow!("val OFI OOB at row {i}"))?;
state_flat.extend_from_slice(ofi_slice);
features.push(fv_f32);
}
// Alignment padding (zeros)
state_flat.extend(std::iter::repeat(0.0_f32).take(pad_cols));
}
let window_prices = vec![prices];
let window_features = vec![features];
// Upload state to GPU as raw CudaSlice<half::bf16> — no GpuTensor, no elementwise kernels.
let state_gpu = crate::cuda_pipeline::clone_htod_f32_to_bf16(stream, &state_flat)
.map_err(|e| anyhow::anyhow!("val state HtoD: {e}"))?;
// cuBLAS forward via fused_ctx — zero GpuTensor involvement
let fused = self.fused_ctx.as_mut()
.ok_or_else(|| anyhow::anyhow!("fused_ctx required for validation forward"))?;
let total_actions = fused.total_actions();
let q_out = fused.compute_q_values(&state_gpu, sample_size)
.map_err(|e| anyhow::anyhow!("cuBLAS validation forward: {e}"))?;
// Download Q-values to CPU for action selection + Sharpe computation.
// q_out is [config.batch_size, total_actions]; allocate to match, then truncate.
let mut host_q = vec![0.0_f32; q_out.len()];
crate::cuda_pipeline::dtoh_bf16_to_f32(stream, q_out, &mut host_q)
.map_err(|e| anyhow::anyhow!("Q-val DtoH: {e}"))?;
host_q.truncate(sample_size * total_actions);
// ── CPU greedy action selection ──
// Direction LUT maps exposure head index → position size direction.
// 0=Short100(-1.0), 1=Short50(-0.5), 2=Flat(0.0), 3=Long50(0.5), 4=Long100(1.0)
const DIRECTION_LUT: [f32; 9] = [-1.0, -0.75, -0.5, -0.25, 0.0, 0.25, 0.5, 0.75, 1.0];
let use_branching = true && total_actions >= 15;
let rewards: Vec<f32> = (0..sample_size).map(|i| {
let row_start = i * total_actions;
let row = &host_q[row_start..row_start + total_actions];
// Greedy exposure action selection (argmax over exposure head, 5 values)
let exposure_idx = if use_branching {
// Branching DQN: exposure head is first 9 Q-values
row.get(..9).unwrap_or(row)
.iter()
.enumerate()
.max_by(|a, b| a.1.partial_cmp(b.1).unwrap_or(std::cmp::Ordering::Equal))
.map(|(j, _)| j)
.expect("argmax must exist for non-empty Q-values")
} else {
// Non-branching: argmax over all Q-values maps to exposure
row.iter()
.enumerate()
.max_by(|a, b| a.1.partial_cmp(b.1).unwrap_or(std::cmp::Ordering::Equal))
.map(|(j, _)| j.min(4))
.expect("argmax must exist for non-empty Q-values")
let hp = &self.hyperparams;
#[allow(clippy::cast_possible_truncation)]
let config = GpuBacktestConfig {
max_position: hp.max_position_absolute as f32,
tx_cost_bps: hp.transaction_cost_multiplier as f32,
spread_cost: (hp.tick_size * hp.contract_multiplier * hp.fill_spread_cost_frac) as f32,
initial_capital: hp.initial_capital as f32,
contract_multiplier: hp.contract_multiplier as f32,
margin_pct: hp.margin_pct as f32,
max_leverage: 0.0, // Disabled: match training env (no leverage cap)
ofi_dim: if ofi_enabled { 8 } else { 0 },
min_hold_bars: hp.min_hold_bars as i32,
bars_per_day: hp.bars_per_day as f32,
trading_days_per_year: hp.trading_days_per_year as f32,
};
let direction = DIRECTION_LUT[exposure_idx.min(8)];
// PnL-based reward: direction × (next_close - cur_close) / cur_close
let cur = current_closes.get(i).copied().unwrap_or(0.0);
let nxt = next_closes.get(i).copied().unwrap_or(cur);
if cur.abs() > 1e-10 {
direction * (nxt - cur) / cur
} else {
0.0
}
}).collect();
let evaluator = GpuBacktestEvaluator::new(
&window_prices,
&window_features,
feature_dim,
config,
stream,
).map_err(|e| anyhow::anyhow!("GpuBacktestEvaluator init: {e}"))?;
// ── CPU Sharpe computation ──
let n = rewards.len() as f64;
let mean_scalar: f64 = rewards.iter().map(|&v| v as f64).sum::<f64>() / n;
let var_scalar: f64 = if n > 1.0 {
rewards.iter().map(|&v| {
let d = v as f64 - mean_scalar;
d * d
}).sum::<f64>() / n
self.gpu_evaluator = Some(evaluator);
}
// ── Extract current weights from agent ──
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 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 {
0.0
};
let std_val = var_scalar.sqrt();
let val_sharpe = if std_val > 1e-10 {
(mean_scalar / std_val) * n.sqrt()
} else {
0.0
crate::cuda_pipeline::gpu_weights::extract_dueling_weights(
vars, eval_stream,
).map_err(|e| anyhow::anyhow!("extract dueling weights: {e}"))?
};
self.set_epsilon(original_epsilon).await?;
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,
shared_h2: network_dims.1,
value_h: network_dims.2,
adv_h: network_dims.3,
num_atoms: hp.num_atoms,
branch_0_size: b0,
branch_1_size: if is_branching { b1 } else { b0 },
branch_2_size: if is_branching { b2 } else { b0 },
v_min: hp.v_min as f32,
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(),
&dqn_cfg,
).map_err(|e| anyhow::anyhow!("GPU backtest evaluate_dqn_graphed: {e}"))?;
// Single window — take its Sharpe directly
let val_sharpe = metrics.first()
.map(|m| m.sharpe as f64)
.unwrap_or(0.0);
Ok(-val_sharpe)
}

22
crates/ml/src/trainers/dqn/trainer/mod.rs
View File

@@ -275,12 +275,10 @@ pub struct DQNTrainer {
/// val_data[i] corresponds to ofi_features[ofi_val_offset + i].
pub(crate) ofi_val_offset: usize,
/// CPU-cached validation features flat vec [sample_size * 42] — built once, reused per epoch
pub(crate) val_features_gpu: Option<Vec<f32>>,
/// CPU-cached validation close prices (current, next) — built once per training run
pub(crate) val_closes_gpu: Option<(Vec<f32>, Vec<f32>)>,
/// CPU-cached validation OFI features flat vec [sample_size * 8] — built once
pub(crate) val_ofi_gpu: Option<Vec<f32>>,
/// GPU backtest evaluator for validation loss (replaces CPU argmax+Sharpe path).
/// Lazy-initialized on first `compute_validation_loss()` call from val_data.
/// Invalidated (set to None) when val_data changes (fold transitions).
pub(crate) gpu_evaluator: Option<crate::cuda_pipeline::gpu_backtest_evaluator::GpuBacktestEvaluator>,
// Phase C: Fill simulation and smart order routing
/// Fill simulator for order type-dependent execution modeling
@@ -487,9 +485,7 @@ impl DQNTrainer {
self.init_from_fxcache(&features, &targets, &[]).await?;
self.set_val_data_from_slices(&val_features, &val_targets, features.len());
self.val_features_gpu = None;
self.val_closes_gpu = None;
self.val_ofi_gpu = None;
self.gpu_evaluator = None;
self.train_fold_from_slices(&features, &targets, checkpoint_callback).await
}
@@ -522,9 +518,7 @@ impl DQNTrainer {
.map(|(f, t)| (*f, *t))
.collect();
self.val_features_gpu = None;
self.val_closes_gpu = None;
self.val_ofi_gpu = None;
self.gpu_evaluator = None;
self.train_with_data_full_loop_slices(&training_data, checkpoint_callback)
.await
@@ -1065,9 +1059,7 @@ impl DQNTrainer {
.map(|(f, t)| (*f, t.to_vec()))
.collect();
self.ofi_val_offset = ofi_val_offset;
self.val_features_gpu = None;
self.val_closes_gpu = None;
self.val_ofi_gpu = None;
self.gpu_evaluator = None;
}
/// Inject pre-uploaded GPU data (e.g. from a `DoubleBufferedLoader`).