fix(examples): evaluate_baseline bf16 closure types for GpuBacktestEvaluator

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-03-29 22:53:17 +02:00
parent 58e09a52ac
commit 9c17b2708d

View File

@@ -142,6 +142,7 @@ use ml::xlstm::{XLSTMConfig, XLSTMTrainableAdapter};
use ml_core::cuda_autograd::GpuTensor;
use std::sync::Arc;
use cudarc::driver::CudaStream;
use half::bf16;
/// Number of bars processed per GPU forward pass.
///
@@ -1240,7 +1241,7 @@ fn evaluate_dqn_fold_gpu(
let eval_stream = stream.clone();
evaluator
.evaluate(
&|states_flat: &cudarc::driver::CudaSlice<f32>, batch_size: usize, state_dim: usize| -> Result<cudarc::driver::CudaSlice<i32>, ml::MLError> {
&|states_flat: &cudarc::driver::CudaSlice<half::bf16>, batch_size: usize, state_dim: usize| -> Result<cudarc::driver::CudaSlice<i32>, ml::MLError> {
// Clone CudaSlice into a GpuTensor, forward through DQN, then argmax
let cloned = ml::cuda_pipeline::clone_cuda_slice_f32(states_flat, &eval_stream)?;
let states_tensor = GpuTensor::new(cloned, vec![batch_size, state_dim])
@@ -1479,13 +1480,14 @@ fn evaluate_ppo_fold_gpu(
let eval_ppo_stream = ppo_stream.clone();
let metrics = evaluator
.evaluate(
&|states_flat: &cudarc::driver::CudaSlice<f32>, batch_size: usize, state_dim: usize| -> Result<cudarc::driver::CudaSlice<i32>, ml::MLError> {
// Download states to host for PPO actor forward
&|states_flat: &cudarc::driver::CudaSlice<half::bf16>, batch_size: usize, state_dim: usize| -> Result<cudarc::driver::CudaSlice<i32>, ml::MLError> {
// Download bf16 states to host, then convert to f32 for PPO actor forward
let n_floats = batch_size * state_dim;
let view = states_flat.slice(..n_floats);
let mut host_states = vec![0.0_f32; n_floats];
eval_ppo_stream.memcpy_dtoh(&view, &mut host_states)
let mut host_bf16 = vec![bf16::ZERO; n_floats];
eval_ppo_stream.memcpy_dtoh(&view, &mut host_bf16)
.map_err(|e| ml::MLError::ModelError(format!("DtoH states: {e}")))?;
let host_states: Vec<f32> = host_bf16.iter().map(|v| v.to_f32()).collect();
// Get action probabilities [batch * 45]
let probs_host = match &ppo.actor {
@@ -1497,18 +1499,20 @@ fn evaluate_ppo_fold_gpu(
}
};
// Upload probs to GPU and collapse 45→5 exposure scores
let mut probs_gpu = eval_ppo_stream.alloc_zeros::<f32>(probs_host.len())
// Upload probs as bf16 to GPU and collapse 45→5 exposure scores
let probs_bf16: Vec<bf16> = probs_host.iter().map(|v| bf16::from_f32(*v)).collect();
let mut probs_gpu = eval_ppo_stream.alloc_zeros::<bf16>(probs_bf16.len())
.map_err(|e| ml::MLError::ModelError(format!("alloc probs: {e}")))?;
eval_ppo_stream.memcpy_htod(&probs_host, &mut probs_gpu)
eval_ppo_stream.memcpy_htod(&probs_bf16, &mut probs_gpu)
.map_err(|e| ml::MLError::ModelError(format!("HtoD probs: {e}")))?;
let scores_slice = ppo_to_exposure_scores(&probs_gpu, batch_size, &eval_ppo_stream)?;
// Argmax over 5 exposure scores per batch element
let mut host_scores = vec![0.0_f32; batch_size * 5];
eval_ppo_stream.memcpy_dtoh(&scores_slice, &mut host_scores)
// Argmax over 5 exposure scores per batch element (download bf16, convert to f32)
let mut host_scores_bf16 = vec![bf16::ZERO; batch_size * 5];
eval_ppo_stream.memcpy_dtoh(&scores_slice, &mut host_scores_bf16)
.map_err(|e| ml::MLError::ModelError(format!("DtoH scores: {e}")))?;
let host_scores: Vec<f32> = host_scores_bf16.iter().map(|v| v.to_f32()).collect();
let mut actions = Vec::with_capacity(batch_size);
for b in 0..batch_size {
let offset = b * 5;
@@ -1648,13 +1652,14 @@ fn evaluate_supervised_fold_gpu(
let metrics = evaluator
.evaluate(
&|states_flat: &cudarc::driver::CudaSlice<f32>, batch_size: usize, state_dim: usize| -> Result<cudarc::driver::CudaSlice<i32>, ml::MLError> {
// Download states to host
&|states_flat: &cudarc::driver::CudaSlice<half::bf16>, batch_size: usize, state_dim: usize| -> Result<cudarc::driver::CudaSlice<i32>, ml::MLError> {
// Download bf16 states to host, then convert to f32 for supervised forward
let n_floats = batch_size * state_dim;
let view = states_flat.slice(..n_floats);
let mut host_states = vec![0.0_f32; n_floats];
eval_sup_stream.memcpy_dtoh(&view, &mut host_states)
let mut host_bf16 = vec![bf16::ZERO; n_floats];
eval_sup_stream.memcpy_dtoh(&view, &mut host_bf16)
.map_err(|e| ml::MLError::ModelError(format!("DtoH states: {e}")))?;
let host_states: Vec<f32> = host_bf16.iter().map(|v| v.to_f32()).collect();
// Extract market features only (strip portfolio dims)
let market_dim = state_dim.saturating_sub(3);
@@ -1690,20 +1695,22 @@ fn evaluate_supervised_fold_gpu(
signals.push(signal);
}
// Upload signals to GPU and convert to 5-exposure action scores
let mut signal_gpu = eval_sup_stream.alloc_zeros::<f32>(signals.len())
// Upload signals as bf16 to GPU and convert to 5-exposure action scores
let signals_bf16: Vec<bf16> = signals.iter().map(|v| bf16::from_f32(*v)).collect();
let mut signal_gpu = eval_sup_stream.alloc_zeros::<bf16>(signals_bf16.len())
.map_err(|e| ml::MLError::ModelError(format!("alloc signals: {e}")))?;
eval_sup_stream.memcpy_htod(&signals, &mut signal_gpu)
eval_sup_stream.memcpy_htod(&signals_bf16, &mut signal_gpu)
.map_err(|e| ml::MLError::ModelError(format!("HtoD signals: {e}")))?;
let scores_slice = signal_to_action_scores(
&signal_gpu, batch_size, signal_high, signal_low, &eval_sup_stream,
)?;
// Argmax over 5 exposure scores
let mut host_scores = vec![0.0_f32; batch_size * 5];
eval_sup_stream.memcpy_dtoh(&scores_slice, &mut host_scores)
// Argmax over 5 exposure scores (download bf16, convert to f32)
let mut host_scores_bf16 = vec![bf16::ZERO; batch_size * 5];
eval_sup_stream.memcpy_dtoh(&scores_slice, &mut host_scores_bf16)
.map_err(|e| ml::MLError::ModelError(format!("DtoH scores: {e}")))?;
let host_scores: Vec<f32> = host_scores_bf16.iter().map(|v| v.to_f32()).collect();
let mut actions = Vec::with_capacity(batch_size);
for b in 0..batch_size {
let offset = b * 5;