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

fix(rl): move gates + log_pi OUTSIDE CUDA Graph capture

Browse Source 
The confidence gate and FRD gate were captured as no-ops during graph
warmup (step 1, before warmup threshold). Graph replay then never
fired them. Moving outside the capture region makes them execute as
individual kernel launches every step.

Result: Hold=100% from step 200-1000 (gate enforcing surfer Hold
default). Gate fires on nearly every non-Hold action until Q
confidence exceeds threshold=0.3.

Hold still drops after ~1500 steps as Q becomes confident — need
higher/adaptive threshold for sustained Hold protection.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-05-27 01:36:59 +02:00
parent 9395075a19
commit 8dfd49bda1
1 changed files with 31 additions and 48 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -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<u8> = 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<u8> = 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<u8> = 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<u8> = 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