PR tensorflow#36901: [xla:gpu] Add support for NCCL scalable initialization

Imported from GitHub PR openxla/xla#36901

Remove `root_device` from `GpuCliqueKey`: root device is not a property of the gpu clique key, but has meaning only during clique initialization.

1. Change CliqueId callback type to return `CliqueIds`
2. Update NCCL Id store to set and get multiple clique ids for all root devices
3. Use NCCL scalable initialization API when XLA generates more that 1 clique id
4. Improve logging and add timing to critical places that tend to deadlock or get stuck
5. Improve logging by adding HumanReadable Devices and Processes APIs for printing large number of devices and processes
Copybara import of the project:

--
7c239c26b18d3a4c86f83dd9ea8b5b2b5d7ccbee by Eugene Zhulenev <ezhulenev@openxla.org>:

[xla:gpu] Add support for NCCL scalable initialization

Merging this change closes tensorflow#36901

PiperOrigin-RevId: 862348585

Author: ezhulenev

third_party/xla/xla/backends/gpu/collectives/BUILD

@@ -397,7 +397,6 @@ cc_library(
         "//xla/pjrt/distributed:key_value_store_interface",
         "//xla/runtime:device_id",
         "//xla/runtime:process_id",
-        "//xla/service/gpu:gpu_executable_run_options",
         "//xla/stream_executor:stream_executor_h",
         "//xla/stream_executor/cuda:nccl_memory_allocator",  # buildcleaner: keep (static registration)
         "//xla/tsl/cuda:nccl",
@@ -408,6 +407,7 @@ cc_library(
         "@com_google_absl//absl/algorithm:container",
         "@com_google_absl//absl/base",
         "@com_google_absl//absl/base:core_headers",
+        "@com_google_absl//absl/container:btree",
         "@com_google_absl//absl/container:flat_hash_map",
         "@com_google_absl//absl/status",
         "@com_google_absl//absl/status:statusor",

third_party/xla/xla/backends/gpu/collectives/gpu_clique_key.cc

@@ -64,12 +64,11 @@ CollectiveStreamId GetCollectiveStreamId(bool is_async,
 GpuCliqueKey::GpuCliqueKey(
     std::vector<GlobalDeviceId> devices, int64_t num_local_participants,
     bool is_p2p, std::vector<std::vector<GlobalDeviceId>> participant_groups,
-    GlobalDeviceId root_device, std::vector<IncarnationId> incarnations)
+    std::vector<IncarnationId> incarnations)
     : CliqueKey(std::move(devices)),
       num_local_participants_(num_local_participants),
       is_p2p_(is_p2p),
       participant_groups_(std::move(participant_groups)),
-      root_device_(root_device),
       incarnations_(std::move(incarnations)) {
   for (std::vector<GlobalDeviceId>& group : participant_groups_) {
     absl::c_sort(group);
@@ -86,8 +85,6 @@ GpuCliqueKey::GpuCliqueKey(
 
 bool GpuCliqueKey::is_p2p() const { return is_p2p_; }
 
-GlobalDeviceId GpuCliqueKey::root_device() const { return root_device_; }
-
 std::vector<std::vector<GlobalDeviceId>> GpuCliqueKey::ParticipantGroups()
     const {
   return participant_groups_;
@@ -105,35 +102,27 @@ bool GpuCliqueKey::IsSubsetOf(const CliqueKey& other) const {
          });
 }
 
-std::vector<GpuCliqueKey> GpuCliqueKey::GetSubKeys(int64_t nroots) const {
-  const auto& devs = devices();
-  int64_t nranks = devs.size();
-  CHECK(nroots <= nranks);
-  int64_t rank_per_root = nranks / nroots;
-  int64_t rank_rem = nranks % nroots;
-  std::vector<GpuCliqueKey> subkeys;
+std::vector<GlobalDeviceId> GpuCliqueKey::GetRootDevices(int64_t nroots) const {
+  int64_t nranks = devices().size();
+  CHECK_LE(nroots, nranks) << "Can't select more root devices than available";
+
+  // If nroots evenly divides nranks, then every root is assigned to
+  // ranks_per_root devices. If nroots doesn't divide nranks, then we increase
+  // the size of the first ranks_rem roots by 1.
+  int64_t ranks_per_root = nranks / nroots;
+  int64_t ranks_rem = nranks % nroots;
+
+  std::vector<GlobalDeviceId> roots;
+  roots.reserve(nroots);
   for (int64_t i = 0; i < nroots; ++i) {
-    GpuCliqueKey subkey(*this);
-    if (i < rank_rem) {
-      subkey.root_device_ = devs[i * (rank_per_root + 1)];
+    if (i < ranks_rem) {
+      roots.push_back(devices()[i * (ranks_per_root + 1)]);
     } else {
-      subkey.root_device_ =
-          devs[rank_rem * (rank_per_root + 1) + (i - rank_rem) * rank_per_root];
+      roots.push_back(devices()[ranks_rem * (ranks_per_root + 1) +
+                                (i - ranks_rem) * ranks_per_root]);
     }
-    subkeys.push_back(subkey);
-  }
-  return subkeys;
-}
-
-// StrJoin for devices that shortens long list of devices for readbility.
-static std::string StrJoinDevices(absl::Span<const GlobalDeviceId> devices,
-                                  absl::string_view separator = ",",
-                                  size_t first = 10, size_t last = 4) {
-  if (devices.size() > first + last) {
-    return absl::StrCat(absl::StrJoin(devices.first(first), separator), "...",
-                        absl::StrJoin(devices.last(last), separator));
   }
-  return absl::StrJoin(devices, separator);
+  return roots;
 }
 
 std::string GpuCliqueKey::ToString() const {
@@ -142,15 +131,15 @@ std::string GpuCliqueKey::ToString() const {
     std::vector<std::string> values;
     values.reserve(participant_groups_.size());
     for (const auto& group : participant_groups_) {
-      values.push_back(absl::StrFormat("[%s]", StrJoinDevices(group)));
+      values.push_back(absl::StrFormat("[%s]", HumanReadableDevices(group)));
     }
     group_string = absl::StrFormat("; groups=[%s]", absl::StrJoin(values, ","));
   }
   return absl::StrFormat(
-      "devices=%d:[%s]; is_p2p=%v%s; root=%lld; local_participants=%lld; "
+      "devices=%d:[%s]; is_p2p=%v%s; local_participants=%lld; "
       "incarnations=[%s]",
-      devices().size(), StrJoinDevices(devices()), is_p2p_, group_string,
-      root_device_.value(), num_local_participants_,
+      devices().size(), HumanReadableDevices(devices()), is_p2p_, group_string,
+      num_local_participants_,
       absl::StrJoin(incarnations_, ", ",
                     [](std::string* out, IncarnationId id) {
                       absl::StrAppend(out, id.value());
@@ -159,19 +148,19 @@ std::string GpuCliqueKey::ToString() const {
 
 void GpuCliqueKey::HashValue(absl::HashState state) const {
   absl::HashState::combine(std::move(state), devices(), participant_groups_,
-                           root_device_, incarnations_);
+                           incarnations_);
 }
 
 bool operator==(const GpuCliqueKey& a, const GpuCliqueKey& b) {
   return a.devices() == b.devices() &&
          a.participant_groups_ == b.participant_groups_ &&
          a.num_local_participants_ == b.num_local_participants_ &&
-         a.root_device_ == b.root_device_ && a.incarnations_ == b.incarnations_;
+         a.incarnations_ == b.incarnations_;
 }
 
 // Constructs a tuple from the clique key for comparison purposes.
 static auto CmpKey(const GpuCliqueKey& key) {
-  return std::make_tuple(key.devices().size(), key.devices(), key.root_device(),
+  return std::make_tuple(key.devices().size(), key.devices(),
                          key.num_local_participants(), key.incarnations());
 }
 

third_party/xla/xla/backends/gpu/collectives/gpu_clique_key.h

@@ -51,7 +51,6 @@ class GpuCliqueKey : public CliqueKey {
       std::vector<GlobalDeviceId> devices, int64_t num_local_participants,
       bool is_p2p = false,
       std::vector<std::vector<GlobalDeviceId>> participant_groups = {},
-      GlobalDeviceId root_device = GlobalDeviceId(-1),
       std::vector<IncarnationId> incarnations = {});
 
   GpuCliqueKey(const GpuCliqueKey&) = default;
@@ -64,21 +63,19 @@ class GpuCliqueKey : public CliqueKey {
 
   std::vector<std::vector<GlobalDeviceId>> ParticipantGroups() const;
 
-  // Device generating the unique id for this key
-  GlobalDeviceId root_device() const;
-
   // Returns true if this clique is a subset of `other`: both cliques have the
   // same `stream_id` and all clique devices are part of `other` clique.
   bool IsSubsetOf(const CliqueKey& other) const final;
 
   // Returns true if this clique will be used with p2p communicators.
   bool is_p2p() const;
 
-  // For multi-root initialization, generate `nroots` copies (subkeys) of the
-  // key each with a different root device. Root devices are distributed evenly
-  // across the ranks. The subkeys are used to exchange the CliqueIds during
-  // clique initialization.
-  std::vector<GpuCliqueKey> GetSubKeys(int64_t nroots) const;
+  // Returns root devices that are responsible for bootstrapping the GPU clique
+  // during initialization. Root devices are distributed evenly across all ranks
+  // in the clique. XLA processes owning the root devices are responsible for
+  // generating clique id and exchanging it with other ranks that share the same
+  // root device (this is done via shared KV store).
+  std::vector<GlobalDeviceId> GetRootDevices(int64_t nroots) const;
 
   // The number of participant devices that are local to the current process (in
   // multi-host environments this likely to be all devices on the same host).
@@ -124,8 +121,6 @@ class GpuCliqueKey : public CliqueKey {
   // situations
   std::vector<std::vector<GlobalDeviceId>> participant_groups_;
 
-  GlobalDeviceId root_device_;
-
   std::vector<IncarnationId> incarnations_;
 };
 

third_party/xla/xla/backends/gpu/collectives/gpu_clique_key_test.cc

@@ -40,8 +40,7 @@ static GpuCliqueKey GetBaseCliqueKey() {
                       /*participant_groups=*/
                       std::vector<std::vector<GlobalDeviceId>>{
                           {GlobalDeviceId(0), GlobalDeviceId(1)},
-                          {GlobalDeviceId(2), GlobalDeviceId(3)}},
-                      /*root_device=*/GlobalDeviceId(0));
+                          {GlobalDeviceId(2), GlobalDeviceId(3)}});
 }
 
 TEST(GpuCliqueKeyTest, IsSubsetOf) {
@@ -60,29 +59,36 @@ TEST(GpuCliqueKeyTest, IsSubsetOf) {
   EXPECT_FALSE(key0.IsSubsetOf(key3));
 }
 
-TEST(GpuCliqueKeyTest, GetSubKeys) {
+TEST(GpuCliqueKeyTest, GetRootDevices) {
   GlobalDeviceId id0 = GlobalDeviceId(0);
   GlobalDeviceId id1 = GlobalDeviceId(1);
   GlobalDeviceId id2 = GlobalDeviceId(2);
   GlobalDeviceId id3 = GlobalDeviceId(3);
 
-  GpuCliqueKey key({id0, id1, id2, id3}, /*num_local_participants=*/1, true);
-  std::array<int64_t, 4> nroots{1, 2, 3, 4};
-  std::vector<std::vector<GlobalDeviceId>> exp_root_devs{
-      {id0}, {id0, id2}, {id0, id2, id3}, {id0, id1, id2, id3}};
-  for (int ridx = 0; ridx < nroots.size(); ++ridx) {
-    int64_t n = nroots[ridx];
-    const auto& subkeys = key.GetSubKeys(n);
-    EXPECT_EQ(subkeys.size(), exp_root_devs[ridx].size());
-    for (int kidx = 0; kidx < subkeys.size(); ++kidx) {
-      GpuCliqueKey exp_subkey(
-          /*devices=*/{id0, id1, id2, id3},
-          /*num_local_participants=*/1,
-          /*is_p2p=*/true,
-          /*participant_groups=*/{},
-          /*root_device=*/exp_root_devs[ridx][kidx]);
-      EXPECT_EQ(subkeys[kidx], exp_subkey);
-    }
+  GpuCliqueKey key({id0, id1, id2, id3}, /*num_local_participants=*/1);
+
+  {
+    std::vector<GlobalDeviceId> roots = key.GetRootDevices(1);
+    std::vector<GlobalDeviceId> expected = {id0};
+    ASSERT_EQ(roots, expected);
+  }
+
+  {
+    std::vector<GlobalDeviceId> roots = key.GetRootDevices(2);
+    std::vector<GlobalDeviceId> expected = {id0, id2};
+    ASSERT_EQ(roots, expected);
+  }
+
+  {
+    std::vector<GlobalDeviceId> roots = key.GetRootDevices(3);
+    std::vector<GlobalDeviceId> expected = {id0, id2, id3};
+    ASSERT_EQ(roots, expected);
+  }
+
+  {
+    std::vector<GlobalDeviceId> roots = key.GetRootDevices(4);
+    std::vector<GlobalDeviceId> expected = {id0, id1, id2, id3};
+    ASSERT_EQ(roots, expected);
   }
 }
 
@@ -179,7 +185,7 @@ TEST(GpuCliqueKeyGettersTest, IsP2P) {
 
 TEST(GpuCliqueKeyGetterTest, ToString) {
   EXPECT_EQ(GetBaseCliqueKey().ToString(),
-            "devices=2:[0,1]; is_p2p=false; groups=[[0,1],[2,3]]; root=0; "
+            "devices=2:[0,1]; is_p2p=false; groups=[[0,1],[2,3]]; "
             "local_participants=2; incarnations=[]");
 }
 
@@ -194,7 +200,7 @@ TEST(GpuCliqueKeyGetterTest, ToStringManyDevices) {
 
   EXPECT_EQ(key.ToString(),
             "devices=100:[0,1,2,3,4,5,6,7,8,9...96,97,98,99]; is_p2p=false; "
-            "root=-1; local_participants=100; incarnations=[]");
+            "local_participants=100; incarnations=[]");
 }
 
 TEST(GpuCliqueIdGettersTest, Data) {

third_party/xla/xla/backends/gpu/collectives/gpu_cliques.cc

@@ -25,7 +25,6 @@ limitations under the License.
 #include <vector>
 
 #include "absl/algorithm/container.h"
-#include "absl/base/const_init.h"
 #include "absl/base/no_destructor.h"
 #include "absl/base/thread_annotations.h"
 #include "absl/container/btree_map.h"
@@ -322,30 +321,13 @@ InitializeGpuClique(GpuCollectives* collectives, se::StreamExecutor* device,
     bool synchronized;
   };
 
-  // Check how many roots are needed to initialize the GpuClique
-  static const int64_t nccl_init_rank_per_root_ratio =
-      xla::GetDebugOptionsFromFlags()
-          .xla_gpu_nccl_init_max_rank_per_root_ratio();
-  int64_t nranks = clique_key.num_devices();
-  int64_t nroots = nccl_init_rank_per_root_ratio != 0
-                       ? CeilOfRatio(nranks, nccl_init_rank_per_root_ratio)
-                       : 1;
-
   // Initializes a GpuClique for given device ranks and returns a lock that
   // gives access to clique communicators.
   auto initialize = [&](absl::Span<const RendezvousArg* const> args)
       -> absl::StatusOr<LockableGpuClique::Lock> {
     tsl::profiler::TraceMe trace("InitializeGpuClique");
 
-    CliqueIds clique_ids;
-    std::vector<GpuCliqueKey> subkeys = clique_key.GetSubKeys(nroots);
-    for (const auto& subkey : subkeys) {
-      VLOG(5) << absl::StreamFormat(
-          "Get CliqueId for sub clique key %s; nroots=%lld", subkey.ToString(),
-          nroots);
-      TF_ASSIGN_OR_RETURN(auto clique_id, clique_id_callback(subkey));
-      clique_ids.Add(clique_id);
-    }
+    TF_ASSIGN_OR_RETURN(CliqueIds clique_ids, clique_id_callback(clique_key));
 
     // Check that all ranks successfully synchronized device activity before
     // trying to instantiate GPU communicators.
@@ -375,7 +357,7 @@ InitializeGpuClique(GpuCollectives* collectives, se::StreamExecutor* device,
         "[%s] [ranks=%s] Create GPU communicators for clique %v; "
         "nroots=%lld; fingerprint(id)=%d, peer_access_enabled=%d",
         DeviceOrdinalsToString(ranks), DeviceRanksToString(ranks), clique_key,
-        nroots, clique_ids.fingerprint(), peer_access_enabled);
+        clique_ids.size(), clique_ids.fingerprint(), peer_access_enabled);
 
     ProcessGpuCliques& state = GetProcessGpuCliques();
 
@@ -426,7 +408,7 @@ InitializeGpuClique(GpuCollectives* collectives, se::StreamExecutor* device,
         "[%s] [ranks=%s] Created GPU communicators for clique %v; "
         "nroots=%lld; fingerprint(id)=%d, peer_access_enabled=%d",
         DeviceOrdinalsToString(ranks), DeviceRanksToString(ranks), clique_key,
-        nroots, clique_ids.fingerprint(), peer_access_enabled);
+        clique_ids.size(), clique_ids.fingerprint(), peer_access_enabled);
 
     // Put constructed clique into the per-process state.
     absl::MutexLock lock(state.mu);

third_party/xla/xla/backends/gpu/collectives/gpu_collectives.h

@@ -50,9 +50,9 @@ class GpuCollectives : public Collectives {
   // Returns the default collectives implementation for the given platform.
   static GpuCollectives* Default(absl::string_view platform_name);
 
-  // A callback to get a unique clique id.
+  // A callback to get a unique clique ids.
   using CliqueIdCallback =  // NOLINT
-      std::function<absl::StatusOr<CliqueId>(const CliqueKey&)>;
+      std::function<absl::StatusOr<CliqueIds>(const CliqueKey&)>;
 
   // Topology describes how exactly the current process fits into the collection
   // of distributed processes, and based on this information the underlying

third_party/xla/xla/backends/gpu/collectives/gpu_collectives_test.cc

@@ -15,6 +15,7 @@ limitations under the License.
 
 #include "xla/backends/gpu/collectives/gpu_collectives.h"
 
+#include <cstddef>
 #include <memory>
 #include <utility>
 #include <vector>
@@ -55,14 +56,19 @@ namespace {
 // Creates a pair of communicators for the given executors.
 static absl::StatusOr<std::vector<std::unique_ptr<GpuCommunicator>>>
 CreateCommunicators(se::StreamExecutor* executor0,
-                    se::StreamExecutor* executor1, bool blocking = true) {
+                    se::StreamExecutor* executor1, bool blocking = true,
+                    size_t num_ids = 1) {
   GpuCollectives::Device device0(executor0);
   GpuCollectives::Device device1(executor1);
 
   GpuCollectives* collectives = GpuCollectives::Default("GPU");
 
-  TF_ASSIGN_OR_RETURN(CliqueId clique_id, collectives->CreateUniqueCliqueId());
-  CliqueIds clique_ids(clique_id);
+  CliqueIds clique_ids;
+  for (size_t i = 0; i < num_ids; ++i) {
+    TF_ASSIGN_OR_RETURN(CliqueId clique_id,
+                        collectives->CreateUniqueCliqueId());
+    clique_ids.Add(clique_id);
+  }
 
   GpuCliqueKey clique_key({GlobalDeviceId(0), GlobalDeviceId(1)},
                           /*num_local_participants=*/2);
@@ -86,6 +92,27 @@ CreateCommunicators(se::StreamExecutor* executor0,
   return gpu_comms;
 }
 
+TEST(GpuCollectivesTest, CreateWithMultipleIds) {
+  ASSERT_OK_AND_ASSIGN(se::Platform * platform,
+                       se::PlatformManager::PlatformWithName("CUDA"));
+
+  if (platform->VisibleDeviceCount() < 2) {
+    GTEST_SKIP() << "Test requires at least 2 GPUs";
+  }
+
+  ASSERT_OK_AND_ASSIGN(se::StreamExecutor * executor0,
+                       platform->ExecutorForDevice(0));
+  ASSERT_OK_AND_ASSIGN(se::StreamExecutor * executor1,
+                       platform->ExecutorForDevice(1));
+
+  ASSERT_OK_AND_ASSIGN(
+      auto comms, CreateCommunicators(executor0, executor1, /*blocking=*/true,
+                                      /*num_ids=*/2));
+
+  EXPECT_TRUE(comms[0]->platform_comm().handle);
+  EXPECT_TRUE(comms[1]->platform_comm().handle);
+}
+
 TEST(GpuCollectivesTest, CreateSymmetricMemory) {
   ASSERT_OK_AND_ASSIGN(se::Platform * platform,
                        se::PlatformManager::PlatformWithName("CUDA"));