[Core][Distributed] improve p2p access check

[youkaichao]

Done

[youkaichao]

Previously, we use the following check for actual p2p access in case cuda driver is broken:

# code partly borrowed from
# https://github.com/turboderp/exllamav2/blob/1c67f97f3d2a968605a9c31ab791a05c85bb7879/exllamav2/compat.py#L10
# License: MIT
def _can_actually_p2p(idx_a, idx_b):
    dev_i = f"cuda:{idx_a}"
    dev_j = f"cuda:{idx_b}"
    a = torch.randn(5, device=dev_i) + 123.0
    b = a.to(dev_j)
    c = b.to(dev_i)
    return torch.all(a == c).cpu().item()

However, pytorch somehow fixes the bug, and it will always return True, no matter whether p2p is available:

import torch
torch.cuda.can_device_access_peer(0, 1) # False
_can_actually_p2p(0, 1) # True

This is reported in #4770 (comment) .

[youkaichao]

cc @hanzhi713

[youkaichao]

@WoosukKwon ready for review

[WoosukKwon]

LGTM! Thanks for the PR! Left some minor comments.

[youkaichao]

@WoosukKwon thanks for the very detailed review!

[youkaichao]

Since we still don't have ci machines with p2p capability, I tested this PR locally.

cc @simon-mo for nvlink machines.