diff --git a/crates/ml/src/cuda_pipeline/mod.rs b/crates/ml/src/cuda_pipeline/mod.rs index 3d01936dd..5123c4fd5 100644 --- a/crates/ml/src/cuda_pipeline/mod.rs +++ b/crates/ml/src/cuda_pipeline/mod.rs @@ -75,15 +75,19 @@ pub fn optimal_launch_dims(n_items: u32, max_threads_per_block: u32) -> (u32, u3 /// Pre-uploaded GPU training data for DQN trainer. /// -/// Holds market features [N, 51] and target prices [N, 4] as GPU tensors. +/// Holds market features [N, 40] and target prices [N, 4] as GPU tensors. /// Portfolio features (3 dims) are computed per-bar by the trainer and /// concatenated on-device via `Tensor::cat`. +/// OFI features (8 dims from MBP-10 order book) are optionally uploaded +/// and concatenated to produce [N, 51] states when present. #[derive(Debug)] pub struct DqnGpuData { /// Market features tensor [num_bars, 40] on GPU (f32) pub features: Tensor, /// Target prices tensor [num_bars, 4] on GPU (f32) pub targets: Tensor, + /// OFI features tensor [num_bars, 8] on GPU (f32), from MBP-10 data + pub ofi_features: Option, /// Number of training bars pub num_bars: usize, /// Feature dimension (40) @@ -143,14 +147,51 @@ impl DqnGpuData { Ok(Self { features, targets, + ofi_features: None, num_bars, feature_dim, }) } + /// Upload OFI features (8 dims per bar) from MBP-10 order book data. + /// + /// Must be called after `upload()` with a slice matching `num_bars` length. + /// Features are: OFI, VPIN, Kyle's Lambda, trade imbalance, etc. + pub fn upload_ofi( + &mut self, + ofi_data: &[[f64; 8]], + device: &Device, + ) -> Result<(), MLError> { + let n = ofi_data.len().min(self.num_bars); + if n == 0 { + return Ok(()); + } + + let mut flat = Vec::with_capacity(self.num_bars * 8); + for i in 0..self.num_bars { + if i < n { + for &v in ofi_data[i].iter() { + flat.push(v as f32); + } + } else { + // Zero-pad if OFI data is shorter than market data + flat.extend_from_slice(&[0.0_f32; 8]); + } + } + + let tensor = Tensor::from_vec(flat, (self.num_bars, 8), device) + .map_err(|e| MLError::ModelError(format!("GPU OFI upload failed: {e}")))? + .to_dtype(self.features.dtype()) + .map_err(|e| MLError::ModelError(format!("GPU OFI dtype cast failed: {e}")))?; + + self.ofi_features = Some(tensor); + Ok(()) + } + /// Estimated VRAM usage in bytes. pub fn vram_bytes(&self) -> usize { - estimate_vram_bytes(self.num_bars * (self.feature_dim + 4)) + let ofi_dim = if self.ofi_features.is_some() { 8 } else { 0 }; + estimate_vram_bytes(self.num_bars * (self.feature_dim + 4 + ofi_dim)) } /// Get market features for a single bar as a [1, 40] tensor slice (zero-copy on GPU). @@ -202,8 +243,10 @@ impl DqnGpuData { ]) } - /// Build a complete 43-dim state tensor by concatenating pre-uploaded market features - /// with per-bar portfolio features on-device. + /// Build a complete state tensor by concatenating pre-uploaded market features + /// with per-bar portfolio features (and OFI features if available) on-device. + /// + /// Returns [1, 43] without OFI or [1, 51] with OFI. pub fn build_state_tensor( &self, bar_idx: usize, @@ -219,15 +262,24 @@ impl DqnGpuData { .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}"))) + if let Some(ref ofi) = self.ofi_features { + let ofi_slice = ofi + .narrow(0, bar_idx, 1) + .map_err(|e| MLError::ModelError(format!("OFI bar {bar_idx} slice failed: {e}")))?; + Tensor::cat(&[&market, &portfolio, &ofi_slice], 1) + .map_err(|e| MLError::ModelError(format!("State cat (with OFI) failed: {e}"))) + } else { + Tensor::cat(&[&market, &portfolio], 1) + .map_err(|e| MLError::ModelError(format!("State cat failed: {e}"))) + } } - /// Build a batch of complete 43-dim state tensors on GPU. + /// Build a batch of complete state tensors on GPU. /// /// Concatenates pre-uploaded [batch_size, 40] market features with /// broadcast [1, 3] portfolio features to produce [batch_size, 43]. - /// Eliminates ~770 Vec allocations per batch vs feature_vector_to_state(). + /// When OFI features are present, also concatenates [batch_size, 8] + /// OFI features to produce [batch_size, 51]. /// /// # Arguments /// * `start` - First bar index into pre-uploaded features @@ -260,13 +312,24 @@ impl DqnGpuData { .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, 40] → [count, 43] + // Broadcast [1, 3] → [count, 3] let portfolio_broadcast = portfolio .broadcast_as((count, 3)) .map_err(|e| MLError::ModelError(format!("Portfolio broadcast failed: {e}")))?; - Tensor::cat(&[&market, &portfolio_broadcast], 1) - .map_err(|e| MLError::ModelError(format!("State cat failed: {e}"))) + if let Some(ref ofi) = self.ofi_features { + // [count, 8] slice from pre-uploaded OFI tensor + let ofi_batch = ofi + .narrow(0, start, count) + .map_err(|e| MLError::ModelError(format!("OFI batch slice failed: {e}")))?; + // [count, 40] + [count, 3] + [count, 8] → [count, 51] + Tensor::cat(&[&market, &portfolio_broadcast, &ofi_batch], 1) + .map_err(|e| MLError::ModelError(format!("State cat (with OFI) failed: {e}"))) + } else { + // [count, 40] + [count, 3] → [count, 43] + Tensor::cat(&[&market, &portfolio_broadcast], 1) + .map_err(|e| MLError::ModelError(format!("State cat failed: {e}"))) + } } } @@ -375,6 +438,7 @@ impl GpuBufferPool { Ok(DqnGpuData { features, targets, + ofi_features: None, num_bars, feature_dim: self.feature_dim, }) diff --git a/crates/ml/src/trainers/dqn/trainer.rs b/crates/ml/src/trainers/dqn/trainer.rs index 79d1133a8..4abbb811d 100644 --- a/crates/ml/src/trainers/dqn/trainer.rs +++ b/crates/ml/src/trainers/dqn/trainer.rs @@ -1583,12 +1583,19 @@ impl DQNTrainer { DqnGpuData::upload(&training_data, &self.device) }; match upload_result { - Ok(gpu_data) => { + Ok(mut gpu_data) => { info!("GPU data pre-uploaded: {} bars x {} features ({:.1} MB)", gpu_data.num_bars, gpu_data.feature_dim, (gpu_data.num_bars * (40 + 4) * 4) as f64 / 1_048_576.0 ); + // Upload OFI features to GPU if available + if let Some(ref ofi) = self.ofi_features { + match gpu_data.upload_ofi(ofi, &self.device) { + Ok(()) => info!("GPU OFI features uploaded: {} bars x 8 dims", ofi.len()), + Err(e) => debug!("GPU OFI upload skipped: {}", e), + } + } self.gpu_data = Some(gpu_data); } Err(e) => { @@ -1877,7 +1884,7 @@ impl DQNTrainer { // Build batch states — GPU path skips ~770 Vec allocs per batch let (batch_tensor, states) = if let Some(ref gpu_data) = self.gpu_data { - // GPU path: build [batch_size, 54] directly from pre-uploaded features + // GPU path: build [batch_size, state_dim] directly from pre-uploaded features let current_price_f32 = { let t = gpu_data.bar_target_values(batch_start)?; if t[2] != 0.0 { t[2] } else { t[0] } @@ -3322,11 +3329,16 @@ impl DQNTrainer { vec![0.0, 0.0, 0.0] // Fallback if no price provided }; - // OFI regime features: 8 features from MBP-10 order book data (or empty if unavailable) + // OFI regime features: 8 features from MBP-10 order book data. + // When OFI is enabled (mbp10_data_dir set), always return 8 features + // (zeros if data didn't load) to match state_dim=51. + let ofi_enabled = self.hyperparams.mbp10_data_dir.is_some(); let regime_features: Vec = if let (Some(ofi), Some(idx)) = (&self.ofi_features, ofi_index) { ofi.get(idx) .map(|f| f.iter().map(|&v| v as f32).collect()) - .unwrap_or_default() + .unwrap_or_else(|| vec![0.0; 8]) + } else if ofi_enabled { + vec![0.0; 8] } else { vec![] };