Autopsy: the SiTU arm OOMs mid-run (steps 200-400) — work log

Lane: situ-oom (autopsy of the ab2000_situ_s{1,2,3} deaths). Worktree wt/situ-oom, off c8df252. The arm died on a SOLO card; the other five arms of the same wave completed 2000 steps. This file names the defect precisely.

Verdict up front: it is not a leak and not a pool bug. The arm ran the TENSOR path of situ_glu (the fused kernel is compiled out — crates/dormouse-core/Cargo.toml:39 deliberately ships burn-situ with default features only), the tensor path keeps ~11 live 32 MiB intermediates per call, the wave ran with checkpointing: Disabled so all of it stayed live until backward, and the arm carried ~+4.2 GiB of live activations per step over the silu control — which put the pool’s high-water past the 16 GB line somewhere in steps 200-400, seed-dependent. The 2026-09-30 review priced the tensor path at “~1% of a step” (situ-2026-09-30.md §7) — by LAUNCHES, under a checkpointing assumption the wave did not run.

1. Forensics (all CPU-side, from ~/logs/)

Three logs, all preset=/tmp/opencode/presets/situ.toml (= configs/small.toml verbatim + use_situ = true; batch 8, seq 512, max_iter = 4, n_experts = 3, d_ffn = 2048, d_model 768, fp32, in-VRAM Engram 25k rows, aux on). params = 9 590 670. Eval window at batch 8 = 81 920 B (the moe log’s eval line prints it; situ never reached an eval — deaths at 200/300/400, first eval at 500).

log	attempts	died at	reserve-fail sizes	where the fails are
ab2000_situ_s1.log	2 (appended)	step 400, then step 300	1 GiB ×7, 32 MiB ×104, 64 MiB ×49	1 GiB before step 0 (benign, see below); 32/64 MiB all AFTER step 200
ab2000_situ_s2.log	≥1, log truncated (25 lines, no header/step lines)	unknown, before step 0 printed	1 GiB ×6	init
ab2000_situ_s3.log	1	step 200	1 GiB ×10, 32 MiB ×8	1 GiB init; 32 MiB after step 100

Death message (all): train failed: step N: device error (loss is not a scalar) — that is crates/dormouse-train/src/lib.rs:1317: the try_into_scalar on the raw loss failed, i.e. the loss tensor’s storage was already dead. Before that, the compute thread had panicked dozens of times:

thread 'DSD-0-0' panicked at vendor/cubecl-fix/cubecl-cuda/src/compute/server.rs:144:35:
failed to reserve 33554432 bytes of device memory: out of device memory allocating 33554432 bytes

The control arms fail reserves too — but only at init. moe s1/s2/s3: 13/15/18 reserve failures, ALL before step 0; attnres s1: 27, all before step 0; mhc, fbs1: same shape. All five completed 2000 steps. The 1 GiB init failures are the runtime’s own calibration/autotune working-set sweeps (cubecl-runtime/src/throughput/curve.rs:185 sweeps up to 1 << 30) — they fail on a tight card, a candidate is refused, the run proceeds. They are noise.

The situ arm’s 32/64 MiB failures are DIFFERENT: they start mid-run (s3 after step 100, s1 after step 200), repeat every step (104 + 49 in s1 across ~150 steps), and end the run the first time the loss read lands on a dead tensor.

2. Why those exact sizes: the tensors are 32 MiB and 64 MiB

N = batch 8 × seq 512 = 4096 positions.

• [4096, 2048] f32 = 4096·2048·4 = 33 554 432 B — exactly the failing 32 MiB. This is every [N, f] intermediate of the expert FFN.
• [4096, 4096] f32 = 67 108 864 B — the failing 64 MiB. This is mid, the gate_up output, which use_situ doubles to d -> 2f (loop_block.rs:115).

So the panicking reserves are the arm’s own activations. The pool asked for another 32 MiB sliced page to hold yet another intermediate and cudaMalloc said no.

3. The allocation profile: tensor path vs fused vs control

crates/dormouse-core/Cargo.toml:39:

burn-situ = { path = "../../vendor/dormouse-fused/crates/burn-situ" }

Default features only, deliberately — the in-tree comment (and situ-2026-09-30.md §7) says the fused kernel “has never executed in any job in this repo” and enabling it wants a CUDA gate first. Consequence: burn_situ::situ_glu’s #[cfg(all(feature = "autodiff", feature = "cuda"))] probes (lib.rs:77-106) are compiled OUT, and every call in training runs situ_glu_tensor (fused_situ.rs:372-383).

Per call, on [4096, 4096] input, situ_glu_tensor materializes: two slice copies (gate, up, 32 MiB each), then per branch div/tanh/mul_scalar (3 × 32 MiB × 2 branches), sigmoid (32), two muls (32 + 32) — ~11 op outputs. The wave ran checkpointing: Disabled (every log header says checkpointing: Disabled), so under burn’s autodiff every op output stays live until backward. Per step: 3 experts × 4 iterations × ~11 × 32 MiB ≈ 4.2 GiB of live situ intermediates, plus the doubled mid (+32 MiB × 12 calls = +0.4 GiB over control).

The arms the wave actually compares against:

path	live intermediates per call	per step (12 calls)
silu control (activation::silu(mid))	1 × 32 MiB	0.4 GiB
situ tensor path (what ran)	~11 × 32 MiB	4.6 GiB
situ fused path (compiled out)	1 × 32 MiB out (+64 MiB dgu at backward)	~0.4-1.2 GiB

The control’s full step graph (KDA + JEPA second forward + everything) fits in what the desktop leaves of 16 GB. +4.2 GiB does not. The pool’s high-water crosses the line at a seed- and timing-dependent step (200, 300, 400 — the same spread repeats across s1’s own two attempts), and from then on every step fails a page allocation until the loss read fails and the trainer exits.

