[NVIDIA GPU] Use memcpy for intra-node all-to-all #15144

terryysun · 2024-07-19T21:19:37Z

The communications of all-to-all rely on NCCL even when it is intra-node. By leveraging memcpy for intra-node communications, all-to-all can have better performance while reducing SM consumption (right now consumed by NCCL).

terryysun · 2024-07-19T21:25:31Z

some screenshots showing the performance difference between NCCL call and memcpy for an example all-to-all.

frgossen

Very nice!

You are adding this for all-gather, all-reduce, all-to-all, collective-broadcast, and collective-permute but I only see a test case for all-to-all.
Can we add test cases for the remaining collectives? would also suggest to split this up into one PR per collective to keep reviews simpler and in case we have to roll back something.

terryysun · 2024-07-23T06:26:37Z

Very nice!

You are adding this for all-gather, all-reduce, all-to-all, collective-broadcast, and collective-permute but I only see a test case for all-to-all. Can we add test cases for the remaining collectives? would also suggest to split this up into one PR per collective to keep reviews simpler and in case we have to roll back something.

Thanks! Right now it's only added for all-to-all, the api changes for the other collectives are just to avoid interim intricacies -- we do plan to add similar support for them in the near future, but the flag is not consumed in this PR.

frgossen

Thanks for clarifying. All makes sense to me, just a few minot comments.

frgossen · 2024-07-29T13:38:48Z

xla/service/gpu/runtime/nccl_collective_thunk.h

@@ -71,6 +71,21 @@ struct NcclCollectiveConfig {
  bool IsDegenerate(int64_t replica_count, int64_t partition_count) const;
 };

