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

feat(ppo): pre-upload rollout states to GPU, eliminate per-step tensor allocs

Browse Source 
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-02-28 02:48:33 +01:00
parent 3883284219
commit 3cfa99de3f
1 changed files with 35 additions and 22 deletions
Download Patch File Download Diff File Expand all files Collapse all files

57
crates/ml/src/trainers/ppo.rs
View File

@@ -23,6 +23,7 @@ use crate::gpu::memory_profile::{self, resolve_batch_size};
use crate::ppo::gae::GAEConfig;
use crate::ppo::ppo::{PPOConfig, PPO};
use crate::ppo::trajectories::{Trajectory, TrajectoryBatch, TrajectoryStep};
use crate::cuda_pipeline::PpoGpuData;
use crate::MLError;
/// PPO training hyperparameters (matches gRPC PpoParams)
@@ -501,15 +502,25 @@ impl PpoTrainer {
// Track position for PnL calculation
let mut position: i8 = 0; // -1: short, 0: neutral, 1: long
// Phase 1: Pre-upload all states to GPU (single allocation, reused per-step)
let gpu_data = PpoGpuData::upload(market_data, &self.device).ok();
if gpu_data.is_some() {
tracing::debug!("PPO: Pre-uploaded {} states to GPU", num_steps);
}
for step_idx in 0..num_steps {
let state = &market_data[step_idx];
// Select action using policy
let action_probs = model.actor.action_probabilities(&Tensor::from_vec(
state.clone(),
(1, state.len()),
&self.device,
)?)?;
// Use pre-uploaded GPU tensor or fall back to per-step allocation
let state_tensor = if let Some(ref gpu) = gpu_data {
gpu.step_state(step_idx)
.map_err(|e| MLError::ModelError(format!("PPO GPU state access failed: {e}")))?
} else {
Tensor::from_vec(state.clone(), (1, state.len()), &self.device)?
};
// Actor forward pass — use pre-uploaded tensor
let action_probs = model.actor.action_probabilities(&state_tensor)?;
let action_idx = self.sample_action(&action_probs)?;
let action = TradingAction::from_int(action_idx as u8).unwrap_or(TradingAction::Hold);
@@ -522,11 +533,7 @@ impl PpoTrainer {
let value = model
.critic
.forward(&Tensor::from_vec(
state.clone(),
(1, state.len()),
&self.device,
)?)?
.forward(&state_tensor)? // Reuse same GPU tensor — eliminates second clone+alloc
.flatten_all()?
.to_vec1::<f32>()?[0]; // Flatten [1, 1] to vec, take first element
@@ -576,6 +583,12 @@ impl PpoTrainer {
let model = self.model.lock().await;
let mut all_trajectories = Vec::new();
// Phase 1: Pre-upload all states to GPU (single allocation, reused per-step)
let gpu_data = PpoGpuData::upload(market_data, &self.device).ok();
if gpu_data.is_some() {
tracing::debug!("PPO vectorized: Pre-uploaded {} states to GPU", num_steps);
}
for env_id in 0..num_envs {
let start_idx = env_id * steps_per_env;
let end_idx = ((env_id + 1) * steps_per_env).min(num_steps);
@@ -591,12 +604,16 @@ impl PpoTrainer {
for step_idx in start_idx..end_idx {
let state = &market_data[step_idx];
// Get action from policy
let action_probs = model.actor.action_probabilities(&Tensor::from_vec(
state.clone(),
(1, state.len()),
&self.device,
)?)?;
// Use pre-uploaded GPU tensor or fall back to per-step allocation
let state_tensor = if let Some(ref gpu) = gpu_data {
gpu.step_state(step_idx)
.map_err(|e| MLError::ModelError(format!("PPO GPU state access failed: {e}")))?
} else {
Tensor::from_vec(state.clone(), (1, state.len()), &self.device)?
};
// Get action from policy — use pre-uploaded tensor
let action_probs = model.actor.action_probabilities(&state_tensor)?;
let action_idx = self.sample_action(&action_probs)?;
let action =
@@ -608,11 +625,7 @@ impl PpoTrainer {
let value = model
.critic
.forward(&Tensor::from_vec(
state.clone(),
(1, state.len()),
&self.device,
)?)?
.forward(&state_tensor)? // Reuse same GPU tensor — eliminates second clone+alloc
.flatten_all()?
.to_vec1::<f32>()?[0];