diff --git a/crates/ml/examples/evaluate_baseline.rs b/crates/ml/examples/evaluate_baseline.rs index c9afde925..22e9266ea 100644 --- a/crates/ml/examples/evaluate_baseline.rs +++ b/crates/ml/examples/evaluate_baseline.rs @@ -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, batch_size: usize, state_dim: usize| -> Result, ml::MLError> { + &|states_flat: &cudarc::driver::CudaSlice, batch_size: usize, state_dim: usize| -> Result, 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, batch_size: usize, state_dim: usize| -> Result, ml::MLError> { - // Download states to host for PPO actor forward + &|states_flat: &cudarc::driver::CudaSlice, batch_size: usize, state_dim: usize| -> Result, 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 = 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::(probs_host.len()) + // Upload probs as bf16 to GPU and collapse 45→5 exposure scores + let probs_bf16: Vec = probs_host.iter().map(|v| bf16::from_f32(*v)).collect(); + let mut probs_gpu = eval_ppo_stream.alloc_zeros::(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 = 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, batch_size: usize, state_dim: usize| -> Result, ml::MLError> { - // Download states to host + &|states_flat: &cudarc::driver::CudaSlice, batch_size: usize, state_dim: usize| -> Result, 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 = 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::(signals.len()) + // Upload signals as bf16 to GPU and convert to 5-exposure action scores + let signals_bf16: Vec = signals.iter().map(|v| bf16::from_f32(*v)).collect(); + let mut signal_gpu = eval_sup_stream.alloc_zeros::(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 = 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;