+template <typename T>
+absl::StatusOr<const int64_t> GetCurrentId(


Rather than templating, would it make sense to impl this base on the NcclCollectiveConfig?

after reorganizing the code turns out we don't need this util, removed.

frgossen · 2024-07-29T13:42:59Z

xla/service/gpu/runtime/nccl_all_to_all_thunk.h

+
+    absl::Status PutRecvPtr(int64_t send_id, int64_t recv_id, void* ptr) {
+      if (!IsInitialized(send_id, recv_id)) {
+        return absl::InternalError(absl::StrCat("Send-receive pair ", send_id, ", ", recv_id,


updated to receive.

frgossen · 2024-07-29T13:45:38Z

xla/service/gpu/runtime/nccl_all_to_all_thunk.cc

  return xla::gpu::RunAllToAll(nccl_api(), config_.has_split_dimension,
                               device_buffers, stream,
-                               comm_wrapper.comm_handle);
+                               comm_wrapper.comm_handle, current_id, use_memcpy, recv_ptr_map_);
 }

 absl::Status RunAllToAll(NcclApi* nccl_api, bool has_split_dimension,


This is implemented in two modes which do not share much code. Can we outline that into two functions and dispatch here?

broke to two functions.

frgossen · 2024-09-09T10:57:22Z

I see you pushed changes and requested a review. Can you reply to the comments and explain how the changes address them? Ty.

terryysun · 2024-09-18T01:09:11Z

I see you pushed changes and requested a review. Can you reply to the comments and explain how the changes address them? Ty.

hey @frgossen sorry for the delayed reply, we were verifying the changes and fixed multiple issues we saw when running realistic models. I've updated the code and replied to the comments accordingly. Could you take another look? Thanks!

frgossen

Thanks!

Imported from GitHub PR #15144 The communications of all-to-all rely on NCCL even when it is intra-node. By leveraging memcpy for intra-node communications, all-to-all can have better performance while reducing SM consumption (right now consumed by NCCL). Copybara import of the project: -- 38720c7 by Terry Sun <[email protected]>: memcpyp2p for local a2a -- 90018f4 by Terry Sun <[email protected]>: use nccl to pass recv ptrs -- f9b75b0 by Terry Sun <[email protected]>: refactor and cleanup Merging this change closes #15144 FUTURE_COPYBARA_INTEGRATE_REVIEW=#15144 from terryysun:terryysun/all2all_memcpyp2p f9b75b0 PiperOrigin-RevId: 678378925

Imported from GitHub PR openxla/xla#15144 The communications of all-to-all rely on NCCL even when it is intra-node. By leveraging memcpy for intra-node communications, all-to-all can have better performance while reducing SM consumption (right now consumed by NCCL). Copybara import of the project: -- 38720c73f5817dbbf5b6d98751140bb53f572690 by Terry Sun <[email protected]>: memcpyp2p for local a2a -- 90018f4a3fe0ed3018767db810518faf9435bc93 by Terry Sun <[email protected]>: use nccl to pass recv ptrs -- f9b75b0e088286ded770b27fff9d020f8e85a648 by Terry Sun <[email protected]>: refactor and cleanup Merging this change closes #15144 FUTURE_COPYBARA_INTEGRATE_REVIEW=openxla/xla#15144 from terryysun:terryysun/all2all_memcpyp2p f9b75b0e088286ded770b27fff9d020f8e85a648 PiperOrigin-RevId: 678378925

Imported from GitHub PR #15144 The communications of all-to-all rely on NCCL even when it is intra-node. By leveraging memcpy for intra-node communications, all-to-all can have better performance while reducing SM consumption (right now consumed by NCCL). Copybara import of the project: -- 38720c7 by Terry Sun <[email protected]>: memcpyp2p for local a2a -- 90018f4 by Terry Sun <[email protected]>: use nccl to pass recv ptrs -- f9b75b0 by Terry Sun <[email protected]>: refactor and cleanup Merging this change closes #15144 FUTURE_COPYBARA_INTEGRATE_REVIEW=#15144 from terryysun:terryysun/all2all_memcpyp2p f9b75b0 PiperOrigin-RevId: 678378925

Imported from GitHub PR openxla/xla#15144 The communications of all-to-all rely on NCCL even when it is intra-node. By leveraging memcpy for intra-node communications, all-to-all can have better performance while reducing SM consumption (right now consumed by NCCL). Copybara import of the project: -- 38720c73f5817dbbf5b6d98751140bb53f572690 by Terry Sun <[email protected]>: memcpyp2p for local a2a -- 90018f4a3fe0ed3018767db810518faf9435bc93 by Terry Sun <[email protected]>: use nccl to pass recv ptrs -- f9b75b0e088286ded770b27fff9d020f8e85a648 by Terry Sun <[email protected]>: refactor and cleanup Merging this change closes #15144 FUTURE_COPYBARA_INTEGRATE_REVIEW=openxla/xla#15144 from terryysun:terryysun/all2all_memcpyp2p f9b75b0e088286ded770b27fff9d020f8e85a648 PiperOrigin-RevId: 678378925

…dress sharing Imported from GitHub PR #17636 This is a followup PR to #15144. A distributed cache is maintained when device addresses are shared across ranks. There are two issues withe the existing implementation: 1. The cache is not guarded by mutex; 2. The cache initialization process have redundant access. These issues can cause race condition or dead lock when the progress on different ranks are very close. Consequently we need to introduce below enhancements: 1. Guard the cache with mutex; 2. Shard the initialization process by rank, so that each rank only handle a piece of the cache and should not have overlapping access in theory. Copybara import of the project: -- a6472fc by Terry Sun <[email protected]>: enhance concurrency handling -- 356ab82 by Terry Sun <[email protected]>: lock mutex -- 29ebb2d by Terry Sun <[email protected]>: bring back test -- 91b911f by Terry Sun <[email protected]>: better lock granularity Merging this change closes #17636 FUTURE_COPYBARA_INTEGRATE_REVIEW=#17636 from terryysun:terryysun/sync_fix 91b911f PiperOrigin-RevId: 679463524

…dress sharing Imported from GitHub PR openxla/xla#17636 This is a followup PR to openxla/xla#15144. A distributed cache is maintained when device addresses are shared across ranks. There are two issues withe the existing implementation: 1. The cache is not guarded by mutex; 2. The cache initialization process have redundant access. These issues can cause race condition or dead lock when the progress on different ranks are very close. Consequently we need to introduce below enhancements: 1. Guard the cache with mutex; 2. Shard the initialization process by rank, so that each rank only handle a piece of the cache and should not have overlapping access in theory. Copybara import of the project: -- a6472fc75fd0411bd8e65f27082e21e9a946ab17 by Terry Sun <[email protected]>: enhance concurrency handling -- 356ab824b95d66c793e361882e95d70689759ffd by Terry Sun <[email protected]>: lock mutex -- 29ebb2de64711bf4b4a08cf1593317228b56f825 by Terry Sun <[email protected]>: bring back test -- 91b911f0aaac0e590636a82956b464436e94ef9f by Terry Sun <[email protected]>: better lock granularity Merging this change closes #17636 FUTURE_COPYBARA_INTEGRATE_REVIEW=openxla/xla#17636 from terryysun:terryysun/sync_fix 91b911f0aaac0e590636a82956b464436e94ef9f PiperOrigin-RevId: 679463524

…dress sharing Imported from GitHub PR #17636 This is a followup PR to #15144. A distributed cache is maintained when device addresses are shared across ranks. There are two issues withe the existing implementation: 1. The cache is not guarded by mutex; 2. The cache initialization process have redundant access. These issues can cause race condition or dead lock when the progress on different ranks are very close. Consequently we need to introduce below enhancements: 1. Guard the cache with mutex; 2. Shard the initialization process by rank, so that each rank only handle a piece of the cache and should not have overlapping access in theory. Copybara import of the project: -- a6472fc by Terry Sun <[email protected]>: enhance concurrency handling -- 356ab82 by Terry Sun <[email protected]>: lock mutex -- 29ebb2d by Terry Sun <[email protected]>: bring back test -- 91b911f by Terry Sun <[email protected]>: better lock granularity Merging this change closes #17636 FUTURE_COPYBARA_INTEGRATE_REVIEW=#17636 from terryysun:terryysun/sync_fix 91b911f PiperOrigin-RevId: 679463524

…dress sharing Imported from GitHub PR openxla/xla#17636 This is a followup PR to openxla/xla#15144. A distributed cache is maintained when device addresses are shared across ranks. There are two issues withe the existing implementation: 1. The cache is not guarded by mutex; 2. The cache initialization process have redundant access. These issues can cause race condition or dead lock when the progress on different ranks are very close. Consequently we need to introduce below enhancements: 1. Guard the cache with mutex; 2. Shard the initialization process by rank, so that each rank only handle a piece of the cache and should not have overlapping access in theory. Copybara import of the project: -- a6472fc75fd0411bd8e65f27082e21e9a946ab17 by Terry Sun <[email protected]>: enhance concurrency handling -- 356ab824b95d66c793e361882e95d70689759ffd by Terry Sun <[email protected]>: lock mutex -- 29ebb2de64711bf4b4a08cf1593317228b56f825 by Terry Sun <[email protected]>: bring back test -- 91b911f0aaac0e590636a82956b464436e94ef9f by Terry Sun <[email protected]>: better lock granularity Merging this change closes #17636 FUTURE_COPYBARA_INTEGRATE_REVIEW=openxla/xla#17636 from terryysun:terryysun/sync_fix 91b911f0aaac0e590636a82956b464436e94ef9f PiperOrigin-RevId: 679463524

…dress sharing Imported from GitHub PR #17636 This is a followup PR to #15144. A distributed cache is maintained when device addresses are shared across ranks. There are two issues withe the existing implementation: 1. The cache is not guarded by mutex; 2. The cache initialization process have redundant access. These issues can cause race condition or dead lock when the progress on different ranks are very close. Consequently we need to introduce below enhancements: 1. Guard the cache with mutex; 2. Shard the initialization process by rank, so that each rank only handle a piece of the cache and should not have overlapping access in theory. Copybara import of the project: -- a6472fc by Terry Sun <[email protected]>: enhance concurrency handling -- 356ab82 by Terry Sun <[email protected]>: lock mutex -- 29ebb2d by Terry Sun <[email protected]>: bring back test -- 91b911f by Terry Sun <[email protected]>: better lock granularity -- cc1d93a by Terry Sun <[email protected]>: guard all accesses Merging this change closes #17636 FUTURE_COPYBARA_INTEGRATE_REVIEW=#17636 from terryysun:terryysun/sync_fix cc1d93a PiperOrigin-RevId: 679463524

…dress sharing Imported from GitHub PR openxla/xla#17636 This is a followup PR to openxla/xla#15144. A distributed cache is maintained when device addresses are shared across ranks. There are two issues withe the existing implementation: 1. The cache is not guarded by mutex; 2. The cache initialization process have redundant access. These issues can cause race condition or dead lock when the progress on different ranks are very close. Consequently we need to introduce below enhancements: 1. Guard the cache with mutex; 2. Shard the initialization process by rank, so that each rank only handle a piece of the cache and should not have overlapping access in theory. Copybara import of the project: -- a6472fc75fd0411bd8e65f27082e21e9a946ab17 by Terry Sun <[email protected]>: enhance concurrency handling -- 356ab824b95d66c793e361882e95d70689759ffd by Terry Sun <[email protected]>: lock mutex -- 29ebb2de64711bf4b4a08cf1593317228b56f825 by Terry Sun <[email protected]>: bring back test -- 91b911f0aaac0e590636a82956b464436e94ef9f by Terry Sun <[email protected]>: better lock granularity -- cc1d93a5f1032a205473961b2c2d3e14bee3a9c6 by Terry Sun <[email protected]>: guard all accesses Merging this change closes #17636 FUTURE_COPYBARA_INTEGRATE_REVIEW=openxla/xla#17636 from terryysun:terryysun/sync_fix cc1d93a5f1032a205473961b2c2d3e14bee3a9c6 PiperOrigin-RevId: 679463524

…dress sharing Imported from GitHub PR #17636 This is a followup PR to #15144. A distributed cache is maintained when device addresses are shared across ranks. There are two issues withe the existing implementation: 1. The cache is not guarded by mutex; 2. The cache initialization process have redundant access. These issues can cause race condition or dead lock when the progress on different ranks are very close. Consequently we need to introduce below enhancements: 1. Guard the cache with mutex; 2. Shard the initialization process by rank, so that each rank only handle a piece of the cache and should not have overlapping access in theory. Copybara import of the project: -- a6472fc by Terry Sun <[email protected]>: enhance concurrency handling -- 356ab82 by Terry Sun <[email protected]>: lock mutex -- 29ebb2d by Terry Sun <[email protected]>: bring back test -- 91b911f by Terry Sun <[email protected]>: better lock granularity -- cc1d93a by Terry Sun <[email protected]>: guard all accesses Merging this change closes #17636 FUTURE_COPYBARA_INTEGRATE_REVIEW=#17636 from terryysun:terryysun/sync_fix cc1d93a PiperOrigin-RevId: 679463524

…dress sharing Imported from GitHub PR openxla/xla#17636 This is a followup PR to openxla/xla#15144. A distributed cache is maintained when device addresses are shared across ranks. There are two issues withe the existing implementation: 1. The cache is not guarded by mutex; 2. The cache initialization process have redundant access. These issues can cause race condition or dead lock when the progress on different ranks are very close. Consequently we need to introduce below enhancements: 1. Guard the cache with mutex; 2. Shard the initialization process by rank, so that each rank only handle a piece of the cache and should not have overlapping access in theory. Copybara import of the project: -- a6472fc75fd0411bd8e65f27082e21e9a946ab17 by Terry Sun <[email protected]>: enhance concurrency handling -- 356ab824b95d66c793e361882e95d70689759ffd by Terry Sun <[email protected]>: lock mutex -- 29ebb2de64711bf4b4a08cf1593317228b56f825 by Terry Sun <[email protected]>: bring back test -- 91b911f0aaac0e590636a82956b464436e94ef9f by Terry Sun <[email protected]>: better lock granularity -- cc1d93a5f1032a205473961b2c2d3e14bee3a9c6 by Terry Sun <[email protected]>: guard all accesses Merging this change closes #17636 FUTURE_COPYBARA_INTEGRATE_REVIEW=openxla/xla#17636 from terryysun:terryysun/sync_fix cc1d93a5f1032a205473961b2c2d3e14bee3a9c6 PiperOrigin-RevId: 679463524

…dress sharing Imported from GitHub PR #17636 This is a followup PR to #15144. A distributed cache is maintained when device addresses are shared across ranks. There are two issues withe the existing implementation: 1. The cache is not guarded by mutex; 2. The cache initialization process have redundant access. These issues can cause race condition or dead lock when the progress on different ranks are very close. Consequently we need to introduce below enhancements: 1. Guard the cache with mutex; 2. Shard the initialization process by rank, so that each rank only handle a piece of the cache and should not have overlapping access in theory. Copybara import of the project: -- a6472fc by Terry Sun <[email protected]>: enhance concurrency handling -- 356ab82 by Terry Sun <[email protected]>: lock mutex -- 29ebb2d by Terry Sun <[email protected]>: bring back test -- 91b911f by Terry Sun <[email protected]>: better lock granularity -- cc1d93a by Terry Sun <[email protected]>: guard all accesses Merging this change closes #17636 FUTURE_COPYBARA_INTEGRATE_REVIEW=#17636 from terryysun:terryysun/sync_fix cc1d93a PiperOrigin-RevId: 679463524

…dress sharing Imported from GitHub PR openxla/xla#17636 This is a followup PR to openxla/xla#15144. A distributed cache is maintained when device addresses are shared across ranks. There are two issues withe the existing implementation: 1. The cache is not guarded by mutex; 2. The cache initialization process have redundant access. These issues can cause race condition or dead lock when the progress on different ranks are very close. Consequently we need to introduce below enhancements: 1. Guard the cache with mutex; 2. Shard the initialization process by rank, so that each rank only handle a piece of the cache and should not have overlapping access in theory. Copybara import of the project: -- a6472fc75fd0411bd8e65f27082e21e9a946ab17 by Terry Sun <[email protected]>: enhance concurrency handling -- 356ab824b95d66c793e361882e95d70689759ffd by Terry Sun <[email protected]>: lock mutex -- 29ebb2de64711bf4b4a08cf1593317228b56f825 by Terry Sun <[email protected]>: bring back test -- 91b911f0aaac0e590636a82956b464436e94ef9f by Terry Sun <[email protected]>: better lock granularity -- cc1d93a5f1032a205473961b2c2d3e14bee3a9c6 by Terry Sun <[email protected]>: guard all accesses Merging this change closes #17636 FUTURE_COPYBARA_INTEGRATE_REVIEW=openxla/xla#17636 from terryysun:terryysun/sync_fix cc1d93a5f1032a205473961b2c2d3e14bee3a9c6 PiperOrigin-RevId: 679463524

…dress sharing Imported from GitHub PR #17636 This is a followup PR to #15144. A distributed cache is maintained when device addresses are shared across ranks. There are two issues withe the existing implementation: 1. The cache is not guarded by mutex; 2. The cache initialization process have redundant access. These issues can cause race condition or dead lock when the progress on different ranks are very close. Consequently we need to introduce below enhancements: 1. Guard the cache with mutex; 2. Shard the initialization process by rank, so that each rank only handle a piece of the cache and should not have overlapping access in theory. Copybara import of the project: -- a6472fc by Terry Sun <[email protected]>: enhance concurrency handling -- 356ab82 by Terry Sun <[email protected]>: lock mutex -- 29ebb2d by Terry Sun <[email protected]>: bring back test -- 91b911f by Terry Sun <[email protected]>: better lock granularity -- cc1d93a by Terry Sun <[email protected]>: guard all accesses Merging this change closes #17636 FUTURE_COPYBARA_INTEGRATE_REVIEW=#17636 from terryysun:terryysun/sync_fix cc1d93a PiperOrigin-RevId: 679464553

…dress sharing Imported from GitHub PR openxla/xla#17636 This is a followup PR to openxla/xla#15144. A distributed cache is maintained when device addresses are shared across ranks. There are two issues withe the existing implementation: 1. The cache is not guarded by mutex; 2. The cache initialization process have redundant access. These issues can cause race condition or dead lock when the progress on different ranks are very close. Consequently we need to introduce below enhancements: 1. Guard the cache with mutex; 2. Shard the initialization process by rank, so that each rank only handle a piece of the cache and should not have overlapping access in theory. Copybara import of the project: -- a6472fc75fd0411bd8e65f27082e21e9a946ab17 by Terry Sun <[email protected]>: enhance concurrency handling -- 356ab824b95d66c793e361882e95d70689759ffd by Terry Sun <[email protected]>: lock mutex -- 29ebb2de64711bf4b4a08cf1593317228b56f825 by Terry Sun <[email protected]>: bring back test -- 91b911f0aaac0e590636a82956b464436e94ef9f by Terry Sun <[email protected]>: better lock granularity -- cc1d93a5f1032a205473961b2c2d3e14bee3a9c6 by Terry Sun <[email protected]>: guard all accesses Merging this change closes #17636 FUTURE_COPYBARA_INTEGRATE_REVIEW=openxla/xla#17636 from terryysun:terryysun/sync_fix cc1d93a5f1032a205473961b2c2d3e14bee3a9c6 PiperOrigin-RevId: 679464553

…dress sharing Imported from GitHub PR #17636 This is a followup PR to #15144. A distributed cache is maintained when device addresses are shared across ranks. There are two issues withe the existing implementation: 1. The cache is not guarded by mutex; 2. The cache initialization process have redundant access. These issues can cause race condition or dead lock when the progress on different ranks are very close. Consequently we need to introduce below enhancements: 1. Guard the cache with mutex; 2. Shard the initialization process by rank, so that each rank only handle a piece of the cache and should not have overlapping access in theory. Copybara import of the project: -- a6472fc by Terry Sun <[email protected]>: enhance concurrency handling -- 356ab82 by Terry Sun <[email protected]>: lock mutex -- 29ebb2d by Terry Sun <[email protected]>: bring back test -- 91b911f by Terry Sun <[email protected]>: better lock granularity -- cc1d93a by Terry Sun <[email protected]>: guard all accesses Merging this change closes #17636 FUTURE_COPYBARA_INTEGRATE_REVIEW=#17636 from terryysun:terryysun/sync_fix cc1d93a PiperOrigin-RevId: 679464553

…dress sharing Imported from GitHub PR openxla/xla#17636 This is a followup PR to openxla/xla#15144. A distributed cache is maintained when device addresses are shared across ranks. There are two issues withe the existing implementation: 1. The cache is not guarded by mutex; 2. The cache initialization process have redundant access. These issues can cause race condition or dead lock when the progress on different ranks are very close. Consequently we need to introduce below enhancements: 1. Guard the cache with mutex; 2. Shard the initialization process by rank, so that each rank only handle a piece of the cache and should not have overlapping access in theory. Copybara import of the project: -- a6472fc75fd0411bd8e65f27082e21e9a946ab17 by Terry Sun <[email protected]>: enhance concurrency handling -- 356ab824b95d66c793e361882e95d70689759ffd by Terry Sun <[email protected]>: lock mutex -- 29ebb2de64711bf4b4a08cf1593317228b56f825 by Terry Sun <[email protected]>: bring back test -- 91b911f0aaac0e590636a82956b464436e94ef9f by Terry Sun <[email protected]>: better lock granularity -- cc1d93a5f1032a205473961b2c2d3e14bee3a9c6 by Terry Sun <[email protected]>: guard all accesses Merging this change closes #17636 FUTURE_COPYBARA_INTEGRATE_REVIEW=openxla/xla#17636 from terryysun:terryysun/sync_fix cc1d93a5f1032a205473961b2c2d3e14bee3a9c6 PiperOrigin-RevId: 679464553

…dress sharing Imported from GitHub PR #17636 This is a followup PR to #15144. A distributed cache is maintained when device addresses are shared across ranks. There are two issues withe the existing implementation: 1. The cache is not guarded by mutex; 2. The cache initialization process have redundant access. These issues can cause race condition or dead lock when the progress on different ranks are very close. Consequently we need to introduce below enhancements: 1. Guard the cache with mutex; 2. Shard the initialization process by rank, so that each rank only handle a piece of the cache and should not have overlapping access in theory. Copybara import of the project: -- a6472fc by Terry Sun <[email protected]>: enhance concurrency handling -- 356ab82 by Terry Sun <[email protected]>: lock mutex -- 29ebb2d by Terry Sun <[email protected]>: bring back test -- 91b911f by Terry Sun <[email protected]>: better lock granularity -- cc1d93a by Terry Sun <[email protected]>: guard all accesses Merging this change closes #17636 FUTURE_COPYBARA_INTEGRATE_REVIEW=#17636 from terryysun:terryysun/sync_fix cc1d93a PiperOrigin-RevId: 679464553

…dress sharing Imported from GitHub PR openxla/xla#17636 This is a followup PR to openxla/xla#15144. A distributed cache is maintained when device addresses are shared across ranks. There are two issues withe the existing implementation: 1. The cache is not guarded by mutex; 2. The cache initialization process have redundant access. These issues can cause race condition or dead lock when the progress on different ranks are very close. Consequently we need to introduce below enhancements: 1. Guard the cache with mutex; 2. Shard the initialization process by rank, so that each rank only handle a piece of the cache and should not have overlapping access in theory. Copybara import of the project: -- a6472fc75fd0411bd8e65f27082e21e9a946ab17 by Terry Sun <[email protected]>: enhance concurrency handling -- 356ab824b95d66c793e361882e95d70689759ffd by Terry Sun <[email protected]>: lock mutex -- 29ebb2de64711bf4b4a08cf1593317228b56f825 by Terry Sun <[email protected]>: bring back test -- 91b911f0aaac0e590636a82956b464436e94ef9f by Terry Sun <[email protected]>: better lock granularity -- cc1d93a5f1032a205473961b2c2d3e14bee3a9c6 by Terry Sun <[email protected]>: guard all accesses Merging this change closes #17636 FUTURE_COPYBARA_INTEGRATE_REVIEW=openxla/xla#17636 from terryysun:terryysun/sync_fix cc1d93a5f1032a205473961b2c2d3e14bee3a9c6 PiperOrigin-RevId: 679464553

…dress sharing Imported from GitHub PR #17636 This is a followup PR to #15144. A distributed cache is maintained when device addresses are shared across ranks. There are two issues withe the existing implementation: 1. The cache is not guarded by mutex; 2. The cache initialization process have redundant access. These issues can cause race condition or dead lock when the progress on different ranks are very close. Consequently we need to introduce below enhancements: 1. Guard the cache with mutex; 2. Shard the initialization process by rank, so that each rank only handle a piece of the cache and should not have overlapping access in theory. Copybara import of the project: -- a6472fc by Terry Sun <[email protected]>: enhance concurrency handling -- 356ab82 by Terry Sun <[email protected]>: lock mutex -- 29ebb2d by Terry Sun <[email protected]>: bring back test -- 91b911f by Terry Sun <[email protected]>: better lock granularity -- cc1d93a by Terry Sun <[email protected]>: guard all accesses Merging this change closes #17636 FUTURE_COPYBARA_INTEGRATE_REVIEW=#17636 from terryysun:terryysun/sync_fix cc1d93a PiperOrigin-RevId: 679464553

…dress sharing Imported from GitHub PR openxla/xla#17636 This is a followup PR to openxla/xla#15144. A distributed cache is maintained when device addresses are shared across ranks. There are two issues withe the existing implementation: 1. The cache is not guarded by mutex; 2. The cache initialization process have redundant access. These issues can cause race condition or dead lock when the progress on different ranks are very close. Consequently we need to introduce below enhancements: 1. Guard the cache with mutex; 2. Shard the initialization process by rank, so that each rank only handle a piece of the cache and should not have overlapping access in theory. Copybara import of the project: -- a6472fc75fd0411bd8e65f27082e21e9a946ab17 by Terry Sun <[email protected]>: enhance concurrency handling -- 356ab824b95d66c793e361882e95d70689759ffd by Terry Sun <[email protected]>: lock mutex -- 29ebb2de64711bf4b4a08cf1593317228b56f825 by Terry Sun <[email protected]>: bring back test -- 91b911f0aaac0e590636a82956b464436e94ef9f by Terry Sun <[email protected]>: better lock granularity -- cc1d93a5f1032a205473961b2c2d3e14bee3a9c6 by Terry Sun <[email protected]>: guard all accesses Merging this change closes #17636 FUTURE_COPYBARA_INTEGRATE_REVIEW=openxla/xla#17636 from terryysun:terryysun/sync_fix cc1d93a5f1032a205473961b2c2d3e14bee3a9c6 PiperOrigin-RevId: 679464553

…dress sharing Imported from GitHub PR #17636 This is a followup PR to #15144. A distributed cache is maintained when device addresses are shared across ranks. There are two issues withe the existing implementation: 1. The cache is not guarded by mutex; 2. The cache initialization process have redundant access. These issues can cause race condition or dead lock when the progress on different ranks are very close. Consequently we need to introduce below enhancements: 1. Guard the cache with mutex; 2. Shard the initialization process by rank, so that each rank only handle a piece of the cache and should not have overlapping access in theory. Copybara import of the project: -- a6472fc by Terry Sun <[email protected]>: enhance concurrency handling -- 356ab82 by Terry Sun <[email protected]>: lock mutex -- 29ebb2d by Terry Sun <[email protected]>: bring back test -- 91b911f by Terry Sun <[email protected]>: better lock granularity -- cc1d93a by Terry Sun <[email protected]>: guard all accesses Merging this change closes #17636 COPYBARA_INTEGRATE_REVIEW=#17636 from terryysun:terryysun/sync_fix cc1d93a PiperOrigin-RevId: 681120763

…dress sharing Imported from GitHub PR openxla/xla#17636 This is a followup PR to openxla/xla#15144. A distributed cache is maintained when device addresses are shared across ranks. There are two issues withe the existing implementation: 1. The cache is not guarded by mutex; 2. The cache initialization process have redundant access. These issues can cause race condition or dead lock when the progress on different ranks are very close. Consequently we need to introduce below enhancements: 1. Guard the cache with mutex; 2. Shard the initialization process by rank, so that each rank only handle a piece of the cache and should not have overlapping access in theory. Copybara import of the project: -- a6472fc75fd0411bd8e65f27082e21e9a946ab17 by Terry Sun <[email protected]>: enhance concurrency handling -- 356ab824b95d66c793e361882e95d70689759ffd by Terry Sun <[email protected]>: lock mutex -- 29ebb2de64711bf4b4a08cf1593317228b56f825 by Terry Sun <[email protected]>: bring back test -- 91b911f0aaac0e590636a82956b464436e94ef9f by Terry Sun <[email protected]>: better lock granularity -- cc1d93a5f1032a205473961b2c2d3e14bee3a9c6 by Terry Sun <[email protected]>: guard all accesses Merging this change closes #17636 PiperOrigin-RevId: 681120763

loislo · 2024-10-02T12:47:49Z

From time to time a test xla/tests/collective_ops_e2e_test.cc:486 xla::(anonymous namespace)::CollectiveOpsTestE2E_AsyncAllToAllMemCpy_Test::TestBody() is failing due to a thread gets stuck.

This is the stacktrace for this thread.

E1002 05:16:35.729531 5068 rendezvous.cc:43] Stack for thread tf_replicas/5068:
@ clock_gettime
@ third_party_gpus_cuda_compat_cuda_compat_data_o_tmpdir_filewrapper_s0_start
@ third_party_gpus_cuda_compat_cuda_compat_data_o_tmpdir_filewrapper_s0_start
@ third_party_gpus_cuda_compat_cuda_compat_data_o_tmpdir_filewrapper_s0_start
@ third_party_gpus_cuda_compat_cuda_compat_data_o_tmpdir_filewrapper_s0_start
@ third_party_gpus_cuda_compat_cuda_compat_data_o_tmpdir_filewrapper_s0_start
@ xla/stream_executor/cuda/cuda_driver.cc:1358 stream_executor::gpu::GpuDriver::SynchronizeStream()
@ xla/stream_executor/cuda/cuda_executor.cc:528 stream_executor::gpu::CudaExecutor::BlockHostUntilDone()
@ xla/stream_executor/stream_common.cc:148 stream_executor::StreamCommon::BlockHostUntilDone()
@ xla/service/gpu/runtime/nccl_all_to_all_thunk.cc:313 xla::gpu::RunMemCpyAllToAll()
@ xla/service/gpu/runtime/nccl_all_to_all_thunk.cc:189 xla::gpu::NcclAllToAllStartThunk::RunNcclCollective()
@ xla/service/gpu/runtime/nccl_collective_thunk.cc:447 xla::gpu::NcclCollectiveThunk::ExecuteOnStream()
@ xla/service/gpu/runtime/sequential_thunk.cc:81 xla::gpu::SequentialThunk::ExecuteOnStream()
@ xla/service/gpu/gpu_executable.cc:481 xla::gpu::(anonymous namespace)::ExecuteThunks()
@ xla/service/gpu/gpu_executable.cc:1011 xla::gpu::GpuExecutable::ExecuteAsyncOnStreamImpl()
@ xla/service/gpu/gpu_executable.cc:798 xla::gpu::GpuExecutable::ExecuteAsyncOnStream()
@ xla/service/executable.cc:84 xla::Executable::ExecuteOnStream()
@ xla/service/hlo_runner.cc:607 xla::HloRunner::ExecuteReplicated()::$_0::operator()()::{lambda()#2}::operator()()
@ Thread::ThreadBody()

We are waiting this thread from the other one that triggers the watchdog code.
@ xla/service/rendezvous.cc:83 xla::internal::AwaitAndLogIfStuck()
@ ./xla/service/rendezvous.h:307 xla::RendezvousSingle<>()
@ ./xla/service/rendezvous.h:336 xla::RendezvousSingle<>()
@ ./xla/service/rendezvous.h:361 xla::RendezvousSingle<>()
@ xla/service/gpu/runtime/nccl_collective_thunk.cc:486 xla::gpu::NcclCollectiveThunk::ExecuteOnStream()
@ xla/service/gpu/runtime/sequential_thunk.cc:81 xla::gpu::SequentialThunk::ExecuteOnStream()
@ xla/service/gpu/gpu_executable.cc:481 xla::gpu::(anonymous namespace)::ExecuteThunks()
@ xla/service/gpu/gpu_executable.cc:1011 xla::gpu::GpuExecutable::ExecuteAsyncOnStreamImpl()
@ xla/service/gpu/gpu_executable.cc:798 xla::gpu::GpuExecutable::ExecuteAsyncOnStream()
@ xla/service/executable.cc:84 xla::Executable::ExecuteOnStream()
@ xla/service/hlo_runner.cc:607 xla::HloRunner::ExecuteReplicated()::$_0::operator()()::{lambda()#2}::operator()()

loislo · 2024-10-02T12:58:27Z

@terryysun it is failing on the machine with p100 gpus

terryysun · 2024-10-02T18:28:20Z

@terryysun it is failing on the machine with p100 gpus

thanks! will file a followup pr to resolve P100 failures and make the flag default, right now it's an optional flag so should not affect anything. it's verified on H100s. working on reproducing the P100 failure on our side

NaiyerRizz requested review from cheshire and thomasjoerg July 22, 2024 05:56

NaiyerRizz self-assigned this Jul 22, 2024

thomasjoerg requested review from frgossen and removed request for thomasjoerg July 22, 2024 08:29

frgossen suggested changes Jul 22, 2024

View reviewed changes

terryysun requested a review from frgossen July 23, 2024 21:31

frgossen suggested changes Jul 29, 2024

View reviewed changes

terryysun force-pushed the terryysun/all2all_memcpyp2p branch from 37c3bb7 to 6e2fcf3 Compare August 5, 2024 22:53

sgerrard requested a review from frgossen August 12, 2024 22:37

terryysun marked this pull request as draft August 13, 2024 17:25

terryysun force-pushed the terryysun/all2all_memcpyp2p branch from 0c384cd to 854d9a3 Compare September 6, 2024 22:36

memcpyp2p for local a2a

38720c7

terryysun force-pushed the terryysun/all2all_memcpyp2p branch from 854d9a3 to 90018f4 Compare September 17, 2024 04:20

use nccl to pass recv ptrs

90018f4

terryysun force-pushed the terryysun/all2all_memcpyp2p branch from 38bd147 to aaadc20 Compare September 18, 2024 00:56

terryysun marked this pull request as ready for review September 18, 2024 01:06

terryysun force-pushed the terryysun/all2all_memcpyp2p branch from aaadc20 to f9b75b0 Compare September 18, 2024 01:15

refactor and cleanup

f9b75b0

frgossen approved these changes Sep 24, 2024

View reviewed changes

copybara-service bot mentioned this pull request Sep 24, 2024

PR #15144: [NVIDIA GPU] Use memcpy for intra-node all-to-all #17558

Open

copybara-service bot mentioned this pull request Sep 24, 2024

PR #15144: [NVIDIA GPU] Use memcpy for intra-node all-to-all tensorflow/tensorflow#76409

Draft

copybara-service bot mentioned this pull request Oct 1, 2024

PR #17636: [NVIDIA GPU] Enhance concurrency handling in cross-rank address sharing #17799

Closed

copybara-service bot mentioned this pull request Oct 1, 2024

PR #17636: [NVIDIA GPU] Enhance concurrency handling in cross-rank address sharing tensorflow/tensorflow#76864

Merged

loislo reopened this Oct 2, 2024

terryysun closed this Oct 10, 2024

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[NVIDIA GPU] Use memcpy for intra-node all-to-all #15144

[NVIDIA GPU] Use memcpy for intra-node all-to-all #15144

terryysun commented Jul 19, 2024

terryysun commented Jul 19, 2024 •

edited

Loading

frgossen left a comment

terryysun commented Jul 23, 2024

frgossen left a comment

frgossen Jul 29, 2024

terryysun Sep 18, 2024

frgossen Jul 29, 2024

terryysun Sep 18, 2024

frgossen Jul 29, 2024

terryysun Sep 18, 2024

frgossen commented Sep 9, 2024

terryysun commented Sep 18, 2024

frgossen left a comment

loislo commented Oct 2, 2024

loislo commented Oct 2, 2024

terryysun commented Oct 2, 2024 •

edited

Loading

[NVIDIA GPU] Use memcpy for intra-node all-to-all #15144

[NVIDIA GPU] Use memcpy for intra-node all-to-all #15144

Conversation

terryysun commented Jul 19, 2024

terryysun commented Jul 19, 2024 • edited Loading

frgossen left a comment

Choose a reason for hiding this comment

terryysun commented Jul 23, 2024

frgossen left a comment

Choose a reason for hiding this comment

frgossen Jul 29, 2024

Choose a reason for hiding this comment

terryysun Sep 18, 2024

Choose a reason for hiding this comment

frgossen Jul 29, 2024

Choose a reason for hiding this comment

terryysun Sep 18, 2024

Choose a reason for hiding this comment

frgossen Jul 29, 2024

Choose a reason for hiding this comment

terryysun Sep 18, 2024

Choose a reason for hiding this comment

frgossen commented Sep 9, 2024

terryysun commented Sep 18, 2024

frgossen left a comment

Choose a reason for hiding this comment

loislo commented Oct 2, 2024

loislo commented Oct 2, 2024

terryysun commented Oct 2, 2024 • edited Loading

terryysun commented Jul 19, 2024 •

edited

Loading

terryysun commented Oct 2, 2024 •

edited

Loading