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

feat: IQN prepare_buffers extraction for parallel stream dispatch

Browse Source 
Separates buffer prep (decode_actions, DtoD copies) from training pipeline
so prep runs on main stream and execute_training_pipeline runs on iqn_stream.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-04-18 16:37:54 +02:00
parent 108a0993c7
commit 80d871ae1d
1 changed files with 36 additions and 21 deletions
Download Patch File Download Diff File Expand all files Collapse all files

57
crates/ml/src/cuda_pipeline/gpu_iqn_head.rs
View File

@@ -574,30 +574,17 @@ impl GpuIqnHead {
})
}
/// Run one GPU-native IQN training step.
/// Prepare IQN buffers: decode flat actions and DtoD copy rewards/dones.
///
/// All data sourced from DQN trainer's GPU buffers via DtoD copies.
/// No CPU arrays, no HtoD uploads for batch data.
///
/// Full pipeline:
/// 1. Decode flat actions → branch actions on GPU
/// 2. DtoD copy rewards/dones from DQN trainer buffers
/// 3. Compute target h_s2 (or use cached)
/// 4. Forward (cuBLAS GEMMs + element-wise kernels)
/// 5. Loss (quantile Huber)
/// 6. Backward (cuBLAS GEMMs + element-wise kernels)
/// 7. Grad norm + Adam
///
/// Returns the batch-mean quantile Huber loss.
pub fn train_iqn_step_gpu(
/// Runs on the IQN head's own stream (the main training stream). Called
/// before `execute_training_pipeline` so that a dedicated override stream
/// can safely read `branch_actions`, `rewards_buf`, and `dones_buf`.
pub fn prepare_buffers(
&mut self,
online_h_s2: &CudaSlice<f32>,
next_states_buf: &CudaSlice<f32>,
target_dueling: &DuelingWeightSet,
dqn_actions_buf: &CudaSlice<i32>,
dqn_rewards_buf: &CudaSlice<f32>,
dqn_dones_buf: &CudaSlice<f32>,
) -> Result<f32, MLError> {
) -> Result<(), MLError> {
let b = self.config.batch_size;
let shared_h1_i32 = self.config.shared_h1 as i32;
let hidden_dim_i32 = self.config.hidden_dim as i32;
@@ -607,7 +594,6 @@ impl GpuIqnHead {
let b2_i32 = self.config.branch_2_size as i32;
let b3_i32 = self.config.branch_3_size as i32;
// 1. GPU decode: flat actions [B] → branch actions [B, 4]
let decode_blocks = (b + 255) / 256;
let decode_config = LaunchConfig {
grid_dim: (decode_blocks as u32, 1, 1),
@@ -633,7 +619,6 @@ impl GpuIqnHead {
.map_err(|e| MLError::ModelError(format!("IQN decode_actions kernel: {e}")))?;
}
// 2. DtoD copy rewards/dones from DQN trainer buffers (f32)
use super::gpu_dqn_trainer::dtod_copy;
let f32_bytes = std::mem::size_of::<f32>();
let r_src = dqn_rewards_buf.raw_ptr();
@@ -644,6 +629,36 @@ impl GpuIqnHead {
let d_dst = self.dones_buf.raw_ptr();
dtod_copy(d_dst, d_src, b * f32_bytes, &self.stream, 0, "iqn_dones_dtod")?;
Ok(())
}
/// Run one GPU-native IQN training step.
///
/// All data sourced from DQN trainer's GPU buffers via DtoD copies.
/// No CPU arrays, no HtoD uploads for batch data.
///
/// Full pipeline:
/// 1. Decode flat actions → branch actions on GPU
/// 2. DtoD copy rewards/dones from DQN trainer buffers
/// 3. Compute target h_s2 (or use cached)
/// 4. Forward (cuBLAS GEMMs + element-wise kernels)
/// 5. Loss (quantile Huber)
/// 6. Backward (cuBLAS GEMMs + element-wise kernels)
/// 7. Grad norm + Adam
///
/// Returns the batch-mean quantile Huber loss.
pub fn train_iqn_step_gpu(
&mut self,
online_h_s2: &CudaSlice<f32>,
next_states_buf: &CudaSlice<f32>,
target_dueling: &DuelingWeightSet,
dqn_actions_buf: &CudaSlice<i32>,
dqn_rewards_buf: &CudaSlice<f32>,
dqn_dones_buf: &CudaSlice<f32>,
) -> Result<f32, MLError> {
// 1-2. Decode actions + DtoD copy rewards/dones
self.prepare_buffers(dqn_actions_buf, dqn_rewards_buf, dqn_dones_buf)?;
// 3-9. Execute core training pipeline
self.execute_training_pipeline(online_h_s2, next_states_buf, target_dueling, None, None)
}