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

fix: capture PER sampling inside parent graph — zero ungraphed PER kernels per step

Browse Source 
Move sample_proportional (prefix scan, binary search, gather x6, IS weights)
from training_loop.rs into run_full_step, captured as per_sample child graph
node. Steps 1+ replay all 13 children via single cuGraphLaunch with no
ungraphed PER kernel dispatches.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-04-19 01:14:53 +02:00
parent ead6d195fa
commit 1d3a533f99
2 changed files with 36 additions and 12 deletions
Download Patch File Download Diff File Expand all files Collapse all files

41
crates/ml/src/trainers/dqn/fused_training.rs
View File

@@ -898,7 +898,6 @@ impl FusedTrainingCtx {
/// After graph launch: only host-side pinned scalar reads (nanoseconds, no GPU ops).
pub(crate) fn run_full_step(
&mut self,
gpu_batch: &crate::dqn::replay_buffer_type::GpuBatch,
agent: &mut DQNAgentType,
_device: &MlDevice,
) -> Result<FusedStepResult> {
@@ -928,6 +927,21 @@ impl FusedTrainingCtx {
self.trainer.target_params_initialized = true;
}
// Phase -1: PER sampling (writes directly to trainer buffers via gather_padded)
let ptrs = agent.primary_dqn_mut().memory.gpu.sample_proportional(self.batch_size)
.map_err(|e| anyhow::anyhow!("PER sample step 0: {e}"))?;
let gpu_batch = crate::dqn::replay_buffer_type::GpuBatch {
states_ptr: ptrs.states_ptr,
actions_ptr: ptrs.actions_ptr,
rewards_ptr: ptrs.rewards_ptr,
next_states_ptr: ptrs.next_states_ptr,
dones_ptr: ptrs.dones_ptr,
weights_ptr: ptrs.weights_ptr,
indices_ptr: ptrs.indices_ptr,
episode_ids_ptr: ptrs.episode_ids_ptr,
batch_size: self.batch_size,
};
// Phase 0: Counter increments + stochastic depth mask (ungraphed)
self.submit_counters_ops()?;
@@ -965,7 +979,7 @@ impl FusedTrainingCtx {
.map_err(|e| anyhow::anyhow!("ddqn ops: {e}"))?;
// Phase 4: Aux ops (HER + EMA + IQL + IQN + attention + CQL)
self.submit_aux_ops(agent, gpu_batch)?;
self.submit_aux_ops(agent, &gpu_batch)?;
// Phase 5b: Post-aux ops (selectivity + denoise + Q-stats + risk_sgd)
self.trainer.submit_post_aux_ops(self.batch_size)
@@ -988,11 +1002,11 @@ impl FusedTrainingCtx {
.map_err(|e| anyhow::anyhow!("maintenance ops: {e}"))?;
// Phase 8: IQL modulate + PER priority update (ungraphed)
self.submit_iql_modulate_ops(agent, gpu_batch)?;
self.submit_iql_modulate_ops(agent, &gpu_batch)?;
self.submit_per_priority_ops(agent)?;
// Capture all children + compose parent after ungraphed step succeeds.
if let Err(e) = self.capture_training_graph(agent, gpu_batch) {
if let Err(e) = self.capture_training_graph(agent, &gpu_batch) {
tracing::warn!("parent graph capture failed (continuing ungraphed): {e}");
}
}
@@ -1519,13 +1533,20 @@ impl FusedTrainingCtx {
/// independently, then composed into the parent via `cuGraphAddChildGraphNode`.
///
/// Children (sequential):
/// counters -> spectral -> forward -> ddqn -> aux -> post_aux ->
/// per_sample -> counters -> spectral -> forward -> ddqn -> aux -> post_aux ->
/// adam_grad -> adam -> maintenance -> iql_modulate -> per_priority
fn capture_training_graph(
&mut self,
agent: &mut DQNAgentType,
gpu_batch: &crate::dqn::replay_buffer_type::GpuBatch,
) -> Result<()> {
// Capture PER sampling child (prefix scan + binary search + gather).
let per_sample = self.capture_child_graph("per_sample", |s| {
agent.primary_dqn_mut().memory.gpu.sample_proportional(s.batch_size)
.map_err(|e| anyhow::anyhow!("per_sample capture: {e}"))
.map(|_| ())
})?;
// Capture counters child (GPU-side step counters + stochastic depth mask).
let counters = self.capture_child_graph("counters", |s| {
s.submit_counters_ops()
@@ -1594,6 +1615,7 @@ impl FusedTrainingCtx {
})?;
// Store children.
self.per_sample_child = Some(per_sample);
self.counters_child = Some(counters);
self.spectral_child = Some(spectral);
self.forward_child = Some(forward);
@@ -1612,14 +1634,14 @@ impl FusedTrainingCtx {
// Create eval-only forward exec for deterministic Q-value replay.
self.create_eval_forward_exec()?;
info!("12 child graphs captured + parent composed (counters + spectral + forward + ddqn + aux + post_aux + adam_grad + adam + maintenance + iql_modulate + per_priority)");
info!("13 child graphs captured + parent composed (per_sample + counters + spectral + forward + ddqn + aux + post_aux + adam_grad + adam + maintenance + iql_modulate + per_priority)");
Ok(())
}
/// Compose parent graph from all child sub-graphs using `cuGraphAddChildGraphNode`.
///
/// Each child is added as a sequential dependency:
/// counters -> spectral -> forward -> ddqn -> aux -> post_aux ->
/// per_sample -> counters -> spectral -> forward -> ddqn -> aux -> post_aux ->
/// adam_grad -> adam -> maintenance -> iql_modulate -> per_priority
/// The parent is instantiated once with `cuGraphInstantiateWithFlags`.
fn compose_parent_graph(&mut self) -> Result<()> {
@@ -1632,7 +1654,10 @@ impl FusedTrainingCtx {
// Collect all children in sequential order.
// Optional children are skipped if None.
let mut children_vec: Vec<(&str, &ChildGraph)> = Vec::with_capacity(12);
let mut children_vec: Vec<(&str, &ChildGraph)> = Vec::with_capacity(13);
if let Some(ref c) = self.per_sample_child {
children_vec.push(("per_sample", c));
}
if let Some(ref c) = self.counters_child {
children_vec.push(("counters", c));
}

7
crates/ml/src/trainers/dqn/trainer/training_loop.rs
View File

@@ -1506,15 +1506,13 @@ impl DQNTrainer {
let mut agent = self.agent.write().await;
for _step in 0..num_training_steps {
let sample_start = std::time::Instant::now();
{
let buffer = agent.memory();
if !buffer.can_sample(self.current_batch_size) {
break;
}
}
let batch = agent.memory_mut().sample(self.current_batch_size)
.map_err(|e| anyhow::anyhow!("PER sample: {e}"))?;
let sample_start = std::time::Instant::now();
// Vaccine batch: only sample when needed (every 5th step for better diversity)
let vaccine_batch = if train_step_count % 5 == 0 {
agent.memory_mut().sample(self.current_batch_size).ok()
@@ -1530,7 +1528,8 @@ impl DQNTrainer {
if let Some(vb) = vaccine_batch {
fused.pending_vaccine_batch = Some(vb);
}
let _fused_result = fused.run_full_step(&batch, &mut *agent, &self.device)
// PER sampling is inside run_full_step (per_sample child graph).
let _fused_result = fused.run_full_step(&mut *agent, &self.device)
.context("Fused CUDA training step failed")?;
let step_us = fused_start.elapsed().as_micros() as u64;