feat(layers): migrate cached-for-backward intermediates to SaveForBackward / live-input recompute (ADR 006, T2.3)

[dndungu]

Implements plan task T2.3 of docs/plan-gpu-training-hardening.md: a field-based audit of every graph.Node / layer Backward in zerfoo for struct fields written during Forward and read during Backward, plus migration of the training-critical set to the ztensor ADR 006 contract (graph.SaverAware / SaveForBackward, landed in zerfoo/ztensor#132 @ 1cc7793) or to live-inputs ... recompute.

Refs #847, zerfoo/ztensor#128.

Why

Nodes caching forward intermediates in struct fields and reading them in Backward are corrupted by GPU arena reuse before Backward runs (#842 LayerNorm variance, #845, the Wolf QK-norm cached inverse). The audit is AST-based (receiver fields assigned in Forward ∩ read in Backward, following same-receiver helper calls), not comment-based.

Audit method

A go/ast walker over every non-test .go file: for each type with pointer-receiver ForwardandBackward, intersect {fields assigned in Forward} with {fields read in Backward} (transitively through same-receiver method calls). 45 types flagged; every one is classified below.

Audit table

Classification: SAFE = host Go value, no device/arena storage. RECOMPUTE = derived from the live inputs ... Backward receives (preferred for cheap intermediates; the #842 LayerNorm pattern). SAVE = expensive to recompute; registered via SaveForBackward (node keeps the Go handle, the contract owns the memory lifetime).

Migrated in this PR (training-critical path: core, activations, normalization, regularization, attention, transpose, reducesum)

Type	Cached field(s)	Class	Resolution
activations.BaseActivation (ReLU/Tanh/Erf/…)	lastInput	RECOMPUTE	cache removed; Backward derives from inputs[0]
activations.Gelu	lastInput	RECOMPUTE	cache removed; Backward derives from inputs[0]; OutputShape from host []int
activations.LeakyReLU	lastInput	RECOMPUTE	cache removed; Backward derives from inputs[0]
activations.Sigmoid	lastOutput (lastInput was dead)	SAVE	SaverAware; output registered each Forward; dead lastInput removed
activations.Softmax	output	SAVE	SaverAware; softmax output registered each Forward
activations.SwiGLU	lastInput; gate, siluX1	RECOMPUTE + SAVE	re-splits live inputs[0]; gate/silu registered; Saver fans into composed Sigmoid
normalization.LayerNormalization	mean, variance, normedInput, inputShape	RECOMPUTE	completes #842: the generic engine-op path (f64/non-mixed) now recomputes stats from inputs[0] like the mixed f32 path; caches removed
normalization.RMSNorm	inputTensor, rms	RECOMPUTE	new rmsRecomputeStats helper; also fixes fused-GPU forward leaving rsqrt nil
normalization.SimplifiedLayerNormalization	invStdDev, normalizedInput	RECOMPUTE	same helper; backward-before-forward now legal
normalization.BatchNormalization	normalized, std, scaleBC	SAVE	Backward only receives X (not scale/mean/var) so stats are not recomputable from its live inputs; SaverAware + register all three
regularization.Dropout	mask	SAVE	random ⇒ not recomputable; SaverAware + register (training mode only)
regularization.FeatureDropout	mask	SAVE	same
attention.ScaledDotProductAttention	q, k, v, attentionWeights	SAVE	not itself a graph.Node (owners call Backward with nil q/k/v); gains SetSaver; owners fan their Saver in so saves attribute to the owning node
attention.AttentionHead	— (composite)	SAVE (fan-out)	SaverAware; fans into SDPA
attention.GroupedQueryAttention	attnOutputFinal	SAVE	SaverAware; registers attnOutputFinal (read by wo backward) + fans into SDPA
attention.GlobalAttention, attention.LocalAttention	— (wrap GQA)	SAVE (fan-out)	SaverAware; fan into GQA
attention.MultiHeadLatentAttention	cachedCKV, cachedAttnOut	SAVE	SaverAware; registers both + fans into SDPA
attention.FusedSDPA	hadMask	SAFE	host bool; no change
core.FFN	inputTensor; w1Output, w3Output, swiGLUOutput	RECOMPUTE + SAVE	input from live inputs[0] (signature previously ignored them); matmul outputs registered; Saver fans into composed SwiGLU
core.FiLM	lastFeature, lastScale (lastBias dead)	RECOMPUTE + SAVE	feature from inputs[0]; generated scale registered; dead caches removed
core.PolynomialExpansion	lastInput	RECOMPUTE	cache removed; previously preferred the stale cache over live inputs — order inverted and cache deleted
core.TemporalConvEncoder	lastConv1Out, lastConv2Out, lastPooled	SAVE	SaverAware + register
core.VariableSelection	lastHidden, lastWeights (lastInput dead)	SAVE	SaverAware + register; dead cache removed
core.RotaryEmbedding	inner	SAFE	lazily-built Go struct handle (its internal cos/sin tables are the embeddings FOLLOW-UP below)
transpose.Transpose	perm	SAFE	host []int lazily defaulted in Forward
reducesum.ReduceSum	—	SAFE	no Forward-written/Backward-read fields (already live-input based)
functional.GELUBackward	(wrapper)	RECOMPUTE	passed the live input to Gelu.Backward instead of forward-seeding the removed cache (also removes a wasted forward per call)
core.Linear/Bias/Add/MatMul/Dense/Gemm	—	SAFE	no cached intermediates; Backward already uses live inputs
timeseries engines (timeseries/*, layers/timeseries/*)	—	SAFE	backward helpers are pure functions of explicit arguments; audit flagged no Forward-written/Backward-read fields

FOLLOW-UP (classified, not migrated here; tracked under #847)

Type	Cached field(s)	Class
core.MatMulNBits	dequantizedWeights, cacheValid	SAVE* — persistent cross-step cache; per-backward pinning does not cover it (contract gap, see below)
core.MixtureOfExperts	cachedAssignments/ExpertOuts/GateWeight/HiddenStates/Indices/Weights	SAVE
core.MoEGate	cachedGateWeight/HiddenStates/Indices/Probs (cachedNumExperts SAFE int)	SAVE
embeddings.RotaryPositionalEmbedding	cosAngles, sinAngles	RECOMPUTE (deterministic tables) or SAVE
embeddings.TokenEmbedding	inputTokenIDs	RECOMPUTE (live inputs)
hrm.HModule, hrm.LModule	fwdCombinedInput (shape/bool fields SAFE)	SAVE
recurrent.SimpleRNN	lastInput (RECOMPUTE); hiddenState	cross-step recurrent state — same contract gap as MatMulNBits
ssm.BCNorm	cachedInput, cachedNorm	RECOMPUTE + SAVE
ssm.MambaBlock / ssm.MIMOMambaBlock / ssm.ComplexSSMState	13–14 cached tensors each	SAVE (large set)
transformer.Block	fwdInput (RECOMPUTE); fwdNorm1Out/fwdNorm2Out/fwdPostAttn/fwdResidual1	SAVE
training/fp8.FP8Linear	lastInput	RECOMPUTE
training/lora.LoraLinear	lastAx	SAVE
training/loss.MSE/BCELoss/CorrLoss/RoutingContrastive	predictions, targets, scores	RECOMPUTE (live inputs)
training/loss.CrossEntropyLoss	predictions, targets (RECOMPUTE); softmaxOutput	SAVE
training/nas.DARTSLayer	lastInput (RECOMPUTE); opOutputs, weights	SAVE

Counts: 45 types flagged ⇒ 27 resolved in this PR (11 RECOMPUTE migrations, 15 SAVE registrations/fan-outs, plus SAFE documentation), 18 FOLLOW-UP types classified above. Re-running the audit on the migrated tree shows the prioritized packages contain only SAVE-registered or SAFE fields.

Wiring

ztensor Builder.Build (1cc7793) calls SetSaver on every SaverAware node, so anything constructed through the graph builder (zerfoo inference/arch_*, Wolf's CrossAsset graph) is wired automatically. Composite layers fan the Saver into children (SwiGLU→Sigmoid, FFN→SwiGLU, AttentionHead/GQA/MLA→SDPA, Global/Local→GQA) so child saves attribute to the parent node and release when the parent's Backward returns. Nodes used outside a Graph tolerate a nil Saver (no-op), keeping CPU behavior identical.

Regression tests

• SAVE:TestDropout_SaveForBackward_RegistersMask / TestFeatureDropout_… (mask registered in training mode, nothing in eval), TestSoftmax_SaveForBackward_RegistersOutput, TestSwiGLU_SaveForBackward_RegistersGateAndSilu — stub graph.Saver recording exactly which tensors are registered.
• RECOMPUTE:TestGelu_Backward_ReadsLiveInputs, TestLayerNorm_MixedBackward_ReadsLiveInput (f32), TestLayerNorm_GenericBackward_ReadsLiveInput (f64, the path newly migrated here) — the input buffer is overwritten in place between Forward and Backward (exactly what arena reuse does) and gradients must match a fresh layer that only ever saw the new values.
• Obsolete cache-dependency tests updated: backward-before-forward now legal for recompute layers (RMSNorm, SimplifiedLN).

Gates

• go build ./... ✓
• go vet ./... ✓
• gofmt ✓ on all touched files (pre-existing drift in untouched files left alone)
• go test -race -count=1 ✓ on all changed packages (layers/activations, layers/normalization, layers/regularization, layers/attention, layers/core, layers/functional)
• go test ./... -count=1 full suite ✓ except cmd/benchTestBenchHarness, which is pre-existing flaky (fails identically with -count=5 on an untouched main checkout; timing-based ThroughputTs > 0 assertion)

Contract gaps found (feeding ztensor#132 open questions)

1. Cross-step caches (MatMulNBits.dequantizedWeights, SimpleRNN.hiddenState): the contract releases pins when the node's Backward returns, but these caches must survive across steps until invalidated. Needs either a pin-until-invalidate API or a guarantee that such caches are allocated outside the arena. This is related to ztensor#132 open question 2 (realloc between save and backward).
2. MarkStepBoundary release hook (ztensor#132 open question 3): not needed by any layer migrated here — next-Forward + end-of-Backward release covers every zerfoo training/inference loop encountered.
3. Direct-Backward callers: several zerfoo wrappers called Backward without passing inputs, relying on caches (functional.GELUBackward, FFN tests). The graph always passes live inputs, but a lint (vet check or contract doc) that Backward implementations must not depend on Forward-time tensor state would prevent regressions; the gradcheck-under-poison harness (T1.2/S2.3.1) is the runtime backstop.

Not done here

• DGX/GPU poison-mode run (S2.3.1) — separate Spark task per plan.
• FOLLOW-UP migrations listed above (ssm/hrm/transformer/loss/lora/nas/MoE/embeddings).