Repository navigation
fix(p2p): validate peer copies before auto-enabling and guard ROCm - #3718
Merged
Merged
Conversation
auto mode now exercises the real runtime peer-copy path before it trusts a pair. It copies random payloads, pre-fills the destination with different random bytes, and compares the exact byte count and content, so a rc=0 copy that moves nothing cannot enable a broken link. It then keeps the pair only when the direct copy beats ordered pinned-host staging. Enable every direction before validating any of them: the validation copy is issued on the destination device, so it also needs the reverse direction enabled to read the source. Validating inside the enablement loop made the first direction look broken on asymmetric topologies. Record the destination stream on both buffers after each raw runtime copy. The copy is invisible to PyTorch's caching allocator, so without this the source block could be reused on the source stream while the destination stream was still reading it, producing an async HSA_STATUS_ERROR_MEMORY_FAULT. Disable the small-leaf nested peer path on ROCm by default. That shape still faults on a remapped expandable segment even after the VMM peer-access grant; the module path copies larger, validated ranges. Set GPTQMODEL_P2P_NESTED_ROCM=1 to opt in on a runtime that carries the fix.
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
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.
Summary
Follow-up to #3489.
automode previously trusted the runtime's peer-copycapability bit and return code. On ROCm with
PYTORCH_ALLOC_CONF=expandable_segments:True,hipMemcpyPeerAsynccan returnhipSuccessand copy nothing (or fault the GPU), so a pair could be enabled andthen produce torn/empty data.
This change makes
autovalidate the real copy path before enabling a pair,fixes an enable/validate ordering bug, keeps async copies alive for PyTorch's
allocator, and disables the unsafe small-leaf nested path on ROCm.
What changed
1. Validate the actual copy, not the capability bit
_validate_peer_copy_pair()now:a no-op copy can never pass even when both devices' RNGs produce the same
sequence;
that shape is reachable;
when the direct copy is at least as fast as the fallback.
A pair that fails validation is marked
degradedand all movers fall back tothe regular
.to()/ host-staging path.2. Enable all directions before validating any
The validation copy is issued on the destination device, so it also needs the
reverse direction enabled to read the source mapping. The old single loop
validated the first direction before the reverse was enabled, making a healthy
link look broken on asymmetric topologies.
enable_peer_access()now enablesevery available direction first, then validates.
3. Keep async buffers alive
The raw
hipMemcpyPeerAsync/cudaMemcpyPeerAsynccall is invisible toPyTorch's caching allocator. Without
record_stream, the source block could befreed and reused on the source stream while the destination stream was still
reading it. That produced an async
HSA_STATUS_ERROR_MEMORY_FAULTon theR9700 rig. Both buffers now record the destination stream.
4. Guard the ROCm nested path
The small-leaf
p2p_move_nestedshape can still fault on a remappedexpandable segment on ROCm, even after the VMM peer-access grant. It is disabled
on ROCm by default and falls back to
.to(). SetGPTQMODEL_P2P_NESTED_ROCM=1to opt in on a runtime that carries the fix.Module moves keep using the peer path.
Validation
Python 3.14.10t,
PYTHON_GIL=0,gptqmodelat this branch.pytest tests/test_p2p.py tests/test_p2p_copy_utils.py tests/test_tensor_layout.py tests/test_looper_helpers.pyGPTQMODEL_P2P_NESTED_ROCM=1The unpatched R9700 run previously hung in
test_p2p_move_nested_is_ordered_on_destination_streamwithHSA_STATUS_ERROR_MEMORY_FAULT. With this change it degrades the pair insteadand the nested tests skip; no GPU fault or hang.
ROCm context
The runtime issue is the ROCr VMM access-validation stub:
Both were reproduced and validated on the same 2x R9700 rig.
autoremainssafe on runtimes without either fix because it validates the real copy path.
Compatibility
No
hf_-prefixed API is touched. Non-hf_behavior changes are intentional:automay now keep a pair disabled when it is broken or slower than hoststaging, and the nested fast path is off on ROCm unless explicitly enabled.