v1.9.0: arena overflow hit on FIRST per-sample alloc (64GB arena) + async overflow still stalls GB10 — crossasset training freezes step 0

[dndungu]

Summary

Follow-up to #115. With v1.9.0 (stream-ordered overflow) the crossasset GB10
training still does not complete one step. Two distinct problems:

1. The arena overflow is hit on the FIRST mini-batch sample's tiny alloc, even
   with ZERFOO_ARENA_SIZE_GB=64. A 64 GB arena should not be exhausted by one
   per-sample forward pass — this looks like an arena sizing/reset/accounting issue.
2. The async overflow (cudaMallocAsync) itself still stalls on GB10 under
   memory pressure — so v1.9.0's fix changed sync→async but the overflow path still
   wedges.

Net: training freezes in the first ComputeGradients (step=0).

Pinned goroutine 1 (v1.9.0, ZERFOO_ARENA_SIZE_GB=64, GB10)

goroutine 1 [syscall]:
runtime.cgocall -> cuda._Cfunc_ccall_wrapper
 -> cuda.MallocAsync(0xc000=48KB)        runtime_purego.go:92   # async overflow
 -> cuda.(*ArenaPool).Alloc              arena.go:186           # overflow path (v1.9.0)
 -> gpuapi.(*CUDAArenaPool).Alloc        cuda_arena.go:59
 -> compute.gpuUnaryOp                   gpu_kernels.go:603
 -> compute.(*GPUEngine).gpuTanh
 -> ... forward pass, first sample

D-state sibling thread: folio_wait_bit_common -> __folio_lock_or_retry (page
fault during the async alloc).

Why (1) is the key puzzle

• Trainer mini-batches: batch=256, but processes one sample per
  ComputeGradients (gradient accumulation), so the per-step working set is a
  single sample's activations — small.
• Engine is CUDAArenaPool (not the MemPool fallback), and NewGPUEngine honors
  ZERFOO_ARENA_SIZE_GB (gpu_engine.go:225). No "arena pool not available" warning.
  So a 64 GB arena was created.
• Yet ArenaPool.Alloc reaches the overflow branch for a 48 KB gpuTanh
  output on the first sample. So either the arena's usable capacity is far below
  64 GB on this path, or it is not being reset (StepScope/MarkStepBoundary) and a
  single forward+backward genuinely fills 64 GB (unlikely for one sample), or there
  is an offset/accounting bug.

Asks

1. Add arena diagnostics (log at engine init + at first overflow): configured
   capacity, current offset, hits/misses/reuses, alloc count, and whether
   Reset/MarkStepBoundary is being driven by the training loop. This will say
   immediately whether the arena is mis-sized vs. legitimately full vs. never reset.
2. Investigate why a single per-sample forward fills a 64 GB arena (or doesn't,
   and the overflow is an accounting bug).
3. Harden the async overflow on GB10 — cudaMallocAsync under unified-memory
   pressure still stalls here; ideally overflow is never reached (proper
   sizing/reset), but the path should not wedge if it is.

Env / repro

NVIDIA GB10 (sm_121), ztensor v1.9.0. Wolf train-crossasset -gpu, full COIN
1m bars, folds=2 epochs=1, batch=256, ZERFOO_ARENA_SIZE_GB=64, mem 96Gi. Hangs in
the first step's forward pass. Capture rig + dumps:
wolf/scripts/t82-stall-watchdog.sh, wolf/.claude/scratch/t82-dump/.

Cross-refs

#115 (async overflow — necessary, still insufficient), #111, #106. Wolf devlog
2026-06-07 has the full chain across v1.8.1 -> v1.8.2 -> v1.9.0.

[dndungu]

