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

fix(ml): wire OFI features into GPU batch state construction

Browse Source 
The GPU path (build_batch_states / build_state_tensor) only concatenated
market[40] + portfolio[3] = 43 dims, while state_dim was set to 51 when
OFI was enabled. This caused shape mismatch [128,43] vs [51,1024].

- Add ofi_features: Option<Tensor> field to DqnGpuData
- Add upload_ofi() method for MBP-10 order book features (8 dims/bar)
- Concatenate OFI in build_batch_states: [count,40]+[count,3]+[count,8]=[count,51]
- Concatenate OFI in build_state_tensor: [1,40]+[1,3]+[1,8]=[1,51]
- Wire trainer to call upload_ofi() after GPU data upload
- Fix CPU path to zero-pad regime_features when OFI enabled but data missing

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-03-07 20:16:32 +01:00
parent 9a74fda3ab
commit b10cdebfb1
2 changed files with 91 additions and 15 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -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<Tensor>,
/// 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,
})

20
crates/ml/src/trainers/dqn/trainer.rs
View File

@@ -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<f32> = 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![]
};