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

feat(ml): BF16 CUDA pipeline (replay buffer, weights, data upload)

Browse Source 
- 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 <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-03-03 16:14:31 +01:00
parent 4accdded0f
commit d0dd3af2f1
3 changed files with 37 additions and 10 deletions
Download Patch File Download Diff File Expand all files Collapse all files

10
crates/ml/src/cuda_pipeline/gpu_replay_buffer.rs
View File

@@ -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;

4
crates/ml/src/cuda_pipeline/gpu_weights.rs
View File

@@ -176,6 +176,8 @@ fn extract_one(
let flat: Vec<f32> = 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::<f32>()
.map_err(|e| MLError::ModelError(format!("to_vec1 {name}: {e}")))?;
let mut buf = stream
@@ -201,6 +203,8 @@ fn sync_one(
let flat: Vec<f32> = 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::<f32>()
.map_err(|e| MLError::ModelError(format!("to_vec1 {name}: {e}")))?;
stream

33
crates/ml/src/cuda_pipeline/mod.rs
View File

@@ -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::<f32>()
.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,