test(parity): Wolf-pattern training loop under poison -- GB10 + CI variants (S2.3.1)

docs/devlog.md

@@ -1,5 +1,41 @@
 # ztensor Development Log
 
+## 2026-06-11: S2.3.1 -- poison-mode full-suite run green on GB10 (parity + Wolf-pattern training loop)
+
+**Type:** validation
+**Tags:** gb10, spark, poison, arena, parity, save-for-backward, S2.3.1, #130, #133, #140
+
+**What ran:** the framework-level proof for zerfoo plan S2.3.1 (UC-GH-004), as a
+single Spark pod on the DGX GB10 -- pod `ztensor-parity-4b4d759c`, image
+`docker.io/library/golang:1.26-bookworm` + hostPath CUDA/kernels, ref
+`test/s231-poison-gb10` (4b4d759c, = v1.11.0/ebf0008b + the new training-loop
+test, PR #140), `ZTENSOR_ARENA_POISON=1`, 64 MiB test arena
+(`SetArenaBytesForTesting`), memory limit 16Gi, GPU serialized (no other pods).
+
+**Results (all PASS, pod exit 0, `VALIDATION_OK`):**
+- `TestParity_GPUvsCPU_ArenaStressSchedules_GPU`: full gradcheck registry,
+  CPU-f32 vs GPU-f32, forward+backward. `no-reset`: 26/26 passed, 0 failed,
+  0 errored. `reset-between-fwd-bwd` (the Wolf per-sample-ResetPool hazard,
+  poison on): 26/26 passed, 0 failed, 0 errored. Per-schedule JSON reports
+  archived from `/home/ndungu/parity/4b4d759c/` (max_abs/max_rel per op all
+  within per-op tolerances; worst observed ~7e-06 on Tanh gradients).
+- `TestParity_GPURedProof_GPU`: raw cached-intermediate fixture flagged red on
+  the real CUDA arena; contract twin green -- the run retains its sensitivity.
+- `TestTrainingLoop_WolfPattern_GPU` (new, PR #140): two-layer net trained
+  3 batches x 4 samples with the exact gr-12 hazard schedule -- per-sample
+  forward+backward, in-place gradient accumulation into persistent non-arena
+  GPUStorage, ResetPool every sample, in-place SGD step per batch -- under
+  poison. All parameters finite and within f32 tolerance (atol 1e-5 /
+  rtol 1e-3) of the byte-identical CPU-engine run; engine Add/Sub preserved
+  the persistent dst storage (no re-homing, the #850/#855 failure mode).
+
+**Meaning:** with v1.11.0 (host-access stream-ordering #137, reset-epoch frees
+#138, dst preservation #134, SaveForBackward/pinning #132, poison #130) the
+lifetime contract holds on real GB10 hardware under the schedule shapes that
+produced the zerfoo#842/#845/#850 corruption class. M2's GB10 leg
+(poison-mode full suite) is done; T2.4 (attention-like integration stress)
+remains the broader follow-up.
+
 ## 2026-06-11: #137 + #138 -- host-access stream ordering and arena reset-epoch frees validated end-to-end on GB10
 
 **Type:** finding

scripts/parity/README.md

@@ -106,3 +106,16 @@ curl -s -X DELETE $SPARK/api/v1/pods/ztensor-parity-$RUNID
 `run.sh` exits non-zero unless BOTH GPU tests genuinely passed (a
 CUDA-unavailable SKIP is a hard failure), so the pod phase is the verdict.
 Test a non-main ref by setting `ZTENSOR_REF` in the manifest env.
+
+## Wolf-pattern training loop (S2.3.1)
+
+`TestTrainingLoop_WolfPattern_GPU` adds the third GB10 gate: a small
+two-layer training loop running the exact Wolf gr-12 hazard schedule
+(per-sample forward+backward, gradients accumulated in place into
+persistent non-arena device buffers, `ResetPool` after every sample,
+in-place SGD step once per batch) under poison with the small arena.
+PASS requires finite parameters AND f32-tolerance agreement with a plain
+CPU-engine run of the byte-identical loop. A CI twin
+(`TestTrainingLoop_WolfPattern_StressCI`) runs the same loop against the
+host-backed-arena StressEngine, so lifetime regressions in the pattern
+are caught before a GB10 round trip.

scripts/parity/run.sh

@@ -35,9 +35,10 @@ grep -- 'parity ' /tmp/out.txt || true
 
 grep -q -- '--- PASS: TestParity_GPUvsCPU_ArenaStressSchedules_GPU' /tmp/out.txt || { echo FATAL: schedule parity not PASS; exit 3; }
 grep -q -- '--- PASS: TestParity_GPURedProof_GPU' /tmp/out.txt || { echo FATAL: GPU red-proof not PASS; exit 5; }
+grep -q -- '--- PASS: TestTrainingLoop_WolfPattern_GPU' /tmp/out.txt || { echo FATAL: Wolf-pattern training loop not PASS; exit 6; }
 grep -q -- '--- SKIP: TestParity_GPUvsCPU_ArenaStressSchedules_GPU' /tmp/out.txt && { echo FATAL: SKIPPED no CUDA; exit 4; }
 test "$code" -eq 0 || { echo "FATAL: go test exit $code"; exit "$code"; }
 
 echo "reports:"
 ls -l "$REPORT_DIR"
-echo "VALIDATION_OK: CPU-vs-GPU parity (both schedules) + GPU red-proof passed on GB10"
+echo "VALIDATION_OK: CPU-vs-GPU parity (both schedules) + GPU red-proof + Wolf-pattern training loop passed on GB10"

testing/parity/training_gpu_test.go

@@ -0,0 +1,286 @@
+package parity
+
+import (
+	"context"
+	"math"
+	"testing"
+
+	"github.com/zerfoo/ztensor/compute"
+	"github.com/zerfoo/ztensor/internal/cuda"
+	"github.com/zerfoo/ztensor/numeric"
+	"github.com/zerfoo/ztensor/tensor"
+)
+
+// TestTrainingLoop_WolfPattern_GPU is the GB10 training-loop proof for
+// poison-mode hardening (zerfoo plan S2.3.1): a small two-layer training
+// loop running the EXACT hazard schedule that bit Wolf (gr-12 /
+// zerfoo#850): per-sample forward+backward with the gradients accumulated
+// into PERSISTENT (non-arena) device buffers, engine ResetPool after every
+// sample, and an in-place optimizer step once per batch -- all on the real
+// CUDA arena with ZTENSOR_ARENA_POISON semantics enabled and a deliberately
+// small arena. Any read of recycled arena memory (a stale activation, a
+// gradient left arena-backed, an accumulator silently re-homed into the
+// pool) surfaces as a deterministic NaN and fails the run.
+//
+// PASS requires (1) every parameter finite after training and (2) the GPU
+// result matching a plain CPU-engine run of the byte-identical loop within
+// f32 reduction tolerance -- catching silent corruption, not just NaN.
+//
+// Runs on the DGX via the Spark pod in scripts/parity/; skips without CUDA.
+func TestTrainingLoop_WolfPattern_GPU(t *testing.T) {
+	if !cuda.Available() {
+		t.Skip("CUDA not available")
+	}
+	restorePoison := cuda.SetArenaPoisonEnabledForTesting(true)
+	defer restorePoison()
+	restoreArena := compute.SetArenaBytesForTesting(gpuParityArenaBytes)
+	defer restoreArena()
+
+	gpuEng, err := compute.NewGPUEngine[float32](numeric.Float32Ops{})
+	if err != nil {
+		t.Fatalf("NewGPUEngine: %v", err)
+	}
+	defer func() { _ = gpuEng.Close() }()
+
+	gpuW1, gpuW2 := runWolfTrainingLoop(t, trainSide{
+		eng:   gpuEng,
+		reset: gpuEng.ResetPool,
+		persistent: func(shape []int, data []float32) *t32 {
+			st, err := tensor.NewGPUStorageFromSlice(append([]float32(nil), data...))
+			if err != nil {
+				t.Fatalf("NewGPUStorageFromSlice: %v", err)
+			}
+			tt, err := tensor.NewWithStorage[float32](shape, st)
+			if err != nil {
+				t.Fatalf("NewWithStorage: %v", err)
+			}
+			return tt
+		},
+	})
+
+	cpuW1, cpuW2 := runWolfTrainingLoop(t, cpuTrainSide(t))
+
+	assertFiniteAndClose(t, "W1", gpuW1, cpuW1)
+	assertFiniteAndClose(t, "W2", gpuW2, cpuW2)
+}
+
+// TestTrainingLoop_WolfPattern_StressCI runs the identical loop in ordinary
+// CI with the host-backed-arena StressEngine as the candidate (GPU lifetime
+// semantics without a GPU) under poison, against the plain CPU reference --
+// catching lifetime regressions in the loop pattern before a GB10 run.
+func TestTrainingLoop_WolfPattern_StressCI(t *testing.T) {
+	enablePoison(t)
+	stress := NewStressEngine(compute.NewCPUEngine[float32](numeric.Float32Ops{}), 1<<20)
+	sW1, sW2 := runWolfTrainingLoop(t, trainSide{
+		eng:        stress,
+		reset:      stress.ResetArena,
+		persistent: hostPersistent(t),
+	})
+
+	cpuW1, cpuW2 := runWolfTrainingLoop(t, cpuTrainSide(t))
+
+	assertFiniteAndClose(t, "W1", sW1, cpuW1)
+	assertFiniteAndClose(t, "W2", sW2, cpuW2)
+}
+
+// cpuTrainSide is the plain CPU reference side: no arena, no reset.
+func cpuTrainSide(t *testing.T) trainSide {
+	t.Helper()
+	return trainSide{
+		eng:        compute.NewCPUEngine[float32](numeric.Float32Ops{}),
+		persistent: hostPersistent(t),
+	}
+}
+
+// hostPersistent allocates GC-owned host tensors, which can never be
+// recycled behind a live reference.
+func hostPersistent(t *testing.T) func(shape []int, data []float32) *t32 {
+	t.Helper()
+	return func(shape []int, data []float32) *t32 {
+		tt, err := tensor.New[float32](shape, append([]float32(nil), data...))
+		if err != nil {
+			t.Fatalf("tensor.New: %v", err)
+		}
+		return tt
+	}
+}
+
+// trainSide abstracts the engine-specific pieces of the loop: how persistent
+// (never-recycled) parameter/accumulator tensors are allocated and how the
+// per-sample arena reset is performed (nil for plain CPU).
+type trainSide struct {
+	eng        compute.Engine[float32]
+	reset      func()
+	persistent func(shape []int, data []float32) *t32
+}
+
+// runWolfTrainingLoop trains a tiny x->MatMul->Tanh->MatMul->Softmax net for
+// several batches with per-sample ResetPool and per-batch SGD, returning the
+// final parameter values (host copies).
+func runWolfTrainingLoop(t *testing.T, side trainSide) ([]float32, []float32) {
+	t.Helper()
+	ctx := context.Background()
+	eng := side.eng
+
+	const (
+		inDim   = 8
+		hidDim  = 8
+		outDim  = 4
+		batches = 3
+		samples = 4
+		lr      = float32(0.05)
+	)
+
+	// Deterministic, engine-independent initialization and data.
+	w1Init := rampData(inDim*hidDim, 0.31)
+	w2Init := rampData(hidDim*outDim, 0.17)
+
+	// Persistent parameters and gradient accumulators: non-arena storage
+	// (raw GPUStorage on the GPU side), so ResetPool never recycles them.
+	w1 := side.persistent([]int{inDim, hidDim}, w1Init)
+	w2 := side.persistent([]int{hidDim, outDim}, w2Init)
+	g1 := side.persistent([]int{inDim, hidDim}, make([]float32, inDim*hidDim))
+	g2 := side.persistent([]int{hidDim, outDim}, make([]float32, hidDim*outDim))
+
+	for b := 0; b < batches; b++ {
+		for s := 0; s < samples; s++ {
+			x := mustNew(t, []int{1, inDim}, rampData(inDim, 0.41+float32(b*samples+s)*0.07))
+			target := oneHot(t, outDim, (b+s)%outDim)
+
+			// Forward (every intermediate is arena-backed on the GPU side).
+			h := mustOp(t, "MatMul(x,W1)")(eng.MatMul(ctx, x, w1))
+			a := mustOp(t, "Tanh")(eng.Tanh(ctx, h))
+			y := mustOp(t, "MatMul(a,W2)")(eng.MatMul(ctx, a, w2))
+			p := mustOp(t, "Softmax")(eng.Softmax(ctx, y, 1))
+
+			// Backward (softmax+cross-entropy: dy = p - target).
+			dy := mustOp(t, "Sub(p,target)")(eng.Sub(ctx, p, target))
+			aT := mustOp(t, "Transpose(a)")(eng.Transpose(ctx, a, []int{1, 0}))
+			dW2 := mustOp(t, "MatMul(aT,dy)")(eng.MatMul(ctx, aT, dy))
+			w2T := mustOp(t, "Transpose(W2)")(eng.Transpose(ctx, w2, []int{1, 0}))
+			da := mustOp(t, "MatMul(dy,W2T)")(eng.MatMul(ctx, dy, w2T))
+			dh := mustOp(t, "TanhPrime")(eng.TanhPrime(ctx, h, da))
+			xT := mustOp(t, "Transpose(x)")(eng.Transpose(ctx, x, []int{1, 0}))
+			dW1 := mustOp(t, "MatMul(xT,dh)")(eng.MatMul(ctx, xT, dh))
+
+			// Accumulate into the persistent buffers IN PLACE (the
+			// zerfoo#855 contract: dst's storage must be preserved).
+			accumulateInPlace(t, ctx, eng, g1, dW1, "G1")
+			accumulateInPlace(t, ctx, eng, g2, dW2, "G2")
+
+			// Wolf's per-sample ResetPool: recycle (and poison) every
+			// arena intermediate of this sample.
+			if side.reset != nil {
+				side.reset()
+			}
+		}
+
+		// Optimizer step, once per batch: W -= lr * G, in place.
+		sgdStepInPlace(t, ctx, eng, w1, g1, lr, "W1")
+		sgdStepInPlace(t, ctx, eng, w2, g2, lr, "W2")
+		if err := eng.Zero(ctx, g1); err != nil {
+			t.Fatalf("Zero(G1): %v", err)
+		}
+		if err := eng.Zero(ctx, g2); err != nil {
+			t.Fatalf("Zero(G2): %v", err)
+		}
+		// Recycle the step's arena temporaries before the next batch.
+		if side.reset != nil {
+			side.reset()
+		}
+	}
+
+	return append([]float32(nil), w1.Data()...), append([]float32(nil), w2.Data()...)
+}
+
+// accumulateInPlace performs accum += grad via the engine and asserts the
+// engine honored the dst contract (no re-homing of the persistent buffer
+// into the arena -- the exact failure mode of zerfoo#850/#855).
+func accumulateInPlace(t *testing.T, ctx context.Context, eng compute.Engine[float32], accum, grad *t32, name string) {
+	t.Helper()
+	before := accum.GetStorage()
+	res, err := eng.Add(ctx, accum, grad, accum)
+	if err != nil {
+		t.Fatalf("Add into %s: %v", name, err)
+	}
+	if res != accum || res.GetStorage() != before {
+		t.Fatalf("engine.Add relocated persistent accumulator %s (storage %p -> %p); dst must be written in place", name, before, res.GetStorage())
+	}
+}
+
+// sgdStepInPlace performs w -= lr*g with the scaled gradient as an arena
+// temporary and the subtraction writing in place into the persistent w.
+func sgdStepInPlace(t *testing.T, ctx context.Context, eng compute.Engine[float32], w, g *t32, lr float32, name string) {
+	t.Helper()
+	step, err := eng.MulScalar(ctx, g, lr)
+	if err != nil {
+		t.Fatalf("MulScalar(%s): %v", name, err)
+	}
+	before := w.GetStorage()
+	res, err := eng.Sub(ctx, w, step, w)
+	if err != nil {
+		t.Fatalf("Sub into %s: %v", name, err)
+	}
+	if res != w || res.GetStorage() != before {
+		t.Fatalf("engine.Sub relocated persistent parameter %s; dst must be written in place", name)
+	}
+}
+
+// assertFiniteAndClose checks every GPU value is finite (poison NaN would
+// land here) and within f32 reduction tolerance of the CPU reference.
+func assertFiniteAndClose(t *testing.T, name string, gpu, cpu []float32) {
+	t.Helper()
+	if len(gpu) != len(cpu) {
+		t.Fatalf("%s: length mismatch gpu=%d cpu=%d", name, len(gpu), len(cpu))
+	}
+	const (
+		atol = 1e-5
+		rtol = 1e-3
+	)
+	for i := range gpu {
+		gv, cv := float64(gpu[i]), float64(cpu[i])
+		if math.IsNaN(gv) || math.IsInf(gv, 0) {
+			t.Fatalf("%s[%d] = %v: non-finite after training under poison (stale arena read)", name, i, gv)
+		}
+		if diff := math.Abs(gv - cv); diff > atol+rtol*math.Abs(cv) {
+			t.Errorf("%s[%d]: gpu=%v cpu=%v |diff|=%v exceeds atol %v + rtol %v", name, i, gv, cv, diff, atol, rtol)
+		}
+	}
+	t.Logf("%s: %d values finite, GPU within atol %v / rtol %v of CPU", name, len(gpu), atol, rtol)
+}
+
+// rampData generates deterministic, well-conditioned values in (-1, 1).
+func rampData(n int, seed float32) []float32 {
+	out := make([]float32, n)
+	for i := range out {
+		out[i] = float32(math.Sin(float64(seed) + 0.7*float64(i)))
+	}
+	return out
+}
+
+func oneHot(t *testing.T, n, hot int) *t32 {
+	t.Helper()
+	data := make([]float32, n)
+	data[hot] = 1
+	return mustNew(t, []int{1, n}, data)
+}
+
+func mustNew(t *testing.T, shape []int, data []float32) *t32 {
+	t.Helper()
+	tt, err := tensor.New[float32](shape, data)
+	if err != nil {
+		t.Fatalf("tensor.New: %v", err)
+	}
+	return tt
+}
+
+// mustOp adapts the (tensor, error) engine-op return for inline use.
+func mustOp(t *testing.T, what string) func(*t32, error) *t32 {
+	t.Helper()
+	return func(res *t32, err error) *t32 {
+		if err != nil {
+			t.Fatalf("%s: %v", what, err)
+		}
+		return res
+	}
+}