From bf8887400750539d684b3570d5f01ad1558b5622 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Thu, 28 May 2026 14:12:57 +0200 Subject: [PATCH] fix(rl): zero loss accumulators + wire LobPtrs to 35 call sites MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Two fixes in the integrated trainer hot path. 1. l_pi / l_q monotonic accumulation (root cause: 90f178ae9 perf removal). The PPO surrogate fwd kernel atomicAdds into ss_pi_loss / ss_pi_loss_entropy, and dqn_distributional_q backward atomicAdds into ss_q_loss. Per the original `BUG FIX` comment in dqn_replay_step, these MUST be zeroed each step or they read as the running sum across all steps. 90f178ae9 (remove K-loop memsets) correctly removed the per-K-iter zeros for buffers that are OVERWRITTEN by backward kernels and reduce_axis0, but the comment in the dead code branch flagged the loss-scalar atomics as a known exception. The exception was lost when the if-false block ceased to execute. Smoke before fix at step 499: l_q=10.96 l_pi=1384.46 Smoke after fix at step 499: l_q=0.024 l_pi=6.65 (both bounded). v_loss writes via reduce_axis0_mapped (single-writer overwrite) and does not need zeroing. 2. LobPtrs dead-code refactor (root cause: 9c1b70ede merge resolution). The 9c1b70ede mega-graph commit introduced LobPtrs to cache pos_d().raw_ptr() / bid_px_d() / ask_px_d() / pos_bytes() once per step instead of dispatching through the RlLobBackend vtable at every raw_launch call site. The struct merged in but the 35 consumer sites did not because of structural divergence between phase-a's integrated.rs and clean-clamp's. Per `feedback_no_hiding`, wire it up or delete — wired across step_with_lobsim, step_with_lobsim_reward_and_train, and step_with_lobsim_gpu_body. Two sites kept the old `lobsim.pos_book_and_market_targets_mut()` pattern because that method returns disjoint borrows including a `&mut CudaSlice` for market_targets — incompatible with the immutable lob snapshot. Co-Authored-By: Claude Opus 4.7 --- crates/ml-alpha/src/trainer/integrated.rs | 142 ++++++++++------------ 1 file changed, 67 insertions(+), 75 deletions(-) diff --git a/crates/ml-alpha/src/trainer/integrated.rs b/crates/ml-alpha/src/trainer/integrated.rs index 819f18a78..4c30b10a2 100644 --- a/crates/ml-alpha/src/trainer/integrated.rs +++ b/crates/ml-alpha/src/trainer/integrated.rs @@ -364,9 +364,9 @@ struct HotPathPtrs { /// Cached raw device pointers for lobsim accessors. Eliminates per-call /// vtable dispatch + CudaSlice borrow overhead on `pos_d()`, `bid_px_d()`, -/// `ask_px_d()` which are called 5-10× per step but always return the same -/// pre-allocated device buffer. Extracted once at step entry, then threaded -/// through all raw_launch sites as plain `u64`. +/// `ask_px_d()` which are called many times per step but always return the +/// same pre-allocated device buffer. Extracted once at step entry, then +/// threaded through all raw_launch sites as plain `u64`. struct LobPtrs { pos: u64, bid_px: u64, @@ -4073,6 +4073,28 @@ impl IntegratedTrainer { // element, and buffers were alloc_zeros'd at construction. The // memsets replayed on every graph launch, burning ~36 × memset // bandwidth per step for zero correctness benefit. + // + // EXCEPTION: the 3 loss scalar accumulators (ss_pi_loss, + // ss_pi_loss_entropy, ss_q_loss) are written via `atomicAdd` by + // the PPO surrogate and DQN distributional Q kernels — see + // ppo_clipped_surrogate.cu:270 and dqn_distributional_q.cu:285. + // These MUST be zeroed each step or the per-step loss values + // accumulate monotonically across all steps (root cause of + // 90f178ae9 perf-regression: l_pi=1384 at step 499 was the sum + // of all per-step l_pi values, not the current step's loss). + // v_loss uses reduce_axis0_mapped (single-writer overwrite) and + // does not need zeroing. + { + let s = self.raw_stream; + unsafe { + raw_memset_d8_zero(self.ss_q_loss_dev_ptr, std::mem::size_of::(), s) + .map_err(|e| anyhow::anyhow!("zero ss_q_loss: {:?}", e))?; + raw_memset_d8_zero(self.ss_pi_loss_dev_ptr, std::mem::size_of::(), s) + .map_err(|e| anyhow::anyhow!("zero ss_pi_loss: {:?}", e))?; + raw_memset_d8_zero(self.ss_pi_loss_entropy_dev_ptr, std::mem::size_of::(), s) + .map_err(|e| anyhow::anyhow!("zero ss_pi_loss_entropy: {:?}", e))?; + } + } // ── Step 4 (Phase R7a): inputs are trainer-owned device buffers, // populated by step_with_lobsim's GPU pipeline before this call. @@ -5415,6 +5437,7 @@ impl IntegratedTrainer { next_snapshots: &[Mbp10RawInput], lobsim: &mut dyn RlLobBackend, ) -> Result { + let lob = LobPtrs::new(lobsim); let b_size = self.cfg.perception.n_batch; if b_size == 0 { anyhow::bail!("step_with_lobsim: empty batch (n_batch = 0)"); @@ -5753,17 +5776,15 @@ impl IntegratedTrainer { // ── Gates + log_pi: OUTSIDE graph capture so they fire every // step (not captured as no-ops during warmup). Device-side kernels. { - let pos_d_ref: &CudaSlice = lobsim.pos_d(); let b_size_i = b_size as i32; - let pos_bytes_i = lobsim.pos_bytes() as i32; let mut args = RawArgs::new(); args.push_ptr(self.actions_d.raw_ptr()); args.push_ptr(self.q_logits_d.raw_ptr()); args.push_ptr(self.atom_supports_d.raw_ptr()); - args.push_ptr(pos_d_ref.raw_ptr()); + args.push_ptr(lob.pos); args.push_ptr(self.isv_dev_ptr); args.push_i32(b_size_i); - args.push_i32(pos_bytes_i); + args.push_i32(lob.pos_bytes); let mut ptrs = args.build_arg_ptrs(); unsafe { raw_launch( @@ -5775,16 +5796,14 @@ impl IntegratedTrainer { } } { - let pos_d_ref: &CudaSlice = lobsim.pos_d(); let b_size_i = b_size as i32; - let pos_bytes_i = lobsim.pos_bytes() as i32; let mut args = RawArgs::new(); args.push_ptr(self.actions_d.raw_ptr()); args.push_ptr(self.frd_logits_d.raw_ptr()); - args.push_ptr(pos_d_ref.raw_ptr()); + args.push_ptr(lob.pos); args.push_ptr(self.isv_dev_ptr); args.push_i32(b_size_i); - args.push_i32(pos_bytes_i); + args.push_i32(lob.pos_bytes); let mut ptrs = args.build_arg_ptrs(); unsafe { raw_launch( @@ -5862,7 +5881,7 @@ impl IntegratedTrainer { // Reads current position_lots from lobsim.pos_d for Flat-from-* // actions; reads actions from self.actions_d (just written by // rl_action_kernel above). - let pos_bytes_i = lobsim.pos_bytes() as i32; + let pos_bytes_i = lob.pos_bytes; let b_size_i = b_size as i32; // Stage ts_ns to device buffer for graph-captured multires kernel. @@ -5929,16 +5948,15 @@ impl IntegratedTrainer { // Hold time = current_step - oldest_active_unit_entry_step // (NOT steps_since_done, which counts flat steps too). { - let pos_d_ref: &CudaSlice = lobsim.pos_d(); let grid_x = ((b_size as u32) + 31) / 32; let mut args = RawArgs::new(); args.push_ptr(self.actions_d.raw_ptr()); args.push_ptr(self.unit_active_d.raw_ptr()); args.push_ptr(self.unit_entry_step_d.raw_ptr()); - args.push_ptr(pos_d_ref.raw_ptr()); + args.push_ptr(lob.pos); args.push_ptr(self.isv_dev_ptr); args.push_i32(b_size_i); - args.push_i32(pos_bytes_i); + args.push_i32(lob.pos_bytes); let mut ptrs = args.build_arg_ptrs(); unsafe { raw_launch( @@ -5954,16 +5972,14 @@ impl IntegratedTrainer { // Runs BEFORE trail_mutate so the structural decay is applied // first, then agent's a7/a8 fine-tuning on top. { - let bid_px_d = lobsim.bid_px_d(); - let ask_px_d = lobsim.ask_px_d(); let mut args = RawArgs::new(); args.push_ptr(self.unit_trail_distance_d.raw_ptr()); args.push_ptr(self.unit_active_d.raw_ptr()); args.push_ptr(self.unit_entry_price_d.raw_ptr()); args.push_ptr(self.unit_initial_r_d.raw_ptr()); args.push_ptr(self.unit_lots_d.raw_ptr()); - args.push_ptr(bid_px_d.raw_ptr()); - args.push_ptr(ask_px_d.raw_ptr()); + args.push_ptr(lob.bid_px); + args.push_ptr(lob.ask_px); args.push_ptr(self.isv_dev_ptr); args.push_i32(b_size_i); let mut ptrs = args.build_arg_ptrs(); @@ -6003,12 +6019,10 @@ impl IntegratedTrainer { // FlatFromLong/Short on per-unit breach. Reads shared lobsim // best book for current mid. { - let bid_px_d = lobsim.bid_px_d(); - let ask_px_d = lobsim.ask_px_d(); let mut args = RawArgs::new(); args.push_ptr(self.actions_d.raw_ptr()); - args.push_ptr(bid_px_d.raw_ptr()); - args.push_ptr(ask_px_d.raw_ptr()); + args.push_ptr(lob.bid_px); + args.push_ptr(lob.ask_px); args.push_ptr(self.unit_active_d.raw_ptr()); args.push_ptr(self.unit_entry_price_d.raw_ptr()); args.push_ptr(self.unit_lots_d.raw_ptr()); @@ -6033,13 +6047,12 @@ impl IntegratedTrainer { // Reads pos_state for aggregate position; writes diag fired-count // to ISV slot 505. { - let pos_d_ref_heat: &CudaSlice = lobsim.pos_d(); let block = (b_size as u32).min(256); let grid = ((b_size as u32) + block - 1) / block; let smem = 256 * std::mem::size_of::() as u32; let mut args = RawArgs::new(); args.push_ptr(self.actions_d.raw_ptr()); - args.push_ptr(pos_d_ref_heat.raw_ptr()); + args.push_ptr(lob.pos); args.push_ptr(self.isv_dev_ptr); args.push_i32(b_size_i); args.push_i32(pos_bytes_i); @@ -6119,8 +6132,9 @@ impl IntegratedTrainer { snapshots: &[Mbp10RawInput], b_size: usize, ) -> Result { + let lob = LobPtrs::new(lobsim); let b_size_i = b_size as i32; - let pos_bytes_i = lobsim.pos_bytes() as i32; + let pos_bytes_i = lob.pos_bytes; // ── Graph B: post-fill reward/EMA/controller pipeline ───────── // ~20 kernels from extract_realized_pnl_delta through @@ -6146,10 +6160,9 @@ impl IntegratedTrainer { // reward_shaping + raw_rewards snapshot + recent_outcome_update // in a single per-batch kernel (7→1 launch). { - let pos_d_ref: &CudaSlice = lobsim.pos_d(); let grid_x = ((b_size as u32) + 31) / 32; let mut args = RawArgs::new(); - args.push_ptr(pos_d_ref.raw_ptr()); + args.push_ptr(lob.pos); args.push_ptr(self.prev_realized_pnl_d.raw_ptr()); args.push_ptr(self.prev_position_lots_d.raw_ptr()); args.push_ptr(self.rewards_d.raw_ptr()); @@ -6244,11 +6257,7 @@ impl IntegratedTrainer { // Trade context features — derived from unit state + current mid. { - let pos_d_ref: &CudaSlice = lobsim.pos_d(); - let bid_px_d = lobsim.bid_px_d(); - let ask_px_d = lobsim.ask_px_d(); let b_size_i = b_size as i32; - let pos_bytes_i = lobsim.pos_bytes() as i32; let grid_x = ((b_size as u32) + 31) / 32; let mut args = RawArgs::new(); args.push_ptr(self.trade_context_d.raw_ptr()); @@ -6258,12 +6267,12 @@ impl IntegratedTrainer { args.push_ptr(self.unit_initial_r_d.raw_ptr()); args.push_ptr(self.unit_peak_unrealized_r_d.raw_ptr()); args.push_ptr(self.unit_lots_d.raw_ptr()); - args.push_ptr(pos_d_ref.raw_ptr()); - args.push_ptr(bid_px_d.raw_ptr()); - args.push_ptr(ask_px_d.raw_ptr()); + args.push_ptr(lob.pos); + args.push_ptr(lob.bid_px); + args.push_ptr(lob.ask_px); args.push_ptr(self.isv_dev_ptr); args.push_i32(b_size_i); - args.push_i32(pos_bytes_i); + args.push_i32(lob.pos_bytes); let mut ptrs = args.build_arg_ptrs(); unsafe { raw_launch( @@ -6277,8 +6286,6 @@ impl IntegratedTrainer { // Multi-resolution streaming features — time-weighted EMA at 3 horizons. { - let bid_px_d = lobsim.bid_px_d(); - let ask_px_d = lobsim.ask_px_d(); let b_size_i = b_size as i32; let grid_x = ((b_size as u32) + 31) / 32; let mut args = RawArgs::new(); @@ -6286,10 +6293,10 @@ impl IntegratedTrainer { args.push_ptr(self.multires_output_d.raw_ptr()); args.push_ptr(self.multires_prev_mid_d.raw_ptr()); args.push_ptr(self.multires_prev_ts_ns_d.raw_ptr()); - args.push_ptr(bid_px_d.raw_ptr()); - args.push_ptr(ask_px_d.raw_ptr()); - args.push_ptr(bid_px_d.raw_ptr()); // bid_sz uses bid_px - args.push_ptr(ask_px_d.raw_ptr()); // ask_sz uses ask_px + args.push_ptr(lob.bid_px); + args.push_ptr(lob.ask_px); + args.push_ptr(lob.bid_px); // bid_sz uses bid_px + args.push_ptr(lob.ask_px); // ask_sz uses ask_px args.push_ptr(self.isv_dev_ptr); args.push_ptr(self.ts_ns_d.raw_ptr()); args.push_i32(b_size_i); @@ -6617,23 +6624,19 @@ impl IntegratedTrainer { // mid-price into ring buffer and update peak (max for long, // min for short) while a position is open. Flat resets ring. { - let pos_d_ref = lobsim.pos_d(); - let bid_px_d = lobsim.bid_px_d(); - let ask_px_d = lobsim.ask_px_d(); let grid_x = ((b_size as u32) + 31) / 32; let b_size_i = b_size as i32; - let pos_bytes_i = lobsim.pos_bytes() as i32; let mut args = RawArgs::new(); - args.push_ptr(pos_d_ref.raw_ptr()); - args.push_ptr(bid_px_d.raw_ptr()); - args.push_ptr(ask_px_d.raw_ptr()); + args.push_ptr(lob.pos); + args.push_ptr(lob.bid_px); + args.push_ptr(lob.ask_px); args.push_ptr(self.hindsight.mid_ring_d.raw_ptr()); args.push_ptr(self.hindsight.ring_write_idx_d.raw_ptr()); args.push_ptr(self.hindsight.peak_mid_d.raw_ptr()); args.push_ptr(self.hindsight.entry_mid_d.raw_ptr()); args.push_ptr(self.hindsight.position_dir_d.raw_ptr()); args.push_i32(b_size_i); - args.push_i32(pos_bytes_i); + args.push_i32(lob.pos_bytes); let mut ptrs = args.build_arg_ptrs(); unsafe { raw_launch( @@ -6804,13 +6807,11 @@ impl IntegratedTrainer { // (ISV[552] steps). If holding would have been better, inject // a synthetic "should-have-held-longer" transition. { - let bid_px_d = lobsim.bid_px_d(); - let ask_px_d = lobsim.ask_px_d(); let smem = (2 * 256 * std::mem::size_of::()) as u32; let cap_i = self.gpu_replay.capacity as i32; let mut args = RawArgs::new(); - args.push_ptr(bid_px_d.raw_ptr()); - args.push_ptr(ask_px_d.raw_ptr()); + args.push_ptr(lob.bid_px); + args.push_ptr(lob.ask_px); args.push_ptr(self.isv_dev_ptr); args.push_ptr(self.hindsight.closed_ring_d.raw_ptr()); args.push_ptr(self.hindsight.closed_h_t_d.raw_ptr()); @@ -7386,6 +7387,7 @@ impl IntegratedTrainer { b_size: usize, seq_len: usize, ) -> Result { + let lob = LobPtrs::new(lobsim); // ── Graph A: pre-snapshot kernel pipeline ────────────────────── if self.prefill_graph.is_some() && !self.mega_graph_enabled { unsafe { @@ -7584,17 +7586,15 @@ impl IntegratedTrainer { // ── Gates + log_pi: OUTSIDE graph capture (GPU data path). { - let pos_d_ref: &CudaSlice = lobsim.pos_d(); let b_size_i = b_size as i32; - let pos_bytes_i = lobsim.pos_bytes() as i32; let mut args = RawArgs::new(); args.push_ptr(self.actions_d.raw_ptr()); args.push_ptr(self.q_logits_d.raw_ptr()); args.push_ptr(self.atom_supports_d.raw_ptr()); - args.push_ptr(pos_d_ref.raw_ptr()); + args.push_ptr(lob.pos); args.push_ptr(self.isv_dev_ptr); args.push_i32(b_size_i); - args.push_i32(pos_bytes_i); + args.push_i32(lob.pos_bytes); let mut ptrs = args.build_arg_ptrs(); unsafe { raw_launch( @@ -7606,16 +7606,14 @@ impl IntegratedTrainer { } } { - let pos_d_ref: &CudaSlice = lobsim.pos_d(); let b_size_i = b_size as i32; - let pos_bytes_i = lobsim.pos_bytes() as i32; let mut args = RawArgs::new(); args.push_ptr(self.actions_d.raw_ptr()); args.push_ptr(self.frd_logits_d.raw_ptr()); - args.push_ptr(pos_d_ref.raw_ptr()); + args.push_ptr(lob.pos); args.push_ptr(self.isv_dev_ptr); args.push_i32(b_size_i); - args.push_i32(pos_bytes_i); + args.push_i32(lob.pos_bytes); let mut ptrs = args.build_arg_ptrs(); unsafe { raw_launch( @@ -7737,7 +7735,7 @@ impl IntegratedTrainer { } } - let pos_bytes_i = lobsim.pos_bytes() as i32; + let pos_bytes_i = lob.pos_bytes; let b_size_i = b_size as i32; // ── Graph A2: post-snapshot / pre-fill kernel pipeline ───────── @@ -7779,13 +7777,12 @@ impl IntegratedTrainer { // Min hold check — true hold time from oldest active unit // (see eval-side launch comment above for rationale). { - let pos_d_ref: &CudaSlice = lobsim.pos_d(); let grid_x = ((b_size as u32) + 31) / 32; let mut args = RawArgs::new(); args.push_ptr(self.actions_d.raw_ptr()); args.push_ptr(self.unit_active_d.raw_ptr()); args.push_ptr(self.unit_entry_step_d.raw_ptr()); - args.push_ptr(pos_d_ref.raw_ptr()); + args.push_ptr(lob.pos); args.push_ptr(self.isv_dev_ptr); args.push_i32(b_size_i); args.push_i32(pos_bytes_i); @@ -7802,16 +7799,14 @@ impl IntegratedTrainer { // Asymmetric trail decay { - let bid_px_d = lobsim.bid_px_d(); - let ask_px_d = lobsim.ask_px_d(); let mut args = RawArgs::new(); args.push_ptr(self.unit_trail_distance_d.raw_ptr()); args.push_ptr(self.unit_active_d.raw_ptr()); args.push_ptr(self.unit_entry_price_d.raw_ptr()); args.push_ptr(self.unit_initial_r_d.raw_ptr()); args.push_ptr(self.unit_lots_d.raw_ptr()); - args.push_ptr(bid_px_d.raw_ptr()); - args.push_ptr(ask_px_d.raw_ptr()); + args.push_ptr(lob.bid_px); + args.push_ptr(lob.ask_px); args.push_ptr(self.isv_dev_ptr); args.push_i32(b_size_i); let mut ptrs = args.build_arg_ptrs(); @@ -7846,12 +7841,10 @@ impl IntegratedTrainer { // Trail stop check { - let bid_px_d = lobsim.bid_px_d(); - let ask_px_d = lobsim.ask_px_d(); let mut args = RawArgs::new(); args.push_ptr(self.actions_d.raw_ptr()); - args.push_ptr(bid_px_d.raw_ptr()); - args.push_ptr(ask_px_d.raw_ptr()); + args.push_ptr(lob.bid_px); + args.push_ptr(lob.ask_px); args.push_ptr(self.unit_active_d.raw_ptr()); args.push_ptr(self.unit_entry_price_d.raw_ptr()); args.push_ptr(self.unit_lots_d.raw_ptr()); @@ -7872,13 +7865,12 @@ impl IntegratedTrainer { // Position heat cap { - let pos_d_ref_heat: &CudaSlice = lobsim.pos_d(); let block = (b_size as u32).min(256); let grid = ((b_size as u32) + block - 1) / block; let smem = 256 * std::mem::size_of::() as u32; let mut args = RawArgs::new(); args.push_ptr(self.actions_d.raw_ptr()); - args.push_ptr(pos_d_ref_heat.raw_ptr()); + args.push_ptr(lob.pos); args.push_ptr(self.isv_dev_ptr); args.push_i32(b_size_i); args.push_i32(pos_bytes_i);