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

refactor: mod.rs bf16 helpers → direct f32 transfers, legacy aliases for incremental migration

Browse Source 
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-04-10 18:06:25 +02:00
parent eb2c16acb9
commit 8fa301dda8
1 changed files with 96 additions and 103 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -5,7 +5,7 @@
//! pre-uploaded buffers instead of creating per-step copies.
//!
//! All modules operate on raw `CudaSlice` via cudarc.
//! `DqnGpuData` and `PpoGpuData` hold GPU-resident `CudaSlice<half::bf16>` buffers.
//! `DqnGpuData` and `PpoGpuData` hold GPU-resident `CudaSlice<f32>` buffers.
use std::sync::Arc;
use cudarc::driver::{CudaSlice, CudaStream, DevicePtr, DevicePtrMut};
@@ -52,72 +52,74 @@ pub fn align_to_tensor_cores(dim: usize) -> usize {
/// Estimate VRAM usage in bytes for pre-uploaded f32 data.
pub fn estimate_vram_bytes(num_elements: usize) -> usize {
num_elements * std::mem::size_of::<half::bf16>()
num_elements * std::mem::size_of::<f32>()
}
// ---------------------------------------------------------------------------
// BF16 ↔ F32 host-side conversion helpers
// F32 host ↔ device transfer helpers
// ---------------------------------------------------------------------------
// All GPU buffers are CudaSlice<half::bf16>. Host code works in f32.
// These helpers bridge the boundary at every HtoD / DtoH transfer.
// All GPU buffers are CudaSlice<f32>. Host code works in f32.
// Direct transfers — no bf16 conversion.
/// Convert an f32 slice to a Vec of bf16 for GPU upload.
#[inline]
pub fn f32_slice_to_bf16(src: &[f32]) -> Vec<half::bf16> {
src.iter().map(|&x| half::bf16::from_f32(x)).collect()
}
/// Upload f32 host data into an existing `CudaSlice<half::bf16>`.
///
/// Converts f32 → bf16 on the host, then calls `memcpy_htod`.
pub fn htod_f32_to_bf16(
/// Upload f32 host data into an existing `CudaSlice<f32>`.
pub fn htod_f32(
stream: &Arc<CudaStream>,
src: &[f32],
dst: &mut CudaSlice<half::bf16>,
dst: &mut CudaSlice<f32>,
) -> Result<(), MLError> {
let bf16_data = f32_slice_to_bf16(src);
stream.memcpy_htod(&bf16_data, dst)
.map_err(|e| MLError::ModelError(format!("htod_f32_to_bf16: {e}")))
stream.memcpy_htod(src, dst)
.map_err(|e| MLError::ModelError(format!("htod_f32: {e}")))
}
/// Upload f32 host data to a new `CudaSlice<half::bf16>` (clone_htod equivalent).
pub fn clone_htod_f32_to_bf16(
/// Upload f32 host data to a new `CudaSlice<f32>` (clone_htod).
pub fn clone_htod_f32(
stream: &Arc<CudaStream>,
src: &[f32],
) -> Result<CudaSlice<half::bf16>, MLError> {
let bf16_data = f32_slice_to_bf16(src);
stream.clone_htod(&bf16_data)
.map_err(|e| MLError::ModelError(format!("clone_htod_f32_to_bf16: {e}")))
) -> Result<CudaSlice<f32>, MLError> {
stream.clone_htod(src)
.map_err(|e| MLError::ModelError(format!("clone_htod_f32: {e}")))
}
/// Download `CudaSlice<half::bf16>` (or `CudaView<half::bf16>`) to an f32 host buffer.
/// Download `CudaSlice<f32>` (or `CudaView<f32>`) to an f32 host buffer.
///
/// Allocates a temporary bf16 host buffer, downloads, then converts to f32.
pub fn dtoh_bf16_to_f32<Src: cudarc::driver::DevicePtr<half::bf16>>(
/// Uses raw cuStreamSynchronize + cuMemcpyDtoH_v2 for CUDA-graph safety
/// (avoids cudarc event tracking which can deadlock inside captured graphs).
pub fn dtoh_f32<Src: cudarc::driver::DevicePtr<f32>>(
stream: &Arc<CudaStream>,
src: &Src,
dst: &mut [f32],
) -> Result<(), MLError> {
let n = dst.len();
let mut bf16_buf = vec![half::bf16::ZERO; n];
let num_bytes = n * std::mem::size_of::<half::bf16>();
let num_bytes = n * std::mem::size_of::<f32>();
let (src_ptr, _guard) = src.device_ptr(stream);
// Sync + raw DtoH to avoid cudarc event tracking (graph-safe)
#[allow(unsafe_code)]
unsafe {
cudarc::driver::sys::cuStreamSynchronize(stream.cu_stream());
cudarc::driver::sys::cuMemcpyDtoH_v2(
bf16_buf.as_mut_ptr().cast(),
dst.as_mut_ptr().cast(),
src_ptr,
num_bytes,
);
}
for (d, &s) in dst.iter_mut().zip(bf16_buf.iter()) {
*d = s.to_f32();
}
Ok(())
}
// Legacy aliases — delegate to the f32 functions above.
// Keeps the 188 existing call sites compiling during incremental migration.
#[inline]
pub fn htod_f32_to_bf16(stream: &Arc<CudaStream>, src: &[f32], dst: &mut CudaSlice<f32>) -> Result<(), MLError> {
htod_f32(stream, src, dst)
}
#[inline]
pub fn clone_htod_f32_to_bf16(stream: &Arc<CudaStream>, src: &[f32]) -> Result<CudaSlice<f32>, MLError> {
clone_htod_f32(stream, src)
}
#[inline]
pub fn dtoh_bf16_to_f32<Src: cudarc::driver::DevicePtr<f32>>(stream: &Arc<CudaStream>, src: &Src, dst: &mut [f32]) -> Result<(), MLError> {
dtoh_f32(stream, src, dst)
}
/// Compute optimal (grid_dim, block_dim) for a 1-D kernel launch.
///
/// Selects the largest block size that divides 32 (warp size) and does not
@@ -167,21 +169,21 @@ pub fn optimal_launch_dims(n_items: u32, max_threads_per_block: u32) -> (u32, u3
/// The fused experience collector kernel (4490 lines, branching+C51+NoisyNets)
// compile_ptx_for_device re-export removed: all kernels now use precompiled cubins.
/// Clone a `CudaSlice<half::bf16>` via DtoD memcpy.
/// Clone a `CudaSlice<f32>` via DtoD memcpy.
///
/// The caller owns the returned slice independently of the source.
pub fn clone_cuda_slice_f32(
src: &CudaSlice<half::bf16>,
src: &CudaSlice<f32>,
stream: &Arc<CudaStream>,
) -> Result<CudaSlice<half::bf16>, crate::MLError> {
) -> Result<CudaSlice<f32>, crate::MLError> {
let n = src.len();
let mut dst = stream.alloc_zeros::<half::bf16>(n).map_err(|e| {
let mut dst = stream.alloc_zeros::<f32>(n).map_err(|e| {
crate::MLError::ModelError(format!("clone_cuda_slice_f32 alloc: {e}"))
})?;
{
let (src_ptr, _src_guard) = src.device_ptr(stream);
let (dst_ptr, _dst_guard) = dst.device_ptr_mut(stream);
let num_bytes = n * std::mem::size_of::<half::bf16>();
let num_bytes = n * std::mem::size_of::<f32>();
#[allow(unsafe_code)]
unsafe {
cudarc::driver::result::memcpy_dtod_async(
@@ -218,16 +220,16 @@ pub fn clone_cuda_slice_u32(
/// Pre-uploaded GPU training data for DQN trainer.
///
/// Holds market features [N * 42] and target prices [N * 4] as GPU-resident
/// `CudaSlice<half::bf16>` buffers. Portfolio features (3 dims) are computed per-bar
/// `CudaSlice<f32>` buffers. Portfolio features (3 dims) are computed per-bar
/// by the trainer and concatenated on-device. OFI features (8 dims from MBP-10
/// order book) are optionally uploaded and appended to the state vector.
pub struct DqnGpuData {
/// Market features [num_bars * 42] on GPU (f32, row-major)
pub features: CudaSlice<half::bf16>,
pub features: CudaSlice<f32>,
/// Target prices [num_bars * 4] on GPU (f32, row-major)
pub targets: CudaSlice<half::bf16>,
pub targets: CudaSlice<f32>,
/// OFI features [num_bars * 8] on GPU (f32), from MBP-10 data
pub ofi_features: Option<CudaSlice<half::bf16>>,
pub ofi_features: Option<CudaSlice<f32>>,
/// Number of training bars
pub num_bars: usize,
/// Feature dimension (42)
@@ -415,19 +417,19 @@ impl DqnGpuData {
/// D2D subrange copy: extract `count` f32 elements at `offset` into a new `CudaSlice`.
pub(crate) fn d2d_subrange(
src: &CudaSlice<half::bf16>,
src: &CudaSlice<f32>,
offset: usize,
count: usize,
stream: &Arc<CudaStream>,
) -> Result<CudaSlice<half::bf16>, MLError> {
let mut dst = stream.alloc_zeros::<half::bf16>(count).map_err(|e| {
) -> Result<CudaSlice<f32>, MLError> {
let mut dst = stream.alloc_zeros::<f32>(count).map_err(|e| {
MLError::ModelError(format!("d2d_subrange alloc: {e}"))
})?;
{
let src_view = src.slice(offset..offset + count);
let (src_ptr, _sg) = src_view.device_ptr(stream);
let (dst_ptr, _dg) = dst.device_ptr_mut(stream);
let nbytes = count * std::mem::size_of::<half::bf16>();
let nbytes = count * std::mem::size_of::<f32>();
#[allow(unsafe_code)]
unsafe {
cudarc::driver::result::memcpy_dtod_async(
@@ -439,7 +441,7 @@ impl DqnGpuData {
}
/// Get market features for a single bar as a [42] `CudaSlice` (D2D copy from GPU buffer).
pub fn bar_features(&self, bar_idx: usize, stream: &Arc<CudaStream>) -> Result<CudaSlice<half::bf16>, MLError> {
pub fn bar_features(&self, bar_idx: usize, stream: &Arc<CudaStream>) -> Result<CudaSlice<f32>, MLError> {
if bar_idx >= self.num_bars {
return Err(MLError::ModelError(format!("Bar {bar_idx} out of range (num_bars={})", self.num_bars)));
}
@@ -447,13 +449,13 @@ impl DqnGpuData {
}
/// Get market features for a batch of bars as a [batch_size * 42] `CudaSlice`.
pub fn batch_features(&self, start: usize, count: usize, stream: &Arc<CudaStream>) -> Result<CudaSlice<half::bf16>, MLError> {
pub fn batch_features(&self, start: usize, count: usize, stream: &Arc<CudaStream>) -> Result<CudaSlice<f32>, MLError> {
let count = count.min(self.num_bars.saturating_sub(start));
Self::d2d_subrange(&self.features, start * self.feature_dim, count * self.feature_dim, stream)
}
/// Get target prices for a single bar as a [4] `CudaSlice`.
pub fn bar_targets(&self, bar_idx: usize, stream: &Arc<CudaStream>) -> Result<CudaSlice<half::bf16>, MLError> {
pub fn bar_targets(&self, bar_idx: usize, stream: &Arc<CudaStream>) -> Result<CudaSlice<f32>, MLError> {
if bar_idx >= self.num_bars {
return Err(MLError::ModelError(format!("Bar {bar_idx} target out of range (num_bars={})", self.num_bars)));
}
@@ -461,16 +463,16 @@ impl DqnGpuData {
}
/// Get target prices for a batch as [batch_size * 4] `CudaSlice`.
pub fn batch_targets(&self, start: usize, count: usize, stream: &Arc<CudaStream>) -> Result<CudaSlice<half::bf16>, MLError> {
pub fn batch_targets(&self, start: usize, count: usize, stream: &Arc<CudaStream>) -> Result<CudaSlice<f32>, MLError> {
let count = count.min(self.num_bars.saturating_sub(start));
Self::d2d_subrange(&self.targets, start * 4, count * 4, stream)
}
/// Get target values for a bar as a GPU-resident [4] `CudaSlice<half::bf16>`.
/// Get target values for a bar as a GPU-resident [4] `CudaSlice<f32>`.
///
/// Tensor order: [current_close_preproc, next_close_preproc, current_close_raw, next_close_raw].
/// Stays on GPU. Callers that need CPU scalars should extract at their own boundary.
pub fn bar_target_values(&self, bar_idx: usize, stream: &Arc<CudaStream>) -> Result<CudaSlice<half::bf16>, MLError> {
pub fn bar_target_values(&self, bar_idx: usize, stream: &Arc<CudaStream>) -> Result<CudaSlice<f32>, MLError> {
self.bar_targets(bar_idx, stream)
}
@@ -483,13 +485,13 @@ impl DqnGpuData {
bar_idx: usize,
portfolio_features: &[f32; 3],
stream: &Arc<CudaStream>,
) -> Result<CudaSlice<half::bf16>, MLError> {
) -> Result<CudaSlice<f32>, MLError> {
let ofi_dim = if self.ofi_features.is_some() { 8 } else { 0 };
let raw_dim = self.feature_dim + 3 + ofi_dim;
let final_dim = self.aligned_state_dim.unwrap_or(raw_dim);
// Allocate output buffer (zeros for padding region)
let mut dst = stream.alloc_zeros::<half::bf16>(final_dim)
let mut dst = stream.alloc_zeros::<f32>(final_dim)
.map_err(|e| MLError::ModelError(format!("build_state alloc: {e}")))?;
// DtoD copy: market features -> dst[0..feature_dim]
@@ -518,7 +520,7 @@ impl DqnGpuData {
count: usize,
portfolio_features: &[f32; 3],
stream: &Arc<CudaStream>,
) -> Result<CudaSlice<half::bf16>, MLError> {
) -> Result<CudaSlice<f32>, MLError> {
let count = count.min(self.num_bars.saturating_sub(start));
if count == 0 {
return Err(MLError::ModelError("Empty batch for state construction".to_owned()));
@@ -528,7 +530,7 @@ impl DqnGpuData {
let final_dim = self.aligned_state_dim.unwrap_or(raw_dim);
// Allocate output buffer (zeros for padding region)
let mut dst = stream.alloc_zeros::<half::bf16>(count * final_dim)
let mut dst = stream.alloc_zeros::<f32>(count * final_dim)
.map_err(|e| MLError::ModelError(format!("build_batch alloc: {e}")))?;
// Upload portfolio features once (3 scalars)
@@ -574,8 +576,8 @@ impl DqnGpuData {
/// DtoD copy `count` f32 elements from `src` (offset 0) into `dst` at element offset.
#[allow(unsafe_code)]
fn dtod_copy_into(
src: &CudaSlice<half::bf16>,
dst: &mut CudaSlice<half::bf16>,
src: &CudaSlice<f32>,
dst: &mut CudaSlice<f32>,
dst_offset_elems: usize,
count: usize,
stream: &Arc<CudaStream>,
@@ -589,9 +591,9 @@ impl DqnGpuData {
/// buffer. The source is sliced via `CudaSlice::slice()` for safety.
#[allow(unsafe_code)]
fn dtod_copy_into_at_offset(
src: &CudaSlice<half::bf16>,
src: &CudaSlice<f32>,
src_offset_elems: usize,
dst: &mut CudaSlice<half::bf16>,
dst: &mut CudaSlice<f32>,
dst_offset_elems: usize,
count: usize,
stream: &Arc<CudaStream>,
@@ -599,12 +601,12 @@ impl DqnGpuData {
if count == 0 {
return Ok(());
}
let nbytes = count * std::mem::size_of::<half::bf16>();
let nbytes = count * std::mem::size_of::<f32>();
let src_view = src.slice(src_offset_elems..src_offset_elems + count);
let (src_ptr, _sg) = src_view.device_ptr(stream);
// Get base pointer of dst, then offset by dst_offset_elems * sizeof(f32)
let (dst_base, _dg) = dst.device_ptr_mut(stream);
let dst_ptr = dst_base + (dst_offset_elems * std::mem::size_of::<half::bf16>()) as u64;
let dst_ptr = dst_base + (dst_offset_elems * std::mem::size_of::<f32>()) as u64;
unsafe {
cudarc::driver::result::memcpy_dtod_async(
dst_ptr, src_ptr, nbytes, stream.cu_stream(),
@@ -620,7 +622,7 @@ impl DqnGpuData {
#[allow(unsafe_code)]
fn htod_copy_into(
src: &[f32],
dst: &mut CudaSlice<half::bf16>,
dst: &mut CudaSlice<f32>,
dst_offset_elems: usize,
stream: &Arc<CudaStream>,
) -> Result<(), MLError> {
@@ -709,7 +711,7 @@ impl GpuBufferPool {
}
}
// Upload the used slice to GPU via f32→bf16 conversion + clone_htod.
// Upload the used slice to GPU via clone_htod.
let features = clone_htod_f32_to_bf16(stream, &self.feature_buf[..feat_len])?;
let targets = clone_htod_f32_to_bf16(stream, &self.target_buf[..targ_len])?;
@@ -731,10 +733,10 @@ impl GpuBufferPool {
/// Pre-uploaded GPU training data for PPO trainer.
///
/// Holds all rollout states as a single [N * state_dim] `CudaSlice<half::bf16>`.
/// Holds all rollout states as a single [N * state_dim] `CudaSlice<f32>`.
pub struct PpoGpuData {
/// All states [num_steps * state_dim] on GPU (f32, row-major)
pub states: CudaSlice<half::bf16>,
pub states: CudaSlice<f32>,
/// Number of training steps
pub num_steps: usize,
/// State dimension
@@ -751,7 +753,7 @@ impl std::fmt::Debug for PpoGpuData {
}
impl PpoGpuData {
/// Upload PPO market data to GPU as a single contiguous `CudaSlice<half::bf16>`.
/// Upload PPO market data to GPU as a single contiguous `CudaSlice<f32>`.
pub fn upload(market_data: &[Vec<f32>], stream: &Arc<CudaStream>) -> Result<Self, MLError> {
let num_steps = market_data.len();
if num_steps == 0 {
@@ -796,7 +798,7 @@ impl PpoGpuData {
}
/// Get a single step's state as a [state_dim] `CudaSlice` (D2D copy from GPU buffer).
pub fn step_state(&self, step_idx: usize, stream: &Arc<CudaStream>) -> Result<CudaSlice<half::bf16>, MLError> {
pub fn step_state(&self, step_idx: usize, stream: &Arc<CudaStream>) -> Result<CudaSlice<f32>, MLError> {
if step_idx >= self.num_steps {
return Err(MLError::ModelError(format!("Step {step_idx} out of range (num_steps={})", self.num_steps)));
}
@@ -804,7 +806,7 @@ impl PpoGpuData {
}
/// Get a batch of states as [batch_size * state_dim] `CudaSlice`.
pub fn batch_states(&self, start: usize, count: usize, stream: &Arc<CudaStream>) -> Result<CudaSlice<half::bf16>, MLError> {
pub fn batch_states(&self, start: usize, count: usize, stream: &Arc<CudaStream>) -> Result<CudaSlice<f32>, MLError> {
let count = count.min(self.num_steps.saturating_sub(start));
DqnGpuData::d2d_subrange(&self.states, start * self.state_dim, count * self.state_dim, stream)
}
@@ -850,11 +852,10 @@ mod tests {
let batch = gpu_data.batch_features(0, 10, &stream).expect("batch");
assert_eq!(batch.len(), 10 * 42);
// Download target values and verify on CPU (bf16 -> f32 conversion)
// Download target values and verify on CPU
let targets_gpu = gpu_data.bar_target_values(0, &stream).expect("targets");
let mut targets_bf16 = vec![half::bf16::ZERO; 4];
stream.memcpy_dtoh(&targets_gpu, &mut targets_bf16).expect("DtoH"); // test readback
let targets_host: Vec<f32> = targets_bf16.iter().map(|v| v.to_f32()).collect();
let mut targets_host = vec![0.0_f32; 4];
dtoh_f32(&stream, &targets_gpu, &mut targets_host).expect("DtoH");
assert!((targets_host.first().copied().unwrap_or(0.0) - 100.0).abs() < 0.5);
}
@@ -962,10 +963,10 @@ mod tests {
#[test]
fn test_vram_estimation() {
// 300K bars × 55 features × 2 bytes (BF16) = ~33 MB
// 300K bars × 55 features × 4 bytes (f32) = ~66 MB
let bytes = estimate_vram_bytes(300_000 * 55);
assert!(bytes < 50_000_000); // Under 50 MB
assert!(bytes > 30_000_000); // Over 30 MB
assert!(bytes < 100_000_000); // Under 100 MB
assert!(bytes > 60_000_000); // Over 60 MB
}
#[test]
@@ -1001,13 +1002,11 @@ mod tests {
let batch = gpu_data.build_batch_states(10, 32, &portfolio, &stream).expect("batch");
assert_eq!(batch.len(), 32 * 45); // 42 market + 3 portfolio
// Download and verify portfolio features in row 0 (bf16 -> f32 conversion)
let mut host_bf16 = vec![half::bf16::ZERO; 32 * 45];
stream.memcpy_dtoh(&batch, &mut host_bf16).expect("DtoH"); // test readback
let host: Vec<f32> = host_bf16.iter().map(|v| v.to_f32()).collect();
// Download and verify portfolio features in row 0
let mut host = vec![0.0_f32; 32 * 45];
dtoh_f32(&stream, &batch, &mut host).expect("DtoH");
// Row 0: features[42..45] should be portfolio [0.95, 0.5, 0.0001]
// BF16 precision: relax tolerance
let pf = &host[42..45];
assert!((pf[0] - 0.95).abs() < 0.02, "portfolio[0]={}", pf[0]);
assert!((pf[1] - 0.5).abs() < 0.02, "portfolio[1]={}", pf[1]);
@@ -1072,11 +1071,10 @@ mod tests {
let g2 = pool.upload_dqn(&data2, &stream).expect("upload2");
assert_eq!(g2.num_bars, 200);
// Download target values and verify (bf16 -> f32 conversion)
// Download target values and verify
let t = g2.bar_target_values(0, &stream).expect("targets");
let mut host_bf16 = vec![half::bf16::ZERO; 4];
stream.memcpy_dtoh(&t, &mut host_bf16).expect("DtoH"); // test readback
let host: Vec<f32> = host_bf16.iter().map(|v| v.to_f32()).collect();
let mut host = vec![0.0_f32; 4];
dtoh_f32(&stream, &t, &mut host).expect("DtoH");
assert!((host[0] - 200.0).abs() < 1.0);
}
@@ -1231,17 +1229,14 @@ mod tests {
let advantages_host: Vec<f32> = (0..total).map(|i| i as f32 * 0.1 - 2.0).collect();
let returns_host: Vec<f32> = (0..total).map(|i| i as f32 * 0.5).collect();
let states_bf16: Vec<half::bf16> = states_host.iter().map(|&v| half::bf16::from_f32(v)).collect();
let log_probs_bf16: Vec<half::bf16> = vec![half::bf16::from_f32(-1.5); total];
let advantages_bf16: Vec<half::bf16> = advantages_host.iter().map(|&v| half::bf16::from_f32(v)).collect();
let returns_bf16: Vec<half::bf16> = returns_host.iter().map(|&v| half::bf16::from_f32(v)).collect();
let log_probs_host: Vec<f32> = vec![-1.5_f32; total];
let batch = PpoExperienceBatch {
states: stream.clone_htod(&states_bf16).unwrap(),
states: stream.clone_htod(&states_host).unwrap(),
actions: stream.clone_htod(&(0..total).map(|i| (i % 45) as i32).collect::<Vec<_>>()).unwrap(),
log_probs: stream.clone_htod(&log_probs_bf16).unwrap(),
advantages: stream.clone_htod(&advantages_bf16).unwrap(),
returns: stream.clone_htod(&returns_bf16).unwrap(),
log_probs: stream.clone_htod(&log_probs_host).unwrap(),
advantages: stream.clone_htod(&advantages_host).unwrap(),
returns: stream.clone_htod(&returns_host).unwrap(),
done_flags: stream.clone_htod(&(0..total).map(|i| if i % timesteps == timesteps - 1 { 1_i32 } else { 0 }).collect::<Vec<_>>()).unwrap(),
n_episodes,
timesteps,
@@ -1322,17 +1317,15 @@ mod tests {
}
let returns_host: Vec<f32> = advantages_host.iter().map(|&a| a + 5.0).collect();
let states_bf16: Vec<half::bf16> = vec![half::bf16::ZERO; total * state_dim];
let log_probs_bf16: Vec<half::bf16> = vec![half::bf16::from_f32(-1.0); total];
let advantages_bf16: Vec<half::bf16> = advantages_host.iter().map(|&v| half::bf16::from_f32(v)).collect();
let returns_bf16: Vec<half::bf16> = returns_host.iter().map(|&v| half::bf16::from_f32(v)).collect();
let states_f32: Vec<f32> = vec![0.0_f32; total * state_dim];
let log_probs_f32: Vec<f32> = vec![-1.0_f32; total];
let batch = PpoExperienceBatch {
states: stream.clone_htod(&states_bf16).unwrap(),
states: stream.clone_htod(&states_f32).unwrap(),
actions: stream.clone_htod(&vec![1_i32; total]).unwrap(),
log_probs: stream.clone_htod(&log_probs_bf16).unwrap(),
advantages: stream.clone_htod(&advantages_bf16).unwrap(),
returns: stream.clone_htod(&returns_bf16).unwrap(),
log_probs: stream.clone_htod(&log_probs_f32).unwrap(),
advantages: stream.clone_htod(&advantages_host).unwrap(),
returns: stream.clone_htod(&returns_host).unwrap(),
done_flags: stream.clone_htod(&(0..total).map(|i| if i % timesteps == timesteps - 1 { 1_i32 } else { 0 }).collect::<Vec<_>>()).unwrap(),
n_episodes,
timesteps,
@@ -1347,7 +1340,7 @@ mod tests {
let r = returns.get(i).copied().unwrap_or(0.0);
let a = advantages.get(i).copied().unwrap_or(0.0);
assert!(
(r - a - 5.0).abs() < 0.1, // BF16 precision: values round-trip with ~1% error
(r - a - 5.0).abs() < 0.1, // f32 precision: values round-trip cleanly
"Value at {} should be ~5.0, got {}",
i,
r - a