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

fix(rl): zero loss accumulators + wire LobPtrs to 35 call sites

Browse Source 
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<i32>` for market_targets — incompatible with the
   immutable lob snapshot.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-05-28 14:12:57 +02:00
parent 8827fc93ee
commit bf88874007
1 changed files with 67 additions and 75 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -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::<f32>(), 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::<f32>(), 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::<f32>(), 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<IntegratedStepStats> {
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<u8> = 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<u8> = 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<u8> = 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<u8> = 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::<i32>() 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<IntegratedStepStats> {
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<u8> = 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<u8> = 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::<i32>()) 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<IntegratedStepStats> {
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<u8> = 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<u8> = 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<u8> = 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<u8> = 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::<i32>() 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);