diff --git a/crates/ml-alpha/src/trainer/integrated.rs b/crates/ml-alpha/src/trainer/integrated.rs index 45f1ff9ac..b541062f4 100644 --- a/crates/ml-alpha/src/trainer/integrated.rs +++ b/crates/ml-alpha/src/trainer/integrated.rs @@ -5445,11 +5445,20 @@ impl IntegratedTrainer { } } - // ── Confidence gate — override opening actions to Hold when C51 - // distributional Q has high uncertainty for the chosen action. - // Fires AFTER π samples an action but BEFORE trail/heat/market - // pipeline. Reads pos_state from lobsim (pre-snapshot = current - // position) to only gate on flat entries. + if capturing_prefill { + let graph = self + .stream + .end_capture( + CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH, + ) + .context("prefill end_capture")? + .ok_or_else(|| anyhow::anyhow!("prefill end_capture returned None"))?; + self.prefill_graph = Some(graph); + } + } // end else (warmup / capture dispatch) + + // ── 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; @@ -5472,8 +5481,6 @@ impl IntegratedTrainer { ).map_err(|e| anyhow::anyhow!("rl_confidence_gate: {:?}", e))?; } } - - // FRD gate — same warmup as confidence gate. { let pos_d_ref: &CudaSlice = lobsim.pos_d(); let b_size_i = b_size as i32; @@ -5495,11 +5502,6 @@ impl IntegratedTrainer { ).map_err(|e| anyhow::anyhow!("rl_frd_gate: {:?}", e))?; } } - - // log pi_old at the (possibly gate-overridden) action — AFTER - // confidence gate + FRD gate so that log π_old matches the action - // that actually executes, giving PPO a correct importance ratio - // when a gate overrides the sampled action to Hold. { let grid_x = ((b_size as u32) + 31) / 32; let b_size_i = b_size as i32; @@ -5515,22 +5517,10 @@ impl IntegratedTrainer { (grid_x.max(1), 1, 1), (32, 1, 1), 0, self.raw_stream, &mut ptrs[..args.len()], - ).map_err(|e| anyhow::anyhow!("log_pi_at_action: {:?}", e))?; + ).map_err(|e| anyhow::anyhow!("log_pi_at_action (post-gate): {:?}", e))?; } } - if capturing_prefill { - let graph = self - .stream - .end_capture( - CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH, - ) - .context("prefill end_capture")? - .ok_or_else(|| anyhow::anyhow!("prefill end_capture returned None"))?; - self.prefill_graph = Some(graph); - } - } // end else (warmup / capture dispatch) - // ── Step 5 (Phase R6): GPU-pure env step. ───────────────────── // No per-batch host loop — actions land in lobsim's // market_targets_d via the `actions_to_market_targets` kernel @@ -7077,7 +7067,19 @@ impl IntegratedTrainer { } } - // Confidence gate + if capturing_prefill { + let graph = self + .stream + .end_capture( + CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH, + ) + .context("prefill end_capture")? + .ok_or_else(|| anyhow::anyhow!("prefill end_capture returned None"))?; + self.prefill_graph = Some(graph); + } + } // end else (warmup / capture dispatch) + + // ── 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; @@ -7097,11 +7099,9 @@ impl IntegratedTrainer { (b_size as u32, 1, 1), (1, 1, 1), 0, self.raw_stream, &mut ptrs[..args.len()], - ).map_err(|e| anyhow::anyhow!("rl_confidence_gate: {:?}", e))?; + ).map_err(|e| anyhow::anyhow!("rl_confidence_gate (gpu): {:?}", e))?; } } - - // FRD gate { let pos_d_ref: &CudaSlice = lobsim.pos_d(); let b_size_i = b_size as i32; @@ -7120,14 +7120,9 @@ impl IntegratedTrainer { (b_size as u32, 1, 1), (1, 1, 1), 0, self.raw_stream, &mut ptrs[..args.len()], - ).map_err(|e| anyhow::anyhow!("rl_frd_gate: {:?}", e))?; + ).map_err(|e| anyhow::anyhow!("rl_frd_gate (gpu): {:?}", e))?; } } - - // log pi_old at the (possibly gate-overridden) action — AFTER - // confidence gate + FRD gate so that log π_old matches the action - // that actually executes, giving PPO a correct importance ratio - // when a gate overrides the sampled action to Hold. { let grid_x = ((b_size as u32) + 31) / 32; let b_size_i = b_size as i32; @@ -7143,22 +7138,10 @@ impl IntegratedTrainer { (grid_x.max(1), 1, 1), (32, 1, 1), 0, self.raw_stream, &mut ptrs[..args.len()], - ).map_err(|e| anyhow::anyhow!("log_pi_at_action: {:?}", e))?; + ).map_err(|e| anyhow::anyhow!("log_pi_at_action (post-gate gpu): {:?}", e))?; } } - if capturing_prefill { - let graph = self - .stream - .end_capture( - CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH, - ) - .context("prefill end_capture")? - .ok_or_else(|| anyhow::anyhow!("prefill end_capture returned None"))?; - self.prefill_graph = Some(graph); - } - } // end else (warmup / capture dispatch) - // ── Step 5: GPU-pure env step (GPU data path). ───────────────── // // apply_snapshot from device: copy batch 0's last-snapshot book