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);