diff --git a/crates/ml-alpha/examples/alpha_rl_train.rs b/crates/ml-alpha/examples/alpha_rl_train.rs index 8146c176f..df1adbc4e 100644 --- a/crates/ml-alpha/examples/alpha_rl_train.rs +++ b/crates/ml-alpha/examples/alpha_rl_train.rs @@ -87,6 +87,7 @@ use ml_alpha::rl::isv_slots::{ RL_TD_KURTOSIS_EMA_INDEX, RL_TD_KURT_STREAM_M2_INDEX, RL_TD_KURT_STREAM_M4_INDEX, RL_TD_KURT_STREAM_MEAN_INDEX, RL_V_GRAD_NORM_EMA_INDEX, }; +use ml_alpha::trainer::diag_staging::DiagStaging; use ml_alpha::trainer::integrated::{ read_slice_d_pub, read_slice_i32_d_pub, read_slice_u8_d_pub, write_slice_i32_d_pub, IntegratedStepStats, IntegratedTrainer, IntegratedTrainerConfig, @@ -440,6 +441,15 @@ fn main() -> Result<()> { let mut diag = BufWriter::new(diag_file); eprintln!("per-step diag JSONL: {}", diag_path.display()); + // Async non-blocking diagnostic staging: separate CUDA stream + + // double-buffered mapped-pinned memory. Replaces 7 blocking + // read_slice_d calls (ISV, rewards, dones, actions, raw_rewards, + // trade_duration, outcome_ema) with async DtoD copies. The host + // reads the PREVIOUS step's buffer while the current step's copies + // run concurrently on the diag stream. One-step delay on diag data. + let mut diag_staging = DiagStaging::new(&dev, cli.n_backtests) + .context("DiagStaging::new")?; + // Per-step scratch (replaced each step via mapped-pinned helpers). // ── Training loop. ─────────────────────────────────────────────── @@ -485,6 +495,12 @@ fn main() -> Result<()> { let t_start = std::time::Instant::now(); 'outer: for step in 0..cli.n_steps { + // Sync the diag stream and swap buffers. On step 0 this waits + // for nothing (buffers are zero-init); from step 1+ it + // completes the PREVIOUS step's async DtoD copies so the host + // can read them below in the diag JSONL section. + diag_staging.sync_and_swap().context("diag sync_and_swap")?; + // Sample n_backtests INDEPENDENT pairs per step. Each pair // becomes one batch slot's K-window — distinct market segments // per batch slot give true gradient-variance reduction (vs @@ -561,21 +577,47 @@ fn main() -> Result<()> { process::exit(2); } - // ── Per-step diag dump (mapped-pinned reads). ──────────────── - // Per `feedback_no_htod_htoh_only_mapped_pinned`: raw - // `stream.memcpy_dtoh` on a regular `&mut [T]` is forbidden; - // the only permitted CPU↔GPU path is mapped-pinned staging - // (cuMemHostAlloc DEVICEMAP → host writes via host_ptr + - // kernel reads dev_ptr + DtoD between them). The - // `read_slice_*_d_pub` helpers in `trainer::integrated` - // encapsulate the pattern. + // Launch async DtoD copies of the 7 high-frequency diag buffers + // into the current staging buffer. Non-blocking on the training + // stream — the copies run on diag_staging's separate stream. + // The host will read these values at the START of the NEXT step + // iteration (after sync_and_swap completes the copy). + diag_staging + .snapshot_async( + &trainer.isv_d, + &trainer.rewards_d, + &trainer.dones_d, + &trainer.actions_d, + &trainer.raw_rewards_d, + &trainer.trade_duration_emit_d, + &trainer.outcome_ema_d, + ) + .context("diag snapshot_async")?; + + // ── Per-step diag dump (async staging reads). ──────────────── + // The 7 high-frequency reads (ISV, rewards, dones, actions, + // raw_rewards, trade_duration, outcome_ema) now come from the + // DiagStaging double-buffer — zero GPU stalls on the training + // stream. On step 0 these report zeros (no prior snapshot); + // from step 1+ they report step N-1's data (one-step delay, + // acceptable for diagnostics). + // + // Remaining reads (position_lots, unit state, pyramid_count, + // close_unit_index, frd_logits) still use blocking helpers — + // they're less frequent or structurally complex. + let rewards_host = diag_staging.read_rewards(); + let dones_host = diag_staging.read_dones(); + let actions_raw = diag_staging.read_actions_raw(); + let actions_host: Vec = actions_raw + .iter() + .map(|f| f.to_bits() as i32) + .collect(); + let raw_rewards_host = diag_staging.read_raw_rewards(); + let trade_duration_host = diag_staging.read_trade_duration(); + let outcome_ema_host = diag_staging.read_outcome_ema(); + + // Remaining blocking reads — structurally complex or low-frequency. let dev_stream = dev.cuda_stream().context("cuda_stream for diag")?; - let actions_host = read_slice_i32_d_pub(dev_stream, &trainer.actions_d, cli.n_backtests) - .context("diag: read actions_d")?; - let rewards_host = read_slice_d_pub(dev_stream, &trainer.rewards_d, cli.n_backtests) - .context("diag: read rewards_d")?; - let dones_host = read_slice_d_pub(dev_stream, &trainer.dones_d, cli.n_backtests) - .context("diag: read dones_d")?; // audit — position state per batch (signed lot count from // `prev_position_lots_d`; positive=long, negative=short, 0=flat). // Helps post-hoc analysis understand whether the agent is @@ -610,15 +652,6 @@ fn main() -> Result<()> { let close_unit_index_host = read_slice_i32_d_pub( dev_stream, &trainer.close_unit_index_d, cli.n_backtests, ).context("diag: read close_unit_index_d")?; - let outcome_ema_host = read_slice_d_pub( - dev_stream, &trainer.outcome_ema_d, cli.n_backtests, - ).context("diag: read outcome_ema_d")?; - let raw_rewards_host = read_slice_d_pub( - dev_stream, &trainer.raw_rewards_d, cli.n_backtests, - ).context("diag: read raw_rewards_d")?; - let trade_duration_host = read_slice_d_pub( - dev_stream, &trainer.trade_duration_emit_d, cli.n_backtests, - ).context("diag: read trade_duration_emit_d")?; let mut act_hist = [0u32; N_ACTIONS]; for &a in &actions_host { diff --git a/crates/ml-alpha/src/trainer/diag_staging.rs b/crates/ml-alpha/src/trainer/diag_staging.rs new file mode 100644 index 000000000..187f4e520 --- /dev/null +++ b/crates/ml-alpha/src/trainer/diag_staging.rs @@ -0,0 +1,258 @@ +//! Async non-blocking diagnostic staging via separate CUDA stream + +//! double-buffered mapped-pinned memory. +//! +//! The training loop reads ~7 device buffers to host every step for JSONL +//! diagnostics. Each blocking `read_slice_d` does `stream.synchronize()` +//! which stalls the GPU training pipeline. At b=256 this wastes ~5ms/step. +//! +//! `DiagStaging` solves this by: +//! 1. Running DtoD copies on a **separate** CUDA stream (non-blocking on +//! the training stream). +//! 2. Using **double-buffered** mapped-pinned memory: the host reads the +//! PREVIOUS step's buffer while the current step's copy runs on the +//! diag stream. +//! +//! Per `feedback_no_htod_htoh_only_mapped_pinned`: mapped-pinned is the +//! only permitted CPU<->GPU path. The DtoD copies land in mapped-pinned +//! device addresses; after `diag_stream.synchronize()` the host reads +//! the same physical pages via host_ptr. +//! +//! One-step latency: step 0 reports zeros (no prior snapshot). From step 1 +//! onward, the host reads step N-1's data. Acceptable for diagnostics. + +use std::sync::Arc; + +use anyhow::{Context, Result}; +use cudarc::driver::{CudaSlice, CudaStream, DevicePtr}; +use ml_core::device::MlDevice; + +use crate::pinned_mem::MappedF32Buffer; +use crate::rl::isv_slots::RL_SLOTS_END; + +/// Double-buffered mapped-pinned staging for async diagnostic reads. +/// +/// Layout within each buffer (all f32, contiguous): +/// `[ISV (RL_SLOTS_END) | rewards (B) | dones (B) | actions (B) | raw_rewards (B) | trade_duration (B) | outcome_ema (B)]` +/// +/// `actions` are stored as raw i32 bit patterns reinterpreted as f32 bytes. +pub struct DiagStaging { + diag_stream: Arc, + bufs: [MappedF32Buffer; 2], + current: usize, + /// Total f32 count per buffer. + pub n_floats: usize, + // Layout offsets (all in f32 elements from start of buffer). + pub isv_offset: usize, + pub isv_len: usize, + pub rewards_offset: usize, + pub dones_offset: usize, + pub actions_offset: usize, + pub raw_rewards_offset: usize, + pub trade_duration_offset: usize, + pub outcome_ema_offset: usize, + pub b_size: usize, +} + +impl DiagStaging { + /// Allocate double-buffered mapped-pinned staging for `b_size` batch slots. + /// + /// Creates a new CUDA stream on the same context as `dev`. The stream is + /// independent of the training stream — copies launched on it do not stall + /// the training pipeline. + pub fn new(dev: &MlDevice, b_size: usize) -> Result { + let ctx = dev.cuda_context().context("DiagStaging needs CUDA context")?; + let diag_stream = ctx + .new_stream() + .map_err(|e| anyhow::anyhow!("DiagStaging new_stream: {e}"))?; + + // Layout: ISV | rewards | dones | actions | raw_rewards | trade_duration | outcome_ema + let isv_len = RL_SLOTS_END; + let n_floats = isv_len + 6 * b_size; + + let buf0 = unsafe { MappedF32Buffer::new(n_floats) } + .map_err(|e| anyhow::anyhow!("DiagStaging buf0: {e}"))?; + let buf1 = unsafe { MappedF32Buffer::new(n_floats) } + .map_err(|e| anyhow::anyhow!("DiagStaging buf1: {e}"))?; + + Ok(Self { + diag_stream, + bufs: [buf0, buf1], + current: 0, + n_floats, + isv_offset: 0, + isv_len, + rewards_offset: isv_len, + dones_offset: isv_len + b_size, + actions_offset: isv_len + 2 * b_size, + raw_rewards_offset: isv_len + 3 * b_size, + trade_duration_offset: isv_len + 4 * b_size, + outcome_ema_offset: isv_len + 5 * b_size, + b_size, + }) + } + + /// Launch async DtoD copies from device buffers into the CURRENT staging + /// buffer. Non-blocking on the training stream — runs on the diag_stream. + /// + /// The cudarc v0.19 `device_ptr()` call on a `CudaSlice` with a different + /// stream automatically inserts inter-stream event waits: the diag_stream + /// will wait for any prior writes to the source slice before the copy + /// begins. This ensures correctness without explicit event management. + pub fn snapshot_async( + &mut self, + isv_d: &CudaSlice, + rewards_d: &CudaSlice, + dones_d: &CudaSlice, + actions_d: &CudaSlice, + raw_rewards_d: &CudaSlice, + trade_duration_d: &CudaSlice, + outcome_ema_d: &CudaSlice, + ) -> Result<()> { + let buf = &self.bufs[self.current]; + let s = self.diag_stream.cu_stream(); + let b = self.b_size; + + unsafe { + // ISV + let (src, _g) = isv_d.device_ptr(&self.diag_stream); + cudarc::driver::result::memcpy_dtod_async( + buf.dev_ptr, + src, + self.isv_len * 4, + s, + ) + .context("diag snapshot isv")?; + + // rewards + let (src, _g) = rewards_d.device_ptr(&self.diag_stream); + cudarc::driver::result::memcpy_dtod_async( + buf.dev_ptr + (self.rewards_offset * 4) as u64, + src, + b * 4, + s, + ) + .context("diag snapshot rewards")?; + + // dones + let (src, _g) = dones_d.device_ptr(&self.diag_stream); + cudarc::driver::result::memcpy_dtod_async( + buf.dev_ptr + (self.dones_offset * 4) as u64, + src, + b * 4, + s, + ) + .context("diag snapshot dones")?; + + // actions (i32 → same byte width as f32, reinterpret on host) + let (src, _g) = actions_d.device_ptr(&self.diag_stream); + cudarc::driver::result::memcpy_dtod_async( + buf.dev_ptr + (self.actions_offset * 4) as u64, + src, + b * 4, + s, + ) + .context("diag snapshot actions")?; + + // raw_rewards + let (src, _g) = raw_rewards_d.device_ptr(&self.diag_stream); + cudarc::driver::result::memcpy_dtod_async( + buf.dev_ptr + (self.raw_rewards_offset * 4) as u64, + src, + b * 4, + s, + ) + .context("diag snapshot raw_rewards")?; + + // trade_duration + let (src, _g) = trade_duration_d.device_ptr(&self.diag_stream); + cudarc::driver::result::memcpy_dtod_async( + buf.dev_ptr + (self.trade_duration_offset * 4) as u64, + src, + b * 4, + s, + ) + .context("diag snapshot trade_duration")?; + + // outcome_ema + let (src, _g) = outcome_ema_d.device_ptr(&self.diag_stream); + cudarc::driver::result::memcpy_dtod_async( + buf.dev_ptr + (self.outcome_ema_offset * 4) as u64, + src, + b * 4, + s, + ) + .context("diag snapshot outcome_ema")?; + } + + Ok(()) + } + + /// Synchronize the diag stream (waits for the CURRENT buffer's copies to + /// complete) then swap buffers. After this call, the buffer that was being + /// written to is now the "complete" buffer readable by the host, and the + /// other buffer becomes the new write target. + pub fn sync_and_swap(&mut self) -> Result<()> { + self.diag_stream + .synchronize() + .context("diag stream sync")?; + self.current = 1 - self.current; + Ok(()) + } + + /// Read the PREVIOUS step's ISV data (the buffer NOT currently targeted + /// for writes). Call AFTER `sync_and_swap`. + pub fn read_isv(&self) -> &[f32] { + let prev = 1 - self.current; + let buf = &self.bufs[prev]; + unsafe { std::slice::from_raw_parts(buf.host_ptr.add(self.isv_offset), self.isv_len) } + } + + /// Read the PREVIOUS step's per-batch rewards. + pub fn read_rewards(&self) -> &[f32] { + let prev = 1 - self.current; + let buf = &self.bufs[prev]; + unsafe { std::slice::from_raw_parts(buf.host_ptr.add(self.rewards_offset), self.b_size) } + } + + /// Read the PREVIOUS step's per-batch done flags. + pub fn read_dones(&self) -> &[f32] { + let prev = 1 - self.current; + let buf = &self.bufs[prev]; + unsafe { std::slice::from_raw_parts(buf.host_ptr.add(self.dones_offset), self.b_size) } + } + + /// Read the PREVIOUS step's per-batch actions as raw f32 bit patterns. + /// The underlying data is i32; caller reinterprets via `f32::to_bits() as i32`. + pub fn read_actions_raw(&self) -> &[f32] { + let prev = 1 - self.current; + let buf = &self.bufs[prev]; + unsafe { std::slice::from_raw_parts(buf.host_ptr.add(self.actions_offset), self.b_size) } + } + + /// Read the PREVIOUS step's per-batch raw (pre-scale) rewards. + pub fn read_raw_rewards(&self) -> &[f32] { + let prev = 1 - self.current; + let buf = &self.bufs[prev]; + unsafe { + std::slice::from_raw_parts(buf.host_ptr.add(self.raw_rewards_offset), self.b_size) + } + } + + /// Read the PREVIOUS step's per-batch trade duration. + pub fn read_trade_duration(&self) -> &[f32] { + let prev = 1 - self.current; + let buf = &self.bufs[prev]; + unsafe { + std::slice::from_raw_parts(buf.host_ptr.add(self.trade_duration_offset), self.b_size) + } + } + + /// Read the PREVIOUS step's per-batch outcome EMA. + pub fn read_outcome_ema(&self) -> &[f32] { + let prev = 1 - self.current; + let buf = &self.bufs[prev]; + unsafe { + std::slice::from_raw_parts(buf.host_ptr.add(self.outcome_ema_offset), self.b_size) + } + } +} diff --git a/crates/ml-alpha/src/trainer/mod.rs b/crates/ml-alpha/src/trainer/mod.rs index 3ba7a2dfe..91dfa5fb3 100644 --- a/crates/ml-alpha/src/trainer/mod.rs +++ b/crates/ml-alpha/src/trainer/mod.rs @@ -1,6 +1,7 @@ //! Trainer module — PerceptionTrainer wraps the full stacked //! Mamba2 -> CfC -> heads pipeline plus 6 AdamW optimizer groups. +pub mod diag_staging; pub mod integrated; pub mod loss; pub mod optim;