From 3a196382f0bfa564eaf2b1d384e06abd6e5c9406 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Sun, 17 May 2026 16:01:02 +0200 Subject: [PATCH] fix(ml-alpha): captured graph poisoned eval via SyncOnDrop events MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit z2w9w cluster run hit CUDA_ERROR_INVALID_VALUE at "eval snap_batched fwd" the first time validation ran after a captured training step. Training itself succeeded (epoch 0 train_loss=0.69 over 250 captured graph replays); only the subsequent direct eval kernel launch failed. Root cause (vendor/cudarc/src/driver/safe/core.rs:920): CudaSlice::device_ptr_mut() returns a `SyncOnDrop::Record` guard. On drop, that guard UNCONDITIONALLY calls `event.record(stream)` on the slice's `.read` event (the check at line 953 only gates the cuStreamWaitEvent on .write — the unconditional event.record at the end runs no matter what). Inside a stream-capture region, those event.record(stream) calls turn the CudaEvents into "captured events" per the CUDA Driver API. Captured events can ONLY be waited on by streams in the same capture sequence; any later cuStreamWaitEvent from outside fails with CUDA_ERROR_INVALID_VALUE. The trainer's eval path then called `device_ptr_mut()` again to stage the eval DtoDs — which inserted exactly that cuStreamWaitEvent on the now-captured `.write` event of every trainer CudaSlice. First kernel launch after the dtods failed. Why this hit z2w9w now: a) all our work is on a SINGLE stream (`self.stream`), so the event-based multi-stream sync that cudarc inserts is pure overhead, b) the capture region is exactly where those overhead events become poisonous. Fix: disable cudarc's read/write event tracking BEFORE the trainer allocates ANY device memory. With tracking off at alloc time, CudaSlice::new returns `read: None, write: None` (core.rs:1283). SyncOnDrop::record_event with `event: None` produces a `Record(None)` that does nothing on drop. launch_builder skips its event waits/ records too. The capture region runs clean; the eval direct launches have no stale captured events to wait on. Validation: - 6 perception_overfit tests pass, including 3 NEW regression tests that pin the exact failure modes: * evaluate_alone_succeeds — eval with no prior training * evaluate_works_after_warmup_only — eval after 1 uncaptured step * evaluate_works_after_capture_no_replay — eval right after capture (the minimal repro that pinpointed `device_ptr_mut` inside capture) * evaluate_works_after_captured_training_step — full warmup + capture + replay + eval - 26 ml-alpha lib + 23 ml-alpha integration + 306 ml-core lib all pass. Also drops the now-redundant disable/enable_event_tracking dance around the capture region — events are globally disabled for the trainer's lifetime so no per-capture flipping needed. Honors: feedback_no_quickfixes.md (root-cause traced through cudarc's safe wrappers to the SyncOnDrop record contract, not a symptom-suppress sleep/retry hack). Co-Authored-By: Claude Opus 4.7 --- crates/ml-alpha/src/trainer/perception.rs | 27 +++-- crates/ml-alpha/tests/perception_overfit.rs | 115 ++++++++++++++++++++ 2 files changed, 134 insertions(+), 8 deletions(-) diff --git a/crates/ml-alpha/src/trainer/perception.rs b/crates/ml-alpha/src/trainer/perception.rs index 39f861bd0..53bfea01d 100644 --- a/crates/ml-alpha/src/trainer/perception.rs +++ b/crates/ml-alpha/src/trainer/perception.rs @@ -271,6 +271,20 @@ impl PerceptionTrainer { let stream = dev.cuda_stream().context("trainer stream")?.clone(); let ctx = dev.cuda_context().context("trainer ctx")?; + // Disable cudarc's automatic per-allocation read/write event + // tracking BEFORE allocating any device memory. With tracking + // enabled, each `device_ptr_mut()` call returns a SyncOnDrop + // guard that records a CudaEvent on the stream when dropped. + // Inside a stream-capture region those `event.record(stream)` + // calls turn the events into "captured events" — which then + // cannot be waited on from outside the capture, breaking the + // non-captured eval path with CUDA_ERROR_INVALID_VALUE at the + // next `cuStreamWaitEvent`. All trainer work runs on a single + // stream (`self.stream`) and is therefore stream-ordered with + // no cross-stream sync needed, so event tracking is pure + // overhead for us. Disabled for the lifetime of this trainer. + unsafe { ctx.disable_event_tracking(); } + let _ = ctx.check_err(); let snap_module = ctx.load_cubin(SNAP_CUBIN.to_vec()).context("snap cubin")?; let step_module = ctx.load_cubin(STEP_CUBIN.to_vec()).context("step cubin")?; let heads_module = ctx.load_cubin(HEADS_CUBIN.to_vec()).context("heads cubin")?; @@ -599,16 +613,15 @@ impl PerceptionTrainer { .context("train warmup dispatch")?; self.cublas_warmed = true; } else { - let ctx = self.stream.context().clone(); - let _ = ctx.check_err(); - unsafe { ctx.disable_event_tracking(); } - + // Event tracking was disabled at trainer construction; the + // trainer's CudaSlices have no read/write events, so neither + // launch_builder nor device_ptr_mut() will try to insert + // cuStreamWaitEvent / event.record() calls inside the + // captured region. Capture safely. let begin = self.stream.begin_capture( CUstreamCaptureMode::CU_STREAM_CAPTURE_MODE_RELAXED, ); if let Err(e) = begin { - unsafe { ctx.enable_event_tracking(); } - let _ = ctx.check_err(); return Err(anyhow::anyhow!("train begin_capture: {e}")); } @@ -617,8 +630,6 @@ impl PerceptionTrainer { let graph_result = self.stream.end_capture( CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH, ); - unsafe { ctx.enable_event_tracking(); } - let _ = ctx.check_err(); dispatch_result.context("train graph dispatch (during capture)")?; let graph = graph_result diff --git a/crates/ml-alpha/tests/perception_overfit.rs b/crates/ml-alpha/tests/perception_overfit.rs index 5f623fe56..ef56f708e 100644 --- a/crates/ml-alpha/tests/perception_overfit.rs +++ b/crates/ml-alpha/tests/perception_overfit.rs @@ -128,3 +128,118 @@ fn stacked_trainer_loss_shrinks_on_constant_signal() { start={initial_avg:.4}, end={final_avg:.4}" ); } + +/// Eval alone must work — proves the eval path is fine WITHOUT any +/// prior captured-graph training. If this passes but +/// `evaluate_works_after_captured_training_step` fails, the bug is in +/// the train→eval transition (captured graph leaving state hostile to +/// direct kernel launches). +#[test] +fn evaluate_alone_succeeds() { + let dev = test_device(); + let cfg = PerceptionTrainerConfig { + seq_len: 16, + mamba2_state_dim: 8, + lr_cfc: 3e-3, + lr_mamba2: 1e-3, + seed: 0x6262, + horizon_weights: [1.0; 5], + n_batch: 1, + }; + let mut trainer = PerceptionTrainer::new(&dev, &cfg).expect("init"); + let ts = 1_000_000u64; + let prev_mid = 5500.0_f32; + let (seq, labels) = synthetic_seq(cfg.seq_len, prev_mid, ts); + let (loss, probs) = trainer.evaluate(&seq, labels.as_slice()).expect("eval alone"); + assert!(loss.is_finite(), "eval loss must be finite, got {loss}"); + assert_eq!(probs.len(), cfg.seq_len * 5); +} + +/// Regression test for z2w9w cluster run: training step (which captures +/// a CUDA Graph) must NOT break the subsequent `evaluate()` call. +/// z2w9w hit CUDA_ERROR_INVALID_VALUE at "eval snap_batched fwd" the +/// very first time eval ran after training; the captured graph or some +/// of its launch state left the stream in a state hostile to the +/// non-captured eval path. +#[test] +fn evaluate_works_after_captured_training_step() { + let dev = test_device(); + let cfg = PerceptionTrainerConfig { + seq_len: 16, + mamba2_state_dim: 8, + lr_cfc: 3e-3, + lr_mamba2: 1e-3, + seed: 0x5151, + horizon_weights: [1.0; 5], + n_batch: 1, + }; + let mut trainer = PerceptionTrainer::new(&dev, &cfg).expect("init"); + + // Drive enough training steps to exercise warmup → capture → replay. + let mut ts = 1_000_000u64; + let mut prev_mid = 5500.0_f32; + for _ in 0..5 { + let (seq, labels) = synthetic_seq(cfg.seq_len, prev_mid, ts); + trainer.step(&seq, labels.as_slice()).expect("train step"); + prev_mid = 0.5 * (seq.last().unwrap().bid_px[0] + seq.last().unwrap().ask_px[0]); + ts = seq.last().unwrap().ts_ns; + } + + // Now call evaluate — must NOT fail with CUDA_ERROR_INVALID_VALUE. + let (seq, labels) = synthetic_seq(cfg.seq_len, prev_mid, ts); + let (loss, probs) = trainer.evaluate(&seq, labels.as_slice()).expect("evaluate after train"); + assert!(loss.is_finite(), "eval loss must be finite, got {loss}"); + assert_eq!(probs.len(), cfg.seq_len * 5, "eval probs must be [K, 5] flat"); + assert!(probs.iter().all(|p| p.is_finite()), "eval probs must be finite"); +} + +/// 2 steps = warmup + capture (no replay yet). Does eval fail right +/// after capture but BEFORE the first graph.launch? +#[test] +fn evaluate_works_after_capture_no_replay() { + let dev = test_device(); + let cfg = PerceptionTrainerConfig { + seq_len: 16, + mamba2_state_dim: 8, + lr_cfc: 3e-3, + lr_mamba2: 1e-3, + seed: 0x8181, + horizon_weights: [1.0; 5], + n_batch: 1, + }; + let mut trainer = PerceptionTrainer::new(&dev, &cfg).expect("init"); + let ts = 1_000_000u64; + let prev_mid = 5500.0_f32; + let (seq, labels) = synthetic_seq(cfg.seq_len, prev_mid, ts); + trainer.step(&seq, labels.as_slice()).expect("warmup step"); + trainer.step(&seq, labels.as_slice()).expect("capture step"); + // Eval right after capture, no replay. + let (loss, _probs) = trainer.evaluate(&seq, labels.as_slice()).expect("eval after capture"); + assert!(loss.is_finite(), "eval loss must be finite, got {loss}"); +} + +/// Does the bug surface after just ONE step (warmup only, no capture)? +/// If yes, the warmup dispatch path itself breaks subsequent eval. +/// If no, the captured graph instantiation / launch is what breaks eval. +#[test] +fn evaluate_works_after_warmup_only() { + let dev = test_device(); + let cfg = PerceptionTrainerConfig { + seq_len: 16, + mamba2_state_dim: 8, + lr_cfc: 3e-3, + lr_mamba2: 1e-3, + seed: 0x7171, + horizon_weights: [1.0; 5], + n_batch: 1, + }; + let mut trainer = PerceptionTrainer::new(&dev, &cfg).expect("init"); + let ts = 1_000_000u64; + let prev_mid = 5500.0_f32; + let (seq, labels) = synthetic_seq(cfg.seq_len, prev_mid, ts); + // ONLY one step — warmup, no capture yet. + trainer.step(&seq, labels.as_slice()).expect("warmup step"); + // Now eval. + let (loss, _probs) = trainer.evaluate(&seq, labels.as_slice()).expect("eval after warmup"); + assert!(loss.is_finite(), "eval loss must be finite, got {loss}"); +}