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

feat(rl): multi-stream — train_stream for PER sample/priority/rebuild

Browse Source 
PER sample, update_priority, and tree_rebuild run on a dedicated
train_stream. The env-step pipeline (encoder through push_flush)
runs on the main stream. Sync points:

  1. Top of step_with_lobsim: train_stream.synchronize() ensures
     previous step's priority/rebuild completed before push touches
     the replay buffer.
  2. Before K-loop: stream.synchronize() ensures push_flush is done
     before train_stream's rl_per_sample reads the replay buffer.
  3. After rl_per_sample: train_stream.synchronize() ensures sampled_*
     buffers are written before step_synthetic/dqn_replay_step reads
     them on the main stream.
  4. After dqn_replay_step (k>0 path): stream.synchronize() ensures
     td_per_sample_d is written before train_stream's priority update.
     (step_synthetic already syncs internally.)

The last K-iter's priority update + tree rebuild are NOT synced at
step end — they overlap with the next step's encoder forward on the
main stream, hiding ~1ms of PER maintenance.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-05-26 01:55:56 +02:00
parent aab7a0b180
commit b06355456a
1 changed files with 41 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

44
crates/ml-alpha/src/trainer/integrated.rs
View File

@@ -411,6 +411,12 @@ pub struct IntegratedTrainer {
pub stream: Arc<CudaStream>,
/// Second CUDA stream for training (PER sample, priority update,
/// tree rebuild). Runs concurrently with the next env-step on
/// `self.stream`, synced at the start of the next step_with_lobsim
/// to ensure priority_tree coherence before the next PER push.
pub train_stream: Arc<CudaStream>,
// ── CUDA Graph capture for the RL step pipeline ─────────────────
// Three-state machine: first step = warmup (eager), second =
// capture, third+ = replay. Split around apply_snapshot (HtoD
@@ -951,6 +957,10 @@ pub struct IntegratedTrainer {
impl IntegratedTrainer {
pub fn new(dev: &MlDevice, cfg: IntegratedTrainerConfig) -> Result<Self> {
let stream: Arc<CudaStream> = dev.cuda_stream().context("integrated stream")?.clone();
let train_stream: Arc<CudaStream> = {
let ctx = dev.cuda_context().context("train_stream ctx")?;
ctx.new_stream().map_err(|e| anyhow::anyhow!("train_stream: {e}"))?
};
let perception =
PerceptionTrainer::new(dev, &cfg.perception).context("PerceptionTrainer::new")?;
@@ -1921,6 +1931,7 @@ impl IntegratedTrainer {
isv_d,
isv_host,
stream,
train_stream,
prefill_graph: None,
postfill_graph: None,
reward_graph: None,
@@ -4550,6 +4561,13 @@ impl IntegratedTrainer {
);
}
// ── Sync train_stream: wait for previous step's PER priority
// update + tree rebuild to complete before this step touches
// the replay buffer (push_ring/push_flush write the same tree).
self.train_stream
.synchronize()
.context("wait previous step training (train_stream)")?;
// ── Step 0: bump device-resident step counter (ISV[548]).
// Must run BEFORE any kernel that reads current_step from ISV.
// Single thread, single block — graph-safe (no scalar args change).
@@ -5782,6 +5800,13 @@ impl IntegratedTrainer {
// settles at K=1 (advantage_var_ratio drops with batch size),
// but the split protects against pathological regimes and
// makes the actor/critic separation explicit.
// Sync env stream before K-loop: ensures push_ring + push_flush
// (launched on self.stream above) are complete before train_stream
// starts reading the replay buffer in rl_per_sample.
self.stream
.synchronize()
.context("sync env stream before K-loop")?;
let mut stats = IntegratedStepStats::default();
for k_iter in 0..k_updates {
// Fresh PER sample per iteration — different transitions
@@ -5795,7 +5820,7 @@ impl IntegratedTrainer {
block_dim: (1, 1, 1),
shared_mem_bytes: 0,
};
let mut launch = self.stream.launch_builder(&self.rl_per_sample_fn);
let mut launch = self.train_stream.launch_builder(&self.rl_per_sample_fn);
launch
.arg(&self.gpu_replay.priority_tree_d)
.arg(&self.gpu_replay.h_t_d)
@@ -5821,6 +5846,13 @@ impl IntegratedTrainer {
}
}
// Sync train_stream: rl_per_sample writes sampled_h_t_d,
// sampled_h_tp1_d etc. which step_synthetic/dqn_replay_step
// read on self.stream. Must complete before the forward pass.
self.train_stream
.synchronize()
.context("sync train_stream after rl_per_sample")?;
if k_iter == 0 {
// First iter: full step_synthetic (env-step-paired).
stats = self.step_synthetic(snapshots)?;
@@ -5830,6 +5862,12 @@ impl IntegratedTrainer {
let _l_q_extra = self
.dqn_replay_step(b_size)
.context("step_with_lobsim: dqn_replay_step (k-loop iter > 0)")?;
// dqn_replay_step has no internal sync — td_per_sample_d
// may still be in-flight on self.stream. Sync before the
// priority update on train_stream reads it.
self.stream
.synchronize()
.context("sync env stream after dqn_replay_step")?;
}
// PER priority update — writes new priorities from per-sample
@@ -5845,7 +5883,7 @@ impl IntegratedTrainer {
shared_mem_bytes: (b_size * std::mem::size_of::<f32>()) as u32,
};
let mut launch =
self.stream.launch_builder(&self.rl_per_update_priority_fn);
self.train_stream.launch_builder(&self.rl_per_update_priority_fn);
launch
.arg(&self.gpu_replay.sample_indices_d)
.arg(&self.td_per_sample_d)
@@ -5872,7 +5910,7 @@ impl IntegratedTrainer {
shared_mem_bytes: 0,
};
let mut launch =
self.stream.launch_builder(&self.rl_per_tree_rebuild_fn);
self.train_stream.launch_builder(&self.rl_per_tree_rebuild_fn);
launch
.arg(&mut self.gpu_replay.priority_tree_d)
.arg(&cap_i);