fix(intervention-b): recapture train graph on boundary-state change

Task 1's attn_pool gate (commit 7dd953aa) was inside the captured-graph
region, so the host branch's decision was frozen at step-2 capture time.
In Sequential mode this would lock the captured graph to whichever
boundary state happened to hold at step 2, breaking subsequent file-
boundary detection — violating pearl_no_host_branches_in_captured_graph.

Fix: track the boundary state at capture time as a new
`train_graph_boundary_state: Option<bool>` field. On each step_batched,
if the current `last_seen_file_boundary` diverges from the captured
state, drop the cached graph and re-capture on the next step.

- Random mode: gate is unconditionally true; state never diverges;
  zero recaptures.
- Sequential mode: ~8 recaptures per epoch (one per file boundary).
  Each recapture costs the same as the original step-2 capture (~50ms).
  Negligible.

Verified clean cargo check -p ml-alpha --all-targets.
This commit is contained in:
jgrusewski
2026-05-21 11:39:29 +02:00
parent 7dd953aa2b
commit f22f3f9488

View File

@@ -548,6 +548,14 @@ pub struct PerceptionTrainer {
/// call runs uncaptured (warmup); second call captures; third+
/// replays. Eliminates ~155 individual kernel launches per step.
train_graph: Option<CudaGraph>,
/// Snapshot of `last_seen_file_boundary` at the time `train_graph` was
/// captured. The captured graph records the attn_pool decision at
/// capture time (per `pearl_no_host_branches_in_captured_graph`); if
/// the current boundary state diverges from the captured one, the
/// graph is dropped and re-captured on the next step. In Random mode
/// this never changes (always `true`), so no re-capture. In Sequential
/// mode this triggers a re-capture per file boundary (~8 per epoch).
train_graph_boundary_state: Option<bool>,
/// Captured CUDA Graph for the inference (`forward_only`) path. Same
/// three-state machine as `train_graph`: first call eager (warmup),
/// second call captures, third+ replays. Mirrors the forward chain
@@ -1339,6 +1347,7 @@ impl PerceptionTrainer {
logit_per_k_d: stream.alloc_zeros::<f32>(k * cfg.n_batch * N_HORIZONS)?,
stg_labels: unsafe { MappedF32Buffer::new(k * cfg.n_batch * N_HORIZONS) }.map_err(|e| anyhow::anyhow!("stg_labels: {e}"))?,
train_graph: None,
train_graph_boundary_state: None,
forward_graph: None,
forward_warmed: false,
cublas_warmed: false,
@@ -1649,8 +1658,22 @@ impl PerceptionTrainer {
}
}
// ── 2. Three-state machine: warmup (first call), capture
// (second call), replay (third+).
// ── 2. Four-state machine: warmup (first), capture (second),
// replay (third+), recapture-on-boundary-change.
// The captured graph records the attn_pool decision at
// capture time (per `pearl_no_host_branches_in_captured_graph`).
// In Sequential mode the decision toggles between sequences
// within vs at file boundaries, so the cached graph is
// invalidated whenever `last_seen_file_boundary` diverges
// from the captured value. Random mode never diverges
// (always `true`), so no recapture.
if let Some(captured_state) = self.train_graph_boundary_state {
if captured_state != self.last_seen_file_boundary {
// Boundary state changed — drop cached graph, fall through to recapture.
self.train_graph = None;
self.train_graph_boundary_state = None;
}
}
if self.train_graph.is_some() {
self.train_graph
.as_ref()
@@ -1687,6 +1710,7 @@ impl PerceptionTrainer {
"train end_capture returned None — no work captured"
))?;
self.train_graph = Some(graph);
self.train_graph_boundary_state = Some(self.last_seen_file_boundary);
}
// ── 3. Sync + read mapped-pinned loss.