Shipped in v1.10.0 (PR #119) — the two ztensor-side pieces of this issue:

1. Arena diagnostics (ask 1): engine init now logs the resolved arena config,
   and ArenaPool emits a one-shot per-reset-epoch first-overflow log:
   path / requested / aligned / capacity / offset / epochAllocs / epochMaxAlloc / hits / misses / reuses / resets / freeListLen. Per-epoch counters are cleared
   by Reset, so the line describes a single forward+backward pass.
2. Async-overflow hardening (ask 3): SetMemPoolReleaseThreshold sets
   cudaMemPoolAttrReleaseThreshold on the device stream-ordered pool (default =
   arena capacity, override ZERFOO_OVERFLOW_POOL_RETAIN_GB), so freed
   cudaMallocAsync blocks stay resident instead of being returned to the OS and
   faulted back in (folio_wait_bit).

Code review ruled out an arena-sizing parse bug: arenaSizeBytes(64) = 68 719 476 736
as int64, the int() cast is 64-bit on arm64, and a failed 64 GB Malloc would log
"arena pool not available" (which the stack contradicts). A real 64 GB arena exists.

GB10 validation (Spark ztensor-issue118-diag-1): diagnostics fire on hardware and
the hardened async path ran 200 overflow rounds with no wedge.

Why this stays open — the remaining gate (problem 1). The synthetic GB10 test
uses a 1 MiB arena and tiny allocs; it does not reproduce the real ~64 GB pressure,
so it can't classify why crossasset's first sample overflows a 64 GB arena. That
needs the new diagnostics under the real train-crossasset run. Read the
arena first-overflow log line:

• epochAllocs ~1 and epochMaxAlloc ~64 GB → a single runaway buffer (fix the
  size computation at its gpu_kernels call site).
• resets=0 after several ComputeGradients → Reset/MarkStepBoundary is not
  being driven per step → activations accumulate across samples (drive the step
  boundary; note ADR 005 assumed full-batch retention, but this issue reports
  per-sample gradient accumulation, so the reset cadence is the prime suspect).
• large epochAllocs with offset ~capacity → legitimately full working set
  (the release-threshold hardening is then the safety net; consider batch sizing).

Once that line is captured, the conditional fix is a small follow-up PR.

[dndungu]

Diagnostic captured — verdict: (c) reset-cadence, Wolf-side. Not a ztensor bug.

Ran v1.10.0 under the real train-crossasset workload on GB10 (ZERFOO_ARENA_SIZE_GB=64). The first-overflow line (had to also mirror it to stderr — see note below):

ztensor arena configured: capacityBytes=68719476736 capacityGB=64 managed=false overflowStream=set
ztensor arena first-overflow: path=async requestedBytes=1048576 alignedBytes=1048576 capacityBytes=68719476736 offsetBytes=68719325952 epochAllocs=1033676 epochMaxAllocBytes=1048576 hits=1033675 misses=1 reuses=237039 resets=0 freeListLen=0

Decision tree:

• NOT (a) runaway buffer: epochMaxAllocBytes=1048576 (1 MB), not ~64 GB.
• (c): resets=0 over 1,033,676 allocations. offsetBytes ≈ capacityBytes (arena full). The Wolf trainer does one sample per ComputeGradients (gradient accumulation) and never reset the arena, so ~1M per-sample activation allocs accumulated and filled 64 GB within the first mini-batch (step 0). ADR 005's per-step-reset assumption was correct; Wolf just wasn't driving it.

Fix is Wolf-side (committed): MarkStepBoundary() after the weight/sample upload, then ResetPool() after each per-sample ComputeGradients. Safe because weights/params/pre-uploaded samples are in separate device buffers and gradients/optimizer state are persistent host tensors (graph.Parameter.AddGradient accumulates on .Data()), so only transient activations are reclaimed. Verifying on GB10 once the box is rebooted (the diagnostic run left a wedged overflow-stall pod).

One ztensor note (separate, minor): the v1.10.0 diagnostic lines are routed to the engine's structured logger, but the crossasset train binary creates the engine via NewGPUEngine with log.Nop(), so both lines are discarded in the real workload — SetArenaOverflowLogger is internal/, unreachable from callers. I captured the line via a temp branch (diag/arena-overflow-stderr-118) that also mirrors both lines to stderr. Worth keeping a stderr fallback (or a logger-accepting constructor) so the instrument fires without a real logger.