Arena ResetPool + memory reuse corrupts save-for-backward tensors → zero/garbage gradients on subsequent steps (GB10)

[dndungu]

Summary

On the GB10 (sm_121, CUDA 13), reusing arena memory after ResetPool() causes a subsequent forward+backward pass to compute a wrong gradient (collapses toward ~0). The save-for-backward (Saver-pinned) activations written by a step's forward get overwritten by that same step's later forward ops when they land on reused arena addresses. This is the same overwrite-on-reuse class as zerfoo#842 (partially worked around for specific nodes by recomputing from live inputs), but it bites the general training loop whenever ResetPool runs between steps.

Repro / evidence (Wolf CrossAsset, f32, batched B=256, qk-norm)

Wolf's batched training loop called ResetPool() once per mini-batch (to reclaim per-step scratch, issue-#118 style). Per-batch gradient L2 of the head weight, GPU vs CPU at identical seeded weights:

batch	CPU grad	GPU grad
0	4.315349	4.315349 (identical)
1	3.737152	2.5e-06
2	2.698082	1.98e-18
3	1.859067	0

• Batch 0 (no preceding ResetPool) is bit-identical CPU/GPU — forward, backward, and on-device AdamW step all correct.
• Every batch whose forward runs after a ResetPool produces a near-zero gradient, even though the forward LOSS is correct. So the forward is fine; the backward reads corrupted (overwritten) saved tensors.
• Disabling ResetPool entirely → GPU trains bit-identical to CPU (loss 0.747481, acc 0.6804). Confirms the corruption is the reset+reuse, not the kernels.
• It is NOT the optimizer state: re-marking the arena floor after the first step (so AdamW m/v sit below the floor) had zero effect; the gradient itself is wrong before opt.Step.

Expected

After ResetPool(), a fresh forward+backward must compute correct gradients. Reused arena memory must not overlap tensors still pinned by the Saver for the in-flight backward.

Likely fix area

Arena allocator / save-for-backward lifetime: either (a) the Saver should pin its tensors in non-resettable storage, or (b) the allocator must not hand out regions that alias live Saver-pinned tensors within a forward/backward pass.

Wolf interim

Hoisted ResetPool to once-per-epoch (feza-ai/wolf#218) — works for the capture-replay (production) path because replay uses a frozen recorded layout. The eager path (capture disabled) still diverges over epochs, which this upstream fix should resolve.

Platform: linux/arm64, NVIDIA GB10, CUDA 13 / sm_121. Models built from the existing Containerfile. Happy to provide the env-gated probe branch and full logs.

[dndungu]

Still reproduces on ztensor v1.19.0 (GB10, sm_121) — current evidence

Reconfirmed 2026-06-21 with a freshly rebuilt libkernels.so from v1.19.0 sources (sm_121, make -C internal/cuda/kernels shared CUDA_ARCH=sm_121, dropout kernel present). Wolf CrossAsset, identical pinned config (-batch-vectorize -batch-size=256 -precision=f32 -seed=42 -qk-norm, 40 epochs, 2 folds, v2-obv cache, 24613 samples), GPU capture-on (captures=1 replays=2559):

epoch0 loss=0.747  epoch1 0.561  epoch2 0.660  epoch3..39 PLATEAU ~0.669
fold-0 acc = 0.6047   (CPU baseline at identical config = 0.7257, loss → 0.0009)

Bit-for-bit the same trajectory as the 2026-06-18 Wolf T-PT.0 observation. Confirms:

• Not fixed in v1.18.0 → v1.19.0. The regression persists across the kernel-lib rebuild.
• Dropout-independent. This is at dropout=0 (no dropout nodes in the graph) — pure tri-path capture-replay. So it is squarely the arena/save-for-backward reuse bug, not anything dropout-related.
• Epoch 0 is correct (loss 0.747 == CPU), then replays collapse the gradient (loss bounces and plateaus at the ~0.669 'learned-a-little-then-stuck' level), matching the per-step grad-collapse table in the issue body.
• Wolf's per-epoch-reset workaround (T15.1, crossasset.go) is still in place and is not sufficient — the overwrite-on-reuse bites within the captured/replayed step graph.

Net: GPU capture-on training is a fast wrong answer on the current stack, which is why the production .zcam stays CPU-trained. Happy to run targeted diagnostics on the GB10 via Spark (e.g. a gradcheck-under-capture harness, or a bisect across kernel-source commits) if that helps localize the exact reuse site.

[dndungu]

FIXED + GB10-validated

Fix merged: zerfoo PR #910 raises the arena reset floor to the capture high-water (ArenaUsedBytes) after EndCapture in CaptureReplayRunner, reserving the captured graph's footprint for its replay lifetime (mirrors the decode path generate/generator.go onCaptured, which always did this).

GB10 acceptance (Wolf CrossAsset, dropout=0, byte-identical pinned config, cache reuse, capture-on):

epoch0 loss=0.747 ... epoch32 0.0052 ... epoch39 loss=0.000889   captures=1 replays=2559  ~12.5k samp/s
Fold 0: acc=0.7257   (== CPU baseline; was 0.6047 before the fix; loss had plateaued ~0.669)

GPU capture-on now trains bit-for-bit to the CPU baseline. Closing once released + propagated.