From d0dd3af2f13972c05038f1d340fac9fd7ec8bedf Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Tue, 3 Mar 2026 16:14:31 +0100 Subject: [PATCH] feat(ml): BF16 CUDA pipeline (replay buffer, weights, data upload) - gpu_replay_buffer: allocate states/next_states with training_dtype(), cast incoming batches at ingestion boundary - gpu_weights: cast BF16 model weights to F32 before extraction for CUDA f32 kernels in both extract_one() and sync_one() - mod.rs: cast DqnGpuData, GpuBufferPool, PpoGpuData uploads to training_dtype(); cast portfolio tensors to match features dtype; cast bar_target_values readback to F32 Co-Authored-By: Claude Opus 4.6 --- .../ml/src/cuda_pipeline/gpu_replay_buffer.rs | 10 ++++-- crates/ml/src/cuda_pipeline/gpu_weights.rs | 4 +++ crates/ml/src/cuda_pipeline/mod.rs | 33 ++++++++++++++----- 3 files changed, 37 insertions(+), 10 deletions(-) diff --git a/crates/ml/src/cuda_pipeline/gpu_replay_buffer.rs b/crates/ml/src/cuda_pipeline/gpu_replay_buffer.rs index d7cd61c22..30101c100 100644 --- a/crates/ml/src/cuda_pipeline/gpu_replay_buffer.rs +++ b/crates/ml/src/cuda_pipeline/gpu_replay_buffer.rs @@ -6,6 +6,7 @@ use candle_core::{Device, DType, Tensor}; use crate::MLError; +use crate::dqn::mixed_precision::training_dtype; use crate::dqn::replay_buffer_type::GpuBatch; /// Configuration for GPU replay buffer @@ -71,8 +72,9 @@ impl GpuReplayBuffer { ))); } - let states = Tensor::zeros(&[cap, sdim], DType::F32, device)?; - let next_states = Tensor::zeros(&[cap, sdim], DType::F32, device)?; + let dtype = training_dtype(device); + let states = Tensor::zeros(&[cap, sdim], dtype, device)?; + let next_states = Tensor::zeros(&[cap, sdim], dtype, device)?; let actions = Tensor::zeros(&[cap], DType::U32, device)?; let rewards = Tensor::zeros(&[cap], DType::F32, device)?; let dones = Tensor::zeros(&[cap], DType::F32, device)?; @@ -177,6 +179,10 @@ impl GpuReplayBuffer { return Ok(()); } + // Cast incoming states to match buffer dtype (e.g. F32 → BF16 on CUDA) + let states = &states.to_dtype(self.states.dtype())?; + let next_states = &next_states.to_dtype(self.next_states.dtype())?; + let cap = self.config.capacity; let cursor = self.write_cursor; diff --git a/crates/ml/src/cuda_pipeline/gpu_weights.rs b/crates/ml/src/cuda_pipeline/gpu_weights.rs index b47314dfe..99a84978f 100644 --- a/crates/ml/src/cuda_pipeline/gpu_weights.rs +++ b/crates/ml/src/cuda_pipeline/gpu_weights.rs @@ -176,6 +176,8 @@ fn extract_one( let flat: Vec = tensor .flatten_all() .map_err(|e| MLError::ModelError(format!("Flatten {name}: {e}")))? + .to_dtype(candle_core::DType::F32) + .map_err(|e| MLError::ModelError(format!("Cast {name} to F32: {e}")))? .to_vec1::() .map_err(|e| MLError::ModelError(format!("to_vec1 {name}: {e}")))?; let mut buf = stream @@ -201,6 +203,8 @@ fn sync_one( let flat: Vec = tensor .flatten_all() .map_err(|e| MLError::ModelError(format!("Flatten {name}: {e}")))? + .to_dtype(candle_core::DType::F32) + .map_err(|e| MLError::ModelError(format!("Cast {name} to F32: {e}")))? .to_vec1::() .map_err(|e| MLError::ModelError(format!("to_vec1 {name}: {e}")))?; stream diff --git a/crates/ml/src/cuda_pipeline/mod.rs b/crates/ml/src/cuda_pipeline/mod.rs index 44da1a49b..95e9af46b 100644 --- a/crates/ml/src/cuda_pipeline/mod.rs +++ b/crates/ml/src/cuda_pipeline/mod.rs @@ -7,6 +7,7 @@ use candle_core::{Device, Tensor}; use crate::MLError; +use crate::dqn::mixed_precision::training_dtype; pub mod double_buffer; pub mod multi_gpu; @@ -130,10 +131,14 @@ impl DqnGpuData { } let features = Tensor::from_vec(flat_features, (num_bars, feature_dim), device) - .map_err(|e| MLError::ModelError(format!("GPU feature upload failed: {e}")))?; + .map_err(|e| MLError::ModelError(format!("GPU feature upload failed: {e}")))? + .to_dtype(training_dtype(device)) + .map_err(|e| MLError::ModelError(format!("GPU feature dtype cast failed: {e}")))?; let targets = Tensor::from_vec(flat_targets, (num_bars, target_dim), device) - .map_err(|e| MLError::ModelError(format!("GPU target upload failed: {e}")))?; + .map_err(|e| MLError::ModelError(format!("GPU target upload failed: {e}")))? + .to_dtype(training_dtype(device)) + .map_err(|e| MLError::ModelError(format!("GPU target dtype cast failed: {e}")))?; Ok(Self { features, @@ -185,6 +190,8 @@ impl DqnGpuData { let slice = self.bar_targets(bar_idx)? .flatten_all() .map_err(|e| MLError::ModelError(format!("Target flatten failed: {e}")))? + .to_dtype(candle_core::DType::F32) + .map_err(|e| MLError::ModelError(format!("Target cast to F32 failed: {e}")))? .to_vec1::() .map_err(|e| MLError::ModelError(format!("Target to_vec1 failed: {e}")))?; Ok([ @@ -208,7 +215,9 @@ impl DqnGpuData { portfolio_features.to_vec(), (1, 3), device, - ).map_err(|e| MLError::ModelError(format!("Portfolio tensor failed: {e}")))?; + ).map_err(|e| MLError::ModelError(format!("Portfolio tensor failed: {e}")))? + .to_dtype(self.features.dtype()) + .map_err(|e| MLError::ModelError(format!("Portfolio dtype cast failed: {e}")))?; Tensor::cat(&[&market, &portfolio], 1) .map_err(|e| MLError::ModelError(format!("State cat failed: {e}"))) @@ -242,12 +251,14 @@ impl DqnGpuData { .narrow(0, start, count) .map_err(|e| MLError::ModelError(format!("Batch feature slice failed: {e}")))?; - // [1, 3] portfolio features — single small upload + // [1, 3] portfolio features — single small upload, cast to match features dtype let portfolio = Tensor::from_vec( portfolio_features.to_vec(), (1, 3), device, - ).map_err(|e| MLError::ModelError(format!("Portfolio tensor failed: {e}")))?; + ).map_err(|e| MLError::ModelError(format!("Portfolio tensor failed: {e}")))? + .to_dtype(self.features.dtype()) + .map_err(|e| MLError::ModelError(format!("Portfolio dtype cast failed: {e}")))?; // Broadcast [1, 3] → [count, 3] then concatenate with [count, 51] → [count, 54] let portfolio_broadcast = portfolio @@ -348,14 +359,18 @@ impl GpuBufferPool { (num_bars, self.feature_dim), device, ) - .map_err(|e| MLError::ModelError(format!("GPU feature upload failed: {e}")))?; + .map_err(|e| MLError::ModelError(format!("GPU feature upload failed: {e}")))? + .to_dtype(training_dtype(device)) + .map_err(|e| MLError::ModelError(format!("GPU feature dtype cast failed: {e}")))?; let targets = Tensor::from_slice( &self.target_buf[..targ_len], (num_bars, self.target_dim), device, ) - .map_err(|e| MLError::ModelError(format!("GPU target upload failed: {e}")))?; + .map_err(|e| MLError::ModelError(format!("GPU target upload failed: {e}")))? + .to_dtype(training_dtype(device)) + .map_err(|e| MLError::ModelError(format!("GPU target dtype cast failed: {e}")))?; Ok(DqnGpuData { features, @@ -416,7 +431,9 @@ impl PpoGpuData { } let states = Tensor::from_vec(flat_states, (num_steps, state_dim), device) - .map_err(|e| MLError::ModelError(format!("GPU state upload failed: {e}")))?; + .map_err(|e| MLError::ModelError(format!("GPU state upload failed: {e}")))? + .to_dtype(training_dtype(device)) + .map_err(|e| MLError::ModelError(format!("GPU state dtype cast failed: {e}")))?; Ok(Self { states,