From b06355456a865bc8505f598e0e7e533a2a11a290 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Tue, 26 May 2026 01:55:56 +0200 Subject: [PATCH] =?UTF-8?q?feat(rl):=20multi-stream=20=E2=80=94=20train=5F?= =?UTF-8?q?stream=20for=20PER=20sample/priority/rebuild?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 --- crates/ml-alpha/src/trainer/integrated.rs | 44 +++++++++++++++++++++-- 1 file changed, 41 insertions(+), 3 deletions(-) diff --git a/crates/ml-alpha/src/trainer/integrated.rs b/crates/ml-alpha/src/trainer/integrated.rs index cfa8fc3e5..49b1317b9 100644 --- a/crates/ml-alpha/src/trainer/integrated.rs +++ b/crates/ml-alpha/src/trainer/integrated.rs @@ -411,6 +411,12 @@ pub struct IntegratedTrainer { pub stream: Arc, + /// 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, + // ── 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 { let stream: Arc = dev.cuda_stream().context("integrated stream")?.clone(); + let train_stream: Arc = { + 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::()) 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);