feat(graph,arena): SaveForBackward contract with arena pinning (ADR 006, T2.1+T2.2)

[dndungu]

Implements ADR 006 decisions 1–3 (decision 4, poison mode, already landed): the save-for-backward contract in the graph plus the arena Pin/Unpin it rests on. Plan tasks T2.1, T2.2, S2.1.1, S2.2.1 of zerfoo docs/plan-gpu-training-hardening.md.

Refs #128.

Why

Nodes cache forward intermediates in struct fields and read them in Backward; the arena reuses/overwrites those buffers first (zerfoo#842 LayerNorm variance, zerfoo#845 gradient buffer, Wolf QK-norm cached inverse). Nothing in ztensor could express "this tensor must survive until backward". Now the graph owns that lifetime, PyTorch-style.

API chosen (T2.1)

• graph.Saver[T] — SaveForBackward(ts ...*tensor.TensorNumeric[T]), the handle a node calls during Forward.
• graph.SaverAware[T] — SetSaver(Saver[T]); Builder.Build wires a per-node Saver into every SaverAware node. A per-node handle (rather than a context.Context value or a "current node" field on the graph) keeps saves correctly attributed under ParallelForward, where several node Forwards run concurrently; the saved-set map carries its own mutex for the same reason. Nodes used outside a Graph get no Saver and must tolerate nil.
• (*Graph[T]).SaveForBackward(node, ts...) is the underlying recording method for callers that orchestrate nodes manually; SavedForBackward(node) exposes the set for tests/diagnostics.
• Recompute-from-live-inputs ... remains the documented alternative for cheap intermediates; both mechanisms (and the deprecation of raw struct-field caches) are documented in graph/save_for_backward.go.

Pinning plumbing: graph → tensor.PinnableStorage (implemented by GPUStorage, which pins the allocation base pointer with the original allocation size) → gpuapi.BackwardPinner (implemented by CUDAArenaPool) → cuda.ArenaPool.Pin/Unpin. CPU storages do not implement PinnableStorage, so CPU engines are completely unaffected. UnpinForBackward targets the pointer captured at pin time, so the pin stays balanced even when the executor's refcount release frees a saved intermediate mid-step.

Lifecycle: the graph unpins a node's saved set immediately after that node's Backward returns (success or error), releases all remaining sets at the end of Backward (nodes that never received a gradient), and at the start of the next Forward — so forward-only inference loops retain at most one pass of saved intermediates and never accumulate pins.

Reset-vs-pinned semantics chosen (T2.2)

Raise the floor.Reset() rewinds the bump offset to the end of the highest pinned span instead of the reset floor. This is the simplest provably-safe option: nothing at or below the rewound offset can ever be re-issued by the bump allocator, and the free-list is cleared as before, so no alias of a pinned span can exist.

• Watermark consequence (documented in arena_pin.go): all bytes between the reset floor and the highest pinned span — including dead, unpinned buffers below it — stay retained until the pins release; the next Reset after the last Unpin reclaims them. Under poison mode those retained-dead bytes are still NaN-filled at Reset (poisonUnpinnedSpanLocked skips only the pinned spans). Pinned bytes are bounded by what backward genuinely needs; PinnedBytes()/PinnedHighWaterBytes() (also surfaced as GPUEngine.ArenaPinnedBytes()) monitor the cost.
• FreeArena on a pinned span is deferred (no poison, no free-list entry) until the last Unpin, which applies the poison fill and free-list insert exactly as if the free had happened then. Reset drops deferred frees wholesale — a deferred block can extend above the raised floor, and releasing it later would alias the bump path.
• Pins are refcounted per buffer (map keyed by base-pointer offset, under the existing lock). Refcount underflow warns via a swappable sink and never panics; non-arena pointers are no-ops.

Migrated consumer

graph/checkpoint.gocheckpointNode: it cached the segment's forward inputs in seg.savedInputs and re-read them in Backward for recomputation — exactly the bug shape. It now implements SaverAware and registers the saved inputs via SaveForBackward; the struct field keeps only the Go handles, the contract owns the memory lifetime.

Regression-test story

TestSaveForBackward_WolfHazard_* (graph package) drives a real host-backed cuda.ArenaPool through the Wolf hazard schedule — Forward, arena.Reset() (per-sample ResetPool), Backward — with a LayerNorm-variance-shaped node:

• using SaveForBackward: Backward reads the correct value (2.5), pins drop to zero after Backward, and the next Reset reclaims+poisons the span;
• using a raw struct-field cache under poison mode: Backward reads NaN — proving the contract is what fixes the bug class and the poison mode is what exposes it.

Arena-level tests (S2.2.1) cover: pinned-survives-Reset (unpinned sibling poisoned), raised-floor retention of dead bytes, deferred FreeArena with poison-on-release, refcounted re-pin, underflow warning, non-arena no-op, PinnedBytes/high-water tracking, Reset-drops-deferred-frees, and zero fills when poison is off. Tensor-level tests cover the GPUStorage→pool delegation incl. unpin-after-free. All host-backed, no GPU needed in CI (poison-test pattern; hooks exported from internal/cuda only).

Gates

• go build ./... ✓
• CI vet package set ✓ (internal/cuda is excluded by CI; new files are vet-clean)
• gofmt ✓ on all new files (pre-existing formatting drift in graph/graph.go/checkpoint.go left untouched)
• go test -race -timeout 300s -count=1 ./... ✓

Open questions for the E2.3 migration wave

1. zerfoo-side op nodes (LayerNorm, attention, AdamW buffers) need the same migration: implement SaverAware and route every struct-field intermediate through the Saver. Audit list = every Backward reading a field written in Forward.
2. Saved tensors must not be reallocated (TrySet with a different length) between save and backward — acceptable contract, or should PinForBackward be re-taken on realloc?
3. MarkStepBoundary currently has no release hook; next-Forward + end-of-Backward release covers the known loops. If a caller needs pins dropped at an explicit step boundary without a Forward, we can add Graph.ReleaseSaved() trivially.