weight-sync: use PJRT event-completion path for D2H (fix silent rank-1 corruption) - #664
Open
lokic233 wants to merge 1 commit into
Open
weight-sync: use PJRT event-completion path for D2H (fix silent rank-1 corruption)#664lokic233 wants to merge 1 commit into
lokic233 wants to merge 1 commit into
Conversation
…1 corruption) WeightSynchronizerBase::D2hLayer issued device->host copies via the legacy CopyRawDeviceToHost (xla::Future) path and detiled inside copy_future.Map(...). On the TPU C-API under concurrency that future can report ready before the DMA host-write lands, so a nondeterministic subset of small rank-1 (norm/scale) tensors publish silent zeros while d2h() returns success. Fix: route D2H through the canonical PJRT event-completion path already used by the KV-cache raw-transfer helper (IssueD2hShard -> CopyRawDeviceToHostEvent / supports_event() -> FromEvents -> PJRT_Event_Await), with 64-byte-aligned temp buffers, JoinPjRtCopyFutures().Await(), and detile AFTER await. No API change. Silicon-validated on HEAD (jax 0.11.0 / libtpu 0.0.44, tpu7x): the concurrent rank-1 reproducer went from firing (5-18 zero leaves at N=40, api success) to 392/392 leaves ok / 0 zero / 0 wrong across N=4/20/40.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Problem
WeightSynchronizerBase::D2hLayerissues device→host copies via the legacyCopyRawDeviceToHost(xla::Future) path and detiles insidecopy_future.Map(...), wrapping the result withPjRtCopyFuture::FromFuture. On the TPU C-API, under concurrency, that future can report readiness before the DMA host-write has landed. As a result a nondeterministic subset of small rank-1 (norm/scale) tensors publish silent zeros whiled2h()returns success — correctness-fatal and fail-silent.The sibling KV-cache raw-transfer path already uses the canonical PJRT event-completion contract (
IssueD2hShard→CopyRawDeviceToHostEvent/supports_event()); the weight-sync path did not.Reproducer
Many concurrent rank-1 bf16 leaves, each with distinct non-zero content; call
d2h(), then inspect each host staging buffer (classify ok / zero / wrong) and capture the API result. Onset at N≥4; failure count scales with concurrency.Root cause
The legacy
CopyRawDeviceToHostxla::Futuredoes not reliably gate on DMA host-write completion for small concurrent copies, so the detile (inside.Map) reads a partially/never-written temp buffer.Fix (correctness only, no API change)
Route D2H through the same canonical PJRT event-completion path the KV path uses:
raiden::IssueD2hShard→CopyRawDeviceToHostEvent(whensupports_event()) →FromEvents→PJRT_Event_Await, with 64-byte-aligned temp buffers,JoinPjRtCopyFutures().Await(), and detile only AFTER the await. The public signature is unchanged; the returned future is already-awaited.Validation (silicon)
Built and run on current HEAD (
a6507fe), jax 0.11.0 / libtpu 0.0.44, tpu7x:Across N=4/20/40 (multiple reps): 392/392 leaves ok, 0 zero, 0 wrong,
d2h()success every rep. A concurrent multi-leaf byte-verify regression test is included (API success alone is insufficient — destination bytes must be checked).This PR is correctness-only and deliberately excludes an unrelated bf16 detile performance change (submitted separately).
Authored with agentic assistance (Navi) and validated on TPU silicon before submission.