Skip to content

test(parity): Wolf-pattern attention stress test (zerfoo plan T2.4) - #141

Merged
dndungu merged 2 commits into
mainfrom
test/wolf-pattern-stress
Jun 12, 2026
Merged

test(parity): Wolf-pattern attention stress test (zerfoo plan T2.4)#141
dndungu merged 2 commits into
mainfrom
test/wolf-pattern-stress

Conversation

@dndungu

@dndungu dndungu commented Jun 12, 2026

Copy link
Copy Markdown
Contributor

Summary

Implements T2.4 of zerfoo docs/plan-gpu-training-hardening.md (verifies UC-GH-4, UC-GH-7): the Wolf-pattern integration stress test on an attention-shaped synthetic graph, building on the linear-net Wolf-pattern training loop from #140.

Graph shape (the increment over #140)

  • attentionNode: single-head self-attention Q=x@Wq, K=x@Wk, V=x@Wv; A=softmax(Q@K^T/sqrt(d)); out=A@V built as a graph.Node with trainable graph.Parameters. Forward registers Q, K, V, and the attention weights A via SaveForBackward (ADR 006). A is multi-consumer in backward: dV, dA, and the softmax denominator-style row reduction.
  • residualNormNode: residual + RMS normalization, saving the pre-norm activation and the normalization inverse (the Wolf QK-norm cached-inverse shape).
  • The input x feeds both nodes, so backward exercises the graph's converging-gradient accumulation path.

Schedule (the hazard that bit Wolf, sharpened)

Per sample: Graph.Forward → host-materialized loss gradient → arena reset BETWEEN forward and backward (every saved lifetime is load-bearing) → Graph.Backward with in-place gradient accumulation into persistent graph.Parameters → per-sample ResetPool. In-place SGD step per batch; 2 epochs × 2 batches × 3 samples; all under ZTENSOR_ARENA_POISON with a deliberately small arena.

Assertions

  • all parameters and accumulated gradients finite (poison NaN lands here),
  • GPU/StressEngine result within f32 tolerance (atol 1e-5 / rtol 1e-3) of the byte-identical plain-CPU run,
  • persistent parameter/gradient storage never re-homed (zerfoo#850/#855 dst contract),
  • zero leaked save-for-backward pins after training.

Red-proof

Removing the SaveForBackward calls makes TestAttentionTraining_WolfPattern_StressCI fail deterministically with Wq.Gradient[0] = NaN: non-finite under poison (stale arena read) — the saved lifetimes are load-bearing, the test has teeth.

CI / GB10

  • TestAttentionTraining_WolfPattern_StressCI runs in plain CI (host-backed-arena StressEngine, poison on).
  • TestAttentionTraining_WolfPattern_GPU skips without CUDA; scripts/parity/run.sh gains a hard PASS gate (exit 7) so the Spark GB10 pod fails unless it passes on real hardware. GB10 evidence will be added to the devlog in this PR before merge.

Quality gates: go build ./..., CI-scoped go vet, gofmt clean, go test -race ./... full suite green locally.

dndungu added 2 commits June 11, 2026 22:24
Adds TestAttentionTraining_WolfPattern_GPU / _StressCI: the Wolf gr-12
hazard schedule on an attention-shaped synthetic graph (Q@K^T -> scale ->
softmax -> @v plus residual + RMS-norm), built on the graph/node APIs with
SaveForBackward (ADR 006) for the multi-consumer saved intermediates
(Q/K/V/attention weights, pre-norm activation, normalization inverse).

Schedule: per-sample Forward+Backward with an arena reset BETWEEN forward
and backward (saved lifetimes load-bearing) plus the per-sample ResetPool,
in-place gradient accumulation into persistent graph.Parameters, in-place
SGD per batch, 2 epochs x 2 batches x 3 samples, all under
ZTENSOR_ARENA_POISON. Asserts params/grads finite, GPU within f32
tolerance of the byte-identical CPU run, persistent storage never
re-homed, and zero leaked pins.

Red-proof: removing the SaveForBackward calls makes the CI variant fail
with poison NaN at the first gradient check.

scripts/parity/run.sh gains the corresponding GB10 PASS gate (exit 7).
@dndungu
dndungu merged commit ea0fdfb into main Jun 12, 2026
1 check passed
@dndungu
dndungu deleted the test/wolf-pattern-stress branch June 12, 2026 17:37
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant