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:
@@ -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.
|
||||
|
||||
Reference in New Issue
Block a user