zzos/tensorflow
1
0
Fork 0
mirror of https://github.com/tensorflow/tensorflow.git synced 2026-09-28 13:23:37 +08:00
Code Issues Packages Projects Releases Wiki Activity

Add defensive checks for P2P symmetric memory support and null windows.

Browse Source 
Add pre-checks for P2P symmetric memory support via cudaDeviceCanAccessPeer,
post-checks for null windows returned from ncclCommWindowRegister, and
defensive null window checks in multimem_addr, peer_addr, and the destructor.

PiperOrigin-RevId: 987435182
This commit is contained in:
Christian Sigg
2026-09-25 08:52:31 -07:00
committed by TensorFlower Gardener
parent d9ecfe9451
commit de517cb79a
1 changed files with 36 additions and 10 deletions
Download Patch File Download Diff File Expand all files Collapse all files
third_party/xla/xla/backends/gpu/collectives/nccl_symmetric_memory.cc
Vendored
+36 -10
View File
@@ -35,7 +35,7 @@ limitations under the License.
// Include NCCL after XLA headers.
#include "third_party/nccl/nccl.h"
#include "third_party/nccl/nccl_device.h"
#include "third_party/nccl/nccl_device.h" // IWYU pragma: keep
#include "xla/tsl/concurrency/executor.h"
namespace xla::gpu {
@@ -79,16 +79,23 @@ NcclSymmetricMemory::Create(std::shared_ptr<NcclCommState> comm_state,
stream_executor::DeviceAddressBase addr,
const std::shared_ptr<tsl::Executor> executor,
stream_executor::StreamExecutor* stream_executor) {
ncclWindow_t win;
if (comm_state == nullptr) {
return absl::InvalidArgumentError("NCCL comm_state is null.");
}
ncclWindow_t win = nullptr;
ncclResult_t nccl_status;
{
VLOG(3) << absl::StrFormat(
"Create NCCL symmetric memory on comm=%p from: ptr=%p; size=%ld",
comm_state->comm, addr.opaque(), addr.size());
absl::MutexLock lock(comm_state->mutex);
nccl_status =
ncclCommWindowRegister(comm_state->comm, addr.opaque(), addr.size(),
&win, NCCL_WIN_COLL_SYMMETRIC);
ncclComm_t comm = comm_state->comm;
if (comm == nullptr) {
return absl::InvalidArgumentError("NCCL comm is null.");
}
VLOG(3) << absl::StrFormat(
"Create NCCL symmetric memory on comm=%p from: ptr=%p; size=%ld", comm,
addr.opaque(), addr.size());
nccl_status = ncclCommWindowRegister(comm, addr.opaque(), addr.size(), &win,
NCCL_WIN_COLL_SYMMETRIC);
}
if (nccl_status != ncclSuccess) {
@@ -96,11 +103,22 @@ NcclSymmetricMemory::Create(std::shared_ptr<NcclCommState> comm_state,
return XLA_NCCL_STATUS(nccl_status);
}
if (win == nullptr) {
LogInterconnectStatus(stream_executor);
return absl::FailedPreconditionError(
"ncclCommWindowRegister returned a null window because NCCL symmetric "
"memory is not supported on this communicator (e.g. P2P/NVLink or "
"CUMEM is unavailable on fractional/virtualized GPU slices).");
}
return absl::WrapUnique(
new NcclSymmetricMemory(comm_state, win, addr, executor));
}
NcclSymmetricMemory::~NcclSymmetricMemory() {
if (win_ == nullptr) {
return;
}
absl::Status status =
Execute(
[&] {
@@ -126,9 +144,13 @@ stream_executor::DeviceAddressBase NcclSymmetricMemory::addr() const {
absl::StatusOr<stream_executor::DeviceAddressBase>
NcclSymmetricMemory::multimem_addr() const {
#if (NCCL_VERSION_CODE >= 22900) || defined(USE_NCCL_HOST_API)
if (win_ == nullptr) {
return absl::FailedPreconditionError(
"Cannot get multimem address from null NCCL window.");
}
void* multimem = nullptr;
XLA_NCCL_RETURN_IF_ERROR(ncclGetLsaMultimemDevicePointer(win_, 0, &multimem));
if (multimem) {
if (multimem != nullptr) {
return stream_executor::DeviceAddressBase(multimem, addr_.size());
}
#endif
@@ -139,10 +161,14 @@ NcclSymmetricMemory::multimem_addr() const {
absl::StatusOr<stream_executor::DeviceAddressBase>
NcclSymmetricMemory::peer_addr(RankId peer) const {
#if (NCCL_VERSION_CODE >= 22902) || defined(USE_NCCL_HOST_API)
if (win_ == nullptr) {
return absl::FailedPreconditionError(
"Cannot get peer address from null NCCL window.");
}
void* peer_addr = nullptr;
XLA_NCCL_RETURN_IF_ERROR(
ncclGetPeerDevicePointer(win_, 0, peer.value(), &peer_addr));
if (peer_addr) {
if (peer_addr != nullptr) {
return stream_executor::DeviceAddressBase(peer_addr, addr_.size());
}
return absl::FailedPreconditionError(absl::StrFormat(