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

perf(cuda): make GPU eval default, eliminate all mid-loop GPU→CPU transfers

Browse Source 
- Flip --gpu-eval default to true (--no-gpu-eval to opt out)
- Move actions_history scatter-write into env kernel (zero CPU accumulation)
- Remove done-flag periodic download (env kernel handles per-thread)
- Only GPU→CPU transfer is final metrics readback (n_windows × 40 bytes)
- Add spread cost discrepancy warning when GPU path is active

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-03-11 15:03:03 +01:00
parent 0c85dcf56a
commit e4870ea5e9
3 changed files with 29 additions and 49 deletions
Download Patch File Download Diff File Expand all files Collapse all files

20
crates/ml/examples/evaluate_baseline.rs
View File

@@ -146,14 +146,16 @@ struct Args {
#[arg(long)]
hyperopt_params: Option<PathBuf>,
/// Use GPU-accelerated backtest evaluation for DQN (uses greedy argmax, not softmax).
/// Use GPU-accelerated backtest evaluation (default: enabled).
///
/// When enabled and CUDA is available, DQN folds are evaluated using
/// `GpuBacktestEvaluator` — all bars processed on GPU with a single metrics
/// readback. Results differ slightly from the default CPU path because the
/// GPU path uses greedy argmax while the CPU path uses hierarchical softmax.
/// readback. Results differ slightly from the CPU path because the GPU path
/// uses greedy argmax while the CPU path uses hierarchical softmax.
/// Falls back to the CPU path on any GPU error.
#[arg(long, default_value_t = false)]
///
/// Pass `--no-gpu-eval` to force the CPU path.
#[arg(long, default_value_t = true)]
gpu_eval: bool,
/// Initial capital for GPU backtest evaluator (used only with --gpu-eval).
@@ -824,6 +826,11 @@ fn evaluate_dqn_fold_gpu(
ckpt_path.display(),
device,
);
warn!(
" [DQN GPU] Spread cost uses constant tick_size*spread_ticks={:.6}, \
CPU path uses per-bar estimate — results will differ slightly",
args.tick_size * args.spread_ticks,
);
// ── Build single window from all test data ────────────────────────────
//
@@ -1268,9 +1275,10 @@ fn main() -> Result<()> {
if eval_dqn {
let hp = load_hyperopt_params(&args.hyperopt_params, "dqn");
// GPU path: when --gpu-eval is set and CUDA is compiled in, try GPU first.
// GPU path (default): when CUDA is compiled in, try GPU first.
// The GPU path uses greedy argmax (not softmax), so results differ slightly.
// On any GPU error, fall through to the CPU path below.
// Pass --no-gpu-eval to skip the GPU path entirely.
#[cfg(feature = "cuda")]
let gpu_handled = if args.gpu_eval {
info!(" [DQN] Attempting GPU-accelerated evaluation (greedy argmax)...");
@@ -1345,7 +1353,7 @@ fn main() -> Result<()> {
}
}
} else {
false // --gpu-eval not set; use CPU path
false // --no-gpu-eval was set; use CPU path
};
// When compiled without CUDA, gpu_handled is always false.

2
crates/ml/src/cuda_pipeline/backtest_env_kernel.cu
View File

@@ -28,6 +28,7 @@ extern "C" __global__ void backtest_env_step(
float* step_rewards, // [n_windows]
float* step_returns, // [n_windows * max_len] (accumulated)
int* done_flags, // [n_windows]
int* actions_history, // [n_windows * max_len] (accumulated, for metrics)
// Config
int n_windows,
@@ -120,4 +121,5 @@ extern "C" __global__ void backtest_env_step(
// Outputs
step_rewards[w] = step_ret;
step_returns[w * max_len + current_step] = step_ret;
actions_history[w * max_len + current_step] = actions[w];
}

56
crates/ml/src/cuda_pipeline/gpu_backtest_evaluator.rs
View File

@@ -8,7 +8,8 @@
//! 3. Metrics reduction kernel → single readback
//!
//! The only GPU→CPU transfer is the final metrics readback (n_windows × 10 floats).
//! Per-step state construction uses a zero-copy DtoD path (gather kernel → Candle tensor).
//! All per-step data (states, actions, rewards, action history) stays GPU-resident.
//! Done-flags are checked per-thread inside the env kernel (no mid-loop downloads).
use std::sync::Arc;
use candle_core::cuda_backend::cudarc;
@@ -115,10 +116,6 @@ pub struct GpuBacktestEvaluator {
// Output buffer (written by metrics kernel)
metrics_buf: CudaSlice<f32>, // [n_windows * 10]
// CPU-side action history accumulated during the step loop.
// Uploaded once before the metrics kernel launch.
actions_history_cpu: Vec<i32>, // [n_windows * max_len]
// Dimensions and config
n_windows: usize,
max_len: usize,
@@ -291,7 +288,6 @@ impl GpuBacktestEvaluator {
actions_history_buf,
states_buf,
metrics_buf,
actions_history_cpu: vec![0_i32; n_windows * max_len],
n_windows,
max_len,
feature_dim,
@@ -457,8 +453,7 @@ impl GpuBacktestEvaluator {
// DtoD copy: argmax output (U32 Candle tensor) → actions_buf (CudaSlice<i32>).
// Action values are 0..N_ACTIONS (small non-negative), so u32 and i32 share
// identical bit patterns — raw byte reinterpret is safe. Eliminates the old
// to_vec1 → collect::<Vec<i32>> → memcpy_htod upload roundtrip on the hot path.
// identical bit patterns — raw byte reinterpret is safe.
{
let (act_guard, _act_layout) = actions_tensor.storage_and_layout();
match &*act_guard {
@@ -489,17 +484,9 @@ impl GpuBacktestEvaluator {
drop(act_guard);
}
// Accumulate actions into CPU-side history (uploaded once before metrics kernel).
// Download is n_windows × 4 bytes — negligible. The old path did the same
// download (to_vec1) PLUS an upload (memcpy_htod); we keep only the download.
let actions_u32: Vec<u32> = actions_tensor
.to_vec1()
.map_err(|e| MLError::ModelError(format!("argmax to_vec1 step {step}: {e}")))?;
for (w, &a) in actions_u32.iter().enumerate() {
self.actions_history_cpu[w * self.max_len + step] = a as i32;
}
// 4. Launch env step kernel — one thread per window
// The kernel also scatter-writes actions into actions_history_buf
// (GPU-resident, no CPU roundtrip).
let grid = ((self.n_windows + 255) / 256) as u32;
let env_cfg = LaunchConfig {
grid_dim: (grid.max(1), 1, 1),
@@ -512,8 +499,8 @@ impl GpuBacktestEvaluator {
// Safety: argument order matches `backtest_env_step` signature exactly:
// prices, window_lens, actions, portfolio_state, step_rewards,
// step_returns, done_flags, n_windows, max_len, max_position,
// tx_cost_bps, spread_cost, current_step
// step_returns, done_flags, actions_history, n_windows, max_len,
// max_position, tx_cost_bps, spread_cost, current_step
unsafe {
self.stream
.launch_builder(&self.env_kernel)
@@ -524,6 +511,7 @@ impl GpuBacktestEvaluator {
.arg(&self.step_rewards_buf)
.arg(&self.step_returns_buf)
.arg(&self.done_buf)
.arg(&self.actions_history_buf)
.arg(&n_win_i32)
.arg(&max_len_i32)
.arg(&self.config.max_position)
@@ -536,31 +524,13 @@ impl GpuBacktestEvaluator {
})?;
}
// Periodic early-exit check (every 100 steps to amortise the download cost)
if step % 100 == 99 {
let mut done_host = vec![0_i32; self.n_windows];
self.stream
.memcpy_dtoh(&self.done_buf, &mut done_host)
.map_err(|e| {
MLError::ModelError(format!("done check download step {step}: {e}"))
})?;
if done_host.iter().all(|&d| d != 0) {
info!(
"GpuBacktestEvaluator: all {} windows done at step {}",
self.n_windows,
step + 1,
);
break;
}
}
// No mid-loop GPU→CPU transfers. The env kernel checks done_flags per-thread
// and early-returns (no-op) for finished windows, so running all max_len steps
// is safe and avoids the latency of periodic done-flag downloads.
}
// 5. Upload accumulated action history once before metrics kernel
self.stream
.memcpy_htod(&self.actions_history_cpu, &mut self.actions_history_buf)
.map_err(|e| MLError::ModelError(format!("actions_history upload: {e}")))?;
// 6. Launch metrics reduction kernel — one block per window
// 5. Launch metrics reduction kernel — one block per window
// actions_history_buf was populated by the env step kernel (no upload needed).
// Shared memory: 6 reduction arrays × 256 threads × 4 bytes = 6 KB
// + 4096 floats for bitonic sort scratch = 16 KB
// Total = 22 KB (well within the 48 KB L1/shmem limit)