Skip to content

fix(p2p): validate peer copies before auto-enabling and guard ROCm - #3718

Merged
Qubitium merged 1 commit into
mainfrom
fix/p2p-auto-benefit
Nov 9, 2026
Merged

Qubitium merged 1 commit into
mainfrom
fix/p2p-auto-benefit

Conversation

@Qubitium

@Qubitium Qubitium commented Nov 9, 2026

Copy link
Copy Markdown
Collaborator

Summary

Follow-up to #3489. auto mode previously trusted the runtime's peer-copy
capability bit and return code. On ROCm with
PYTORCH_ALLOC_CONF=expandable_segments:True, hipMemcpyPeerAsync can return
hipSuccess and copy nothing (or fault the GPU), so a pair could be enabled and
then produce torn/empty data.

This change makes auto validate 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:

  • copies random payloads (not a deterministic pattern);
  • pre-fills the destination with different random bytes, complementing them so
    a no-op copy can never pass even when both devices' RNGs produce the same
    sequence;
  • verifies the exact byte count and byte-for-byte content;
  • additionally exercises a multi-tensor batch (64 KiB + 256 KiB + 1 MiB) where
    that shape is reachable;
  • then measures direct vs ordered pinned-host staging and keeps the pair only
    when the direct copy is at least as fast as the fallback.

A pair that fails validation is marked degraded and all movers fall back to
the 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 enables
every available direction first, then validates.

3. Keep async buffers alive

The raw hipMemcpyPeerAsync/cudaMemcpyPeerAsync call is invisible to
PyTorch's caching allocator. Without record_stream, the source block could be
freed and reused on the source stream while the destination stream was still
reading it. That produced an async HSA_STATUS_ERROR_MEMORY_FAULT on the
R9700 rig. Both buffers now record the destination stream.

4. Guard the ROCm nested path

The small-leaf p2p_move_nested shape can still fault on a remapped
expandable segment on ROCm, even after the VMM peer-access grant. It is disabled
on ROCm by default and falls back to .to(). Set
GPTQMODEL_P2P_NESTED_ROCM=1 to opt in on a runtime that carries the fix.
Module moves keep using the peer path.

Validation

Python 3.14.10t, PYTHON_GIL=0, gptqmodel at this branch.

rig command result
2x RTX 4090 pytest tests/test_p2p.py tests/test_p2p_copy_utils.py tests/test_tensor_layout.py tests/test_looper_helpers.py 76 passed
2x R9700, unpatched ROCm same 62 passed, 14 skipped (broken VMM peer copy detected and disabled; no fault)
2x R9700, ROCm with #15333 same 72 passed, 4 skipped (beneficial directions enabled)
2x R9700, ROCm with #15333, GPTQMODEL_P2P_NESTED_ROCM=1 nested subset 5 passed, 2 skipped

The unpatched R9700 run previously hung in
test_p2p_move_nested_is_ordered_on_destination_stream with
HSA_STATUS_ERROR_MEMORY_FAULT. With this change it degrades the pair instead
and 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. auto remains
safe 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:
auto may now keep a pair disabled when it is broken or slower than host
staging, and the nested fast path is off on ROCm unless explicitly enabled.

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.
@Qubitium
Qubitium merged commit c97823c into main Nov 9, 2026
6 checks passed
@Qubitium
Qubitium deleted the fix/p2p-auto-benefit branch November 9, 2026 15:13
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