Two things made this invisible:

1. §7’s cost estimate priced the wrong resource. “~66 extra launches, order 1-2 ms” counted launches, and was written while the trainer ran BalancedCheckpointing (the review’s own §7 text says so). Under checkpointing the intermediates are recomputed, not retained, so the memory multiplier never materialized. This wave ran checkpointing: Disabled (visible in every backend=Device<...> log line), and under NoCheckpointing the ~11× live-set multiplier is the whole story. A cost estimate in launches is not a cost estimate for a memory defect.
2. The arm has no externally visible evidence of which path ran (ADR-0019 class). probe::SITU counts CALLS (loop_block.rs:928), not arrivals at the fused kernel; burn-situ’s own seam_counts() (fused_situ.rs:42) is never read by the trainer; the eval line’s fused kda=... norm=... engram=... fields have no situ counterpart. A reader of the training log cannot tell tensor path from fused — §1.1’s “if you add an arm, add its counter”, unheeded for the path dimension.

4. What it is NOT

• Not the fused kernel — it is compiled out; nothing fused executed.
• Not a per-step leak. Every allocation is a fixed size; nothing grows with the step index. The growth into the wall is the pool’s high-water + page inventory accumulating past the card’s free VRAM, which every arm does; the situ arm just starts ~4 GiB closer to the edge.
• Not the Engram, not the 1 GiB init failures — those are the runtime’s throughput sweeps (throughput/curve.rs:185), benign, and present in every living arm’s log too.
• Not batch-8-specific arithmetic error — the numbers close exactly: 4096·2048·4 = 33 554 432 = the literal in the panic.

5. Fix

FIX-1 (the real one) — enable the fused path, gated exactly as §7 demanded. Two lines, ARMED on the CUDA gate:

• crates/dormouse-core/Cargo.toml:64-73: add "burn-situ/cuda", "burn-situ/autodiff" to the cuda feature list (same shape as burn-rmsnorm/cuda on the line above).
• The gate §7 asks for EXISTS and is written: burn-situ’s situ_fused_matches_tensor + fused_backward_matches_tensor_backward + situ_grad_matches_finite_difference (fused_situ.rs:533, :570, :627), all #[ignore]d pending a live context. They have never run.

I did NOT flip the feature: the gate has zero GPU executions on record, and flipping first would put an never-executed backward kernel into the training path — the 8fa5d4c defect class this repo keeps re-learning. The flip lands in the commit that runs the gate green, not before.

Gate + repro script (ARMED — the card was busy with another lane’s graph_seam_cuda run during this autopsy; run on a free card, pgrep -ax train empty, free -g avail ≥ 25):

# 1. The gate (vendor workspace; ~one cold cuda build):
cd vendor/dormouse-fused
BURN_DEVICE=cuda cargo test -p burn-situ --features cuda,autodiff -- --ignored --nocapture
#    situ_fused_matches_tensor, fused_backward_matches_tensor_backward,
#    situ_grad_matches_finite_difference must all pass; situ_bench prints the speedup.

# 2. Flip the feature (FIX-1), rebuild the trainer, and repro at the exact wave recipe:
cd /home/sehaxe/dormouse-wt/situ-oom
CARGO_TARGET_DIR=/mnt/dormouse-target cargo build --release -p dormouse-train --features cuda
systemd-run --user --scope -p MemoryMax=40G \
  ./target/release/train --preset /tmp/opencode/presets/situ.toml \
  --data /mnt/e43497ab-0ff2-45b4-b45f-28de3339a53e/aria_data/pretrain/real \
  --eval /mnt/e43497ab-0ff2-45b4-b45f-28de3339a53e/aria_data/pretrain/real_eval/eval_tail.bin \
  --batch 8 --seq-len 512 --steps 400 --log /home/sehaxe/logs/situ_oom_repro.log
#    Falsifies the diagnosis if it still dies at 200-400 with 32 MiB reserve panics
#    (with the fused path the 32 MiB intermediates no longer exist).
#    Confirmation signal in the log: no `failed to reserve 33554432` after step 0.

# 3. Optional negative control (tensor path, proves the mechanism, ~10 min):
#    rebuild WITHOUT the feature flip, same command — expect the death to reproduce.

Cost of the un-fixed arm: the tensor path costs ~10-12 extra elementwise launches per call (the §7 estimate — correct as launches) AND ~4 GiB live memory under NoCheckpointing (the part §7 missed). At this recipe only the fused path fits a 16 GB card with the desktop up.

6. Follow-ups (other lanes’ files, named not touched)

1. Add a situ field to the eval line (crates/dormouse-train/src/lib.rs, the eval-line format at ~:1417): situ=<entry>/<fwd>/<bwd> from burn_situ::fused_situ::seam_counts(). Until it exists, tensor-vs-fused is invisible to any reader — this autopsy’s conclusion would have been a one-line grep.
2. Re-cost §7-style “the arm is cheap” claims under the checkpointing strategy actually configured. The wave’s checkpointing: Disabled is a trainer-level switch (dt-graph lane’s territory); every memory estimate made under BalancedCheckpointing assumptions is stale for it.
3. The A/B wave record: the situ arm’s numbers from this wave are void — not because the tensor path computes the wrong thing (it is the reference form), but because the arm cannot complete at this recipe. Re-run after FIX-1 + gate; 3 seeds, same window (81 920 B at batch 8).
4. situ-2026-09-30.md §7’s cost paragraph should carry a pointer here; the glossary’s disagreement list gains an entry (“§7 priced launches, the wave paid memory”).