Skip to content

fix: stream-order GPUStorage host access (Bug 11, Wolf batch-3 gradient NaN) - #137

Merged
dndungu merged 3 commits into
mainfrom
fix/d2h-stream-order
Jun 12, 2026
Merged

fix: stream-order GPUStorage host access (Bug 11, Wolf batch-3 gradient NaN)#137
dndungu merged 3 commits into
mainfrom
fix/d2h-stream-order

Conversation

@dndungu

Copy link
Copy Markdown
Contributor

Problem (Bug 11, zerfoo#850 lineage)

Wolf CrossAsset GB10 f32 training (verify5/6 on image d33ee2da, all prior
fixes in) still hits a deterministic NaN in layer0.ffnB1_biases at batch 3:
forward pass clean, per-sample accumulated gradients finite, yet the AdamW
guard reads NaN at the batch boundary. Bugs 9 (zerfoo#851/#852) and 10
(ztensor#134) removed the arena-lifetime causes; the remaining entry point is
the host read itself.

Root cause

GPUStorage host-access paths (TrySlice/Slice/CopyTo/TrySet/Set)
relied on synchronous cudaMemcpy's implicit legacy-stream ordering against
the engine's (blocking) stream. On cache-coherent unified-memory platforms —
the GB10's NVLink-C2C — a "synchronous" pageable-memory copy may be serviced
without waiting for kernels pending on a non-default stream, and managed
storage is accessed directly through the unified pointer with no CUDA call
at all
. A host read can therefore observe bytes from before a still-async
kernel write. In the Wolf schedule, zerfoo's gradAccumulator.addInto host
fallback reads the bias gradient (grad.Data()TrySlice) immediately
after the async Bias Sum kernel is enqueued — garbage bytes → NaN folded
into the persistent accumulator.

Fix (contract level — any workload benefits)

A host access to GPU memory is now stream-ordered with respect to pending
device work:

  • tensor.RegisterHostAccessSync(deviceID, fn) — per-device sync hook
    registry (same indirection pattern as the arena poison kernel fill; tensor
    cannot import compute).
  • Every GPUStorage host read/write runs the device's hooks before touching
    memory (managed and discrete paths alike); hook errors fail the access.
  • GPUEngine registers stream.Synchronize() at construction
    (capture-guarded: syncing during CUDA graph capture is illegal, and capture
    paths never host-read kernel outputs) and unregisters on Close.
  • Graph.Engine() getter so training utilities can run gradient-maintenance
    ops on the graph's own engine/stream (consumed by the zerfoo follow-up).

Tests

tensor/host_access_sync_test.go: a fake runtime where a pending async write
becomes visible only at sync — TrySlice/Slice/CopyTo read stale bytes and
fail without the fix (verified red), plus write-ordering, error-propagation,
and registry semantics. go build ./..., CI-equivalent go vet, and
go test -race ./... all green locally.

…hooks
A host read or write of GPU storage must be stream-ordered with respect
to kernels pending on the owning engine's stream. The previous code
relied on synchronous cudaMemcpy's implicit legacy-stream ordering,
which does not hold for pageable host memory on cache-coherent
unified-memory platforms (GB10 NVLink-C2C), and managed storage was
accessed directly through the unified pointer with no ordering at all.
Result: a host read could observe bytes from before a still-async
kernel write -- the Wolf CrossAsset batch-3 gradient NaN (Bug 11,
zerfoo#850 lineage).
TrySlice/Slice/CopyTo/TrySet/Set now run the host-access sync hooks
registered for their device (tensor.RegisterHostAccessSync) before
touching memory. Regression tests simulate a pending async write that
becomes visible only at sync and fail without the fix.
Training utilities need to run gradient-maintenance ops (persistent
accumulator adds, zerfoo#850) on the same engine -- and therefore the
same device stream -- as the graph's own kernels. The engine was
already stored on the graph; expose a getter.
…ream
NewGPUEngine registers tensor.RegisterHostAccessSync(dev) backed by
stream.Synchronize() (capture-guarded: syncing during CUDA graph
capture is illegal and capture paths never host-read kernel outputs),
and Close unregisters it before destroying the stream. Completes the
Bug 11 host-access ordering contract: every GPUStorage host read/write
now drains this engine's pending kernels first.
@dndungu
dndungu merged commit 3bf58b9 into mainJun 12, 2026
1 check passed
dndungu added a commit to zerfoo/zerfoo that referenced this pull request Jun 12, 2026
…ph engine
Bug 11 (#850 lineage): the gradAccumulator host fallback round-trips
every device gradient through the host (Data() D2H, add, TrySet H2D) once
per sample. On the GB10's coherent unified memory that host read raced the
still-async Bias Sum kernel writing the gradient -- garbage bytes folded
into the persistent accumulator surfaced as the deterministic batch-3 NaN
in layer0.ffnB1_biases (Wolf verify5/6).
The ordering itself is now guaranteed at the contract level by ztensor's
host-access sync hooks (zerfoo/ztensor#137). This change additionally takes
the host round-trip off the per-sample path: when no engine was configured
via SetEngine, the accumulator derives the graph's own engine for fully
device-resident f32 accumulation (both tensors *tensor.GPUStorage), so the
add runs as an in-place kernel on the same stream as the graph's kernels.
Host-backed, non-f32, and engine-less graphs keep the host fallback
unchanged.
dndungu added a commit to zerfoo/zerfoo that referenced this pull request Jun 12, 2026
…ntract)
Pulls zerfoo/ztensor#137: GPUStorage host reads/writes synchronize the
owning engine's stream before touching device memory, and Graph.Engine()
exposes the graph's compute engine for the accumulator's derived-engine
path.
dndungu added a commit that referenced this pull request Jun 12, 2026
…dst-form accumulation policy
design.md documents the contracts shipped for the GPU training
hardening plan in general terms:
- Host-Access Synchronization: host reads/writes of device memory are
stream-ordered via per-device registered sync hooks
(tensor.RegisterHostAccessSync); the GPU engine registers a
capture-guarded stream.Synchronize() (#137).
- Arena Reset: pinning (ADR 006 reset floor / MarkStepBoundary) and
reset epochs -- FreeAtEpoch drops cross-epoch frees, making
GC-finalizer-driven frees safe across Reset (#138).
- Engine dst-form accumulation policy: results write into dst's
existing storage, never re-homed to the arena.
- Graph features: Engine() accessor and the save-for-backward
lifetime contract (ADR 006).
devlog: end-to-end GB10 validation entry for #137/#138 (two clean
runs, zero NaN, accuracy within 0.05pp of the CPU baseline).
Sign up for freeto 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

@dndungu