wehub/tensorflow--tensorflow
1
0
Fork 0
Code Issues Pull Requests Actions 5 Deployments Models Agent Notes Packages Projects Releases Wiki Activity

chore: import upstream snapshot with attribution
cffconvert / validate (push) Has been skipped
Details
License Check / license-check (push) Failing after 2s
Details

Browse Source
This commit is contained in:
wehub-resource-sync
2026-07-13 12:14:16 +08:00
commit 8a852e4b4e
36502 changed files with 9277225 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tensorflow/core/nccl/BUILD
+144
View File
@@ -0,0 +1,144 @@
# Description:
# Wrap NVIDIA (https://github.com/NVIDIA/nccl) NCCL with tensorflow ops.
# APIs are meant to change over time.
load("@local_config_cuda//cuda:build_defs.bzl", "if_cuda")
load("@local_config_rocm//rocm:build_defs.bzl", "if_rocm")
load("//tensorflow:tensorflow.bzl", "if_cuda_or_rocm", "if_nccl", "tf_copts")
load("//tensorflow:tensorflow.default.bzl", "filegroup", "tf_cuda_cc_test")
load(
"//tensorflow/core/platform:build_config_root.bzl",
"tf_cuda_tests_tags",
)
load("//tensorflow/core/platform:rules_cc.bzl", "cc_library")
package(
# copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
default_visibility = ["//tensorflow:__subpackages__"],
licenses = ["notice"],
)
cc_library(
name = "nccl_lib",
srcs = if_cuda_or_rocm([
"nccl_manager.cc",
"nccl_rewrite.cc",
]),
hdrs = if_cuda_or_rocm([
"nccl_manager.h",
]),
copts = tf_copts(),
deps = if_cuda([
"@xla//xla/tsl/cuda:nccl",
"//tensorflow/core/platform:blocking_counter",
"//tensorflow/core/platform:unbounded_work_queue",
"//tensorflow/core/framework:tensor_proto_cc",
]) + if_rocm([
"@local_config_rocm//rocm:rccl",
"//tensorflow/core:gpu_runtime",
]) + if_cuda_or_rocm([
"@com_google_absl//absl/base",
"@com_google_absl//absl/container:flat_hash_map",
"//tensorflow/core:core_cpu",
"//tensorflow/core:framework",
"//tensorflow/core:gpu_headers_lib",
"//tensorflow/core:lib",
"//tensorflow/core/platform:stream_executor",
"//tensorflow/core/profiler/lib:traceme",
"//tensorflow/core/profiler/lib:connected_traceme",
"//tensorflow/core/profiler/lib:annotated_traceme",
]) + [
"//tensorflow/core/framework:types_proto_cc",
"@com_google_absl//absl/status",
"@com_google_absl//absl/strings",
"@com_google_absl//absl/strings:string_view",
],
alwayslink = 1,
)
tf_cuda_cc_test(
name = "nccl_manager_test",
size = "medium",
srcs = ["nccl_manager_test.cc"],
tags = tf_cuda_tests_tags() + [
"noguitar", # TODO(b/176867216): Flaky.
"manual",
"multi_gpu",
"no_oss",
# TODO(b/147451637): Replace 'no_rocm' with 'rocm_multi_gpu'.
"no_rocm",
"notap",
],
deps = [
"//tensorflow/core:test",
"//tensorflow/core:test_main",
"//tensorflow/core:testlib",
"//tensorflow/core/framework:types_proto_cc",
"@com_google_absl//absl/log",
"@com_google_absl//absl/log:check",
"@com_google_absl//absl/status",
"@com_google_absl//absl/strings",
"@com_google_googletest//:gtest",
"@xla//xla/tsl/protobuf:error_codes_proto_impl_cc",
] + if_cuda_or_rocm([
":nccl_lib",
]) + if_cuda([
"@xla//xla/tsl/cuda:nccl",
"//tensorflow/core:cuda",
]) + if_rocm([
"@local_config_rocm//rocm:rccl",
"//tensorflow/core/common_runtime/gpu:rocm",
]),
)
cc_library(
name = "collective_communicator",
srcs = ["collective_communicator.cc"],
hdrs = ["collective_communicator.h"],
copts = tf_copts() + if_nccl(["-DTENSORFLOW_USE_NCCL=1"]),
visibility = [
"//learning/brain/runtime:__subpackages__",
"//tensorflow:__subpackages__",
],
deps =
[
":loose_headers",
"//tensorflow/core:framework",
"//tensorflow/core/protobuf:for_core_protos_cc",
] + if_nccl([
":nccl_lib",
"@com_google_absl//absl/memory",
"//tensorflow/core/profiler/lib:traceme",
]) + if_cuda([
"//tensorflow/core/platform:tracing",
]),
)
# For a more maintainable build this target should not exist and the headers
# should be split into the existing cc_library targets, but this change was
# automatically done so that we can remove long standing issues and complexity
# in the build system. It's up to the OWNERS of this package to get rid of it or
# not. The use of the textual_hdrs attribute is discouraged, use hdrs instead.
# Here it is used to avoid header parsing errors in packages where the feature
# parse_headers was enabled since loose headers were not being parsed. See
# go/loose-lsc-one-target-approach for more details.
cc_library(
name = "loose_headers",
tags = ["avoid_dep"],
textual_hdrs = ["nccl_manager.h"],
visibility = ["//visibility:private"],
)
filegroup(
name = "mobile_srcs",
srcs = [
"collective_communicator.cc",
],
)
filegroup(
name = "mobile_hdrs",
srcs = [
"collective_communicator.h",
],
)
tensorflow/core/nccl/collective_communicator.cc
+236
View File
@@ -0,0 +1,236 @@
/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include "tensorflow/core/nccl/collective_communicator.h"
#include "tensorflow/core/framework/collective.h"
#include "tensorflow/core/protobuf/config.pb.h"
#if TENSORFLOW_USE_NCCL && (GOOGLE_CUDA || TENSORFLOW_USE_ROCM)
#include "absl/memory/memory.h"
#include "tensorflow/core/nccl/nccl_manager.h"
#include "tensorflow/core/profiler/lib/traceme.h"
namespace tensorflow {
class NcclCommunicator : public NcclCommunicatorInterface {
public:
std::string GenerateCommunicatorKey() override {
return nccl_manager_.GenerateCommunicatorKey();
}
void Enqueue(std::shared_ptr<CollectiveContext> col_ctx,
StatusCallback done) override;
void StartAbort(const absl::Status& s) override;
private:
NcclManager nccl_manager_;
};
namespace {
absl::Status ReductionOp(const std::string& merge_op,
ncclRedOp_t* reduction_op) {
if (merge_op == "Add") {
*reduction_op = ncclSum;
return absl::OkStatus();
} else if (merge_op == "Mul") {
*reduction_op = ncclProd;
return absl::OkStatus();
} else if (merge_op == "Maximum") {
*reduction_op = ncclMax;
return absl::OkStatus();
} else if (merge_op == "Minimum") {
*reduction_op = ncclMin;
return absl::OkStatus();
} else {
return absl::InternalError(absl::StrCat(
"Expected merge_op to be in [Add, Mul, Maximum, Minimum], found ",
merge_op));
}
}
std::string NcclCollectiveKey(const std::string& exec_key, int step_id) {
return absl::StrCat(exec_key, ":", step_id);
}
} // namespace
std::unique_ptr<NcclCommunicatorInterface> MaybeCreateNcclCommunicator(
const ConfigProto& config) {
// Skip creating a NcclCommunicator if there are 0 GPUs configured.
const auto& device_count = config.device_count();
auto item = device_count.find("GPU");
if (item != device_count.end() && item->second == 0) {
return nullptr;
}
return absl::make_unique<NcclCommunicator>();
}
void NcclCommunicator::Enqueue(std::shared_ptr<CollectiveContext> col_ctx,
StatusCallback done) {
const CollectiveParams* col_params = col_ctx->col_params.get();
const int num_global_devices = col_params->group.group_size;
const int num_local_devices = col_params->group.num_devices_per_task.at(
col_params->group.members[col_params->default_rank].task);
const std::string nccl_collective_key =
NcclCollectiveKey(col_ctx->exec_key, col_ctx->step_id);
auto* compute_stream = col_ctx->op_ctx->op_device_context()->stream();
auto* gpu_info =
col_ctx->op_ctx->device()->tensorflow_accelerator_device_info();
auto participant = absl::make_unique<NcclManager::Participant>(
compute_stream->parent(), compute_stream, gpu_info, col_ctx->input,
col_ctx->output, col_ctx->col_params->default_rank,
/*done_callback=*/nullptr);
CancellationManager* cancel_mgr = col_ctx->op_ctx->cancellation_manager();
if (cancel_mgr == nullptr) {
participant->done_callback = std::move(done);
} else {
CancellationToken cancel_token = cancel_mgr->get_cancellation_token();
bool already_cancelled =
!cancel_mgr->RegisterCallback(cancel_token, [this]() {
nccl_manager_.StartAbort(absl::CancelledError("op cancelled"));
nccl_manager_.Reset();
});
if (already_cancelled) {
done(absl::CancelledError("op cancelled"));
return;
}
participant->done_callback = [cancel_mgr, cancel_token,
done =
std::move(done)](const absl::Status& s) {
// Do not block on deregistration since this can be invoked by
// NcclManager::StartAbort() in the cancellation callback.
cancel_mgr->TryDeregisterCallback(cancel_token);
done(s);
};
}
NcclManager::Context context(
nccl_collective_key, num_local_devices, num_global_devices,
col_params->group.runtime_details.communicator_key,
col_params->source_rank);
VLOG(1) << "NcclCommunicator::Enqueue type " << col_params->instance.type
<< " num_tasks " << col_params->group.num_tasks << " current task "
<< col_params->group.members[col_params->default_rank].task
<< " num local devices " << num_local_devices
<< " num global devices " << num_global_devices << " device "
<< col_ctx->device_name << " instance "
<< col_params->instance.instance_key;
// `AddTo*` performs consistency checks for the NCCL call and enqueues the
// `Participant` struct locally. When all local participants with this
// `nccl_collective_key` have called `AddToAllReduce` and
// `SignalMultiNodeReady`, all devices at this worker are ready to process
// this NCCL op.
//
// The `NcclManager` uses a dedicated CUDA stream for NCCL kernels. At this
// point, it synchronizes the NCCL stream with the compute stream, and then
// enqueues the NCCL kernel on the NCCL stream.
switch (col_params->instance.type) {
case REDUCTION_COLLECTIVE: {
ncclRedOp_t reduction_op;
absl::Status s =
ReductionOp(col_params->merge_op->type_string(), &reduction_op);
if (!s.ok()) {
participant->done_callback(s);
return;
}
nccl_manager_.AddToAllReduce(std::move(participant), context,
reduction_op);
break;
}
case GATHER_COLLECTIVE: {
nccl_manager_.AddToAllGather(std::move(participant), context);
break;
}
case BROADCAST_COLLECTIVE: {
if (col_params->is_source) {
nccl_manager_.AddBroadcastSend(std::move(participant), context);
} else {
nccl_manager_.AddBroadcastRecv(std::move(participant), context);
}
break;
}
case REDUCE_SCATTER_COLLECTIVE: {
ncclRedOp_t reduction_op;
absl::Status s =
ReductionOp(col_params->merge_op->type_string(), &reduction_op);
if (!s.ok()) {
participant->done_callback(s);
return;
}
nccl_manager_.AddToReduceScatter(std::move(participant), context,
reduction_op);
break;
}
case ALL_TO_ALL_COLLECTIVE: {
nccl_manager_.AddToAllToAll(std::move(participant), context);
break;
}
default: {
participant->done_callback(absl::InternalError(absl::StrCat(
"Unexpected CollectiveType ", col_params->instance.type)));
return;
}
}
// NOTE(ayushd): We need to synchronize NCCL launches across nodes to prevent
// deadlocks. In the current implementation, we define a deterministic
// sequential launch order between potentially concurrent collective instances
// by introducing control information during static graph analysis in
// graph/collective_order.cc. This can be either in the form of explicit
// control edges or via `wait_for` attribute on the collective op.
//
// The other end of the design spectrum would have a distinguished node
// dynamically signal the next collective to launch to all other participants.
// This has higher degree of runtime coordination, but it may be able to
// achieve better performance if the (arbitrary) static execution order
// assigned in the first approach turns out to not be good from a scheduling
// perspective. e.g. consider a graph in which c1, c2, and c3 are three
// concurrent collective instances, and the static ordering assigns c1 -> c2
// -> c3. In practice, it could turn out that c3 is always ready to execute
// before c1 or c2.
{
// `WaitForDependencies` may block if the collective instances on which this
// op depends have not yet launched. When this function returns, this op is
// ready to go.
tsl::profiler::TraceMe activity("WaitForDependencies",
tsl::profiler::TraceMeLevel::kInfo);
col_ctx->col_exec->WaitForDependencies(*col_params);
nccl_manager_.SignalMultiNodeReady(nccl_collective_key);
}
{
// When all devices at this worker have called `SignalMultiNodeReady`, the
// `NcclManager` will enqueue the NCCL kernel on the NCCL stream. Thus the
// implementation of `UnblockDependencies` keeps track of the number of
// devices that have launched.
tsl::profiler::TraceMe activity("Schedule",
tsl::profiler::TraceMeLevel::kInfo);
col_ctx->col_exec->UnblockDependencies(*col_params);
}
}
void NcclCommunicator::StartAbort(const absl::Status& s) {
nccl_manager_.StartAbort(s);
}
} // namespace tensorflow
#else
namespace tensorflow {
std::unique_ptr<NcclCommunicatorInterface> MaybeCreateNcclCommunicator(
const ConfigProto& config) {
return nullptr;
}
} // namespace tensorflow
#endif // TENSORFLOW_USE_NCCL && (GOOGLE_CUDA || TENSORFLOW_USE_ROCM)
tensorflow/core/nccl/collective_communicator.h
+30
View File
@@ -0,0 +1,30 @@
/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#ifndef TENSORFLOW_CORE_NCCL_COLLECTIVE_COMMUNICATOR_H_
#define TENSORFLOW_CORE_NCCL_COLLECTIVE_COMMUNICATOR_H_
#include "tensorflow/core/framework/collective.h"
#include "tensorflow/core/protobuf/config.pb.h"
namespace tensorflow {
// Creates a NcclCommunicator if built with NCCL support (unless configured to
// use no GPU devices), otherwise it returns nullptr.
std::unique_ptr<NcclCommunicatorInterface> MaybeCreateNcclCommunicator(
const ConfigProto& config);
} // namespace tensorflow
#endif // TENSORFLOW_CORE_NCCL_COLLECTIVE_COMMUNICATOR_H_
tensorflow/core/nccl/nccl_manager.cc
+997
View File
@@ -0,0 +1,997 @@
/* Copyright 2016 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include "tensorflow/core/nccl/nccl_manager.h"
#include <utility>
#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
#include "absl/base/call_once.h"
#include "tensorflow/core/framework/types.h"
#include "tensorflow/core/lib/core/refcount.h"
#include "tensorflow/core/lib/core/threadpool.h"
#include "tensorflow/core/platform/blocking_counter.h"
#include "tensorflow/core/platform/env.h"
#include "tensorflow/core/platform/unbounded_work_queue.h"
#include "tensorflow/core/profiler/lib/annotated_traceme.h"
#include "tensorflow/core/profiler/lib/connected_traceme.h"
#include "tensorflow/core/profiler/lib/traceme.h"
#if TENSORFLOW_USE_ROCM
#include "tensorflow/core/platform/rocm.h"
#endif
namespace tensorflow {
#if TENSORFLOW_USE_ROCM
// Local hipify of cuda symbols
#define cudaError_t hipError_t
#define cudaStream_t hipStream_t
#define cudaGetErrorString hipGetErrorString
#define cudaGetDevice hipGetDevice
#define cudaSetDevice hipSetDevice
#define cudaSuccess hipSuccess
int NcclManager::instance_count = 0;
#endif
#define NCCL_RETURN_IF_ERROR(...) \
do { \
ncclResult_t nccl_status = (__VA_ARGS__); \
if (nccl_status != ncclSuccess) { \
return errors::Internal("NCCL: ", ncclGetErrorString(nccl_status), \
". Set NCCL_DEBUG=WARN for detail."); \
} \
} while (0)
#define CUDA_RETURN_IF_ERROR(...) \
do { \
cudaError_t cuda_status = (__VA_ARGS__); \
if (cuda_status != cudaSuccess) { \
return errors::Internal("CUDA: ", cudaGetErrorString(cuda_status)); \
} \
} while (0)
// Contains data for a single stream used for nccl communication; this includes
// a background thread that calls NcclManager::LoopKernelLaunches.
struct NcclManager::NcclStream : public core::RefCounted {
public:
NcclStream() = default;
~NcclStream() = default;
se::StreamExecutor* executor = nullptr;
// The stream on which to run the nccl collective.
// This is a different stream than the tensorflow compute stream.
#if TENSORFLOW_USE_ROCM
// On ROCm, we borrow the nccl stream from the device context.
se::Stream* stream = nullptr;
#else
std::unique_ptr<se::Stream> stream;
#endif
// `mu` protects access to `pending_launches_`, which is the list of
// collectives ready but whose kernels are yet to be launched. When the
// NcclManager object that owns this NcclStream object is destroyed, it
// signals `cv` to unblock the thread waiting on more collectives.
mutex mu;
condition_variable cv;
// Has (collective, participant_idx) pairs.
std::deque<std::pair<Collective*, int>> pending_launches_ TF_GUARDED_BY(mu);
bool shutdown_requested TF_GUARDED_BY(mu) = false;
};
struct NcclManager::CommunicatorMember {
public:
CommunicatorMember() {}
~CommunicatorMember() {
if (nccl_comm != nullptr) ncclCommDestroy(nccl_comm);
}
ncclComm_t nccl_comm = nullptr;
// Owned by NcclManager::device_to_comm_streams_ and LoopKernelLaunches.
NcclStream* nccl_stream = nullptr;
};
struct NcclManager::Communicator {
public:
explicit Communicator(std::vector<CommunicatorMember> members,
const std::string& key)
: num_devices(members.size()), members(std::move(members)), key(key) {}
const int num_devices;
std::vector<CommunicatorMember> members;
const std::string key;
};
namespace {
static constexpr DataTypeSet kValidDataTypes =
ToSet(DT_HALF) | ToSet(DT_FLOAT) | ToSet(DT_DOUBLE) | ToSet(DT_INT32) |
ToSet(DT_INT64);
ncclDataType_t ToNcclType(DataType t) {
switch (t) {
case DT_HALF:
return ncclHalf;
case DT_FLOAT:
return ncclFloat;
case DT_DOUBLE:
return ncclDouble;
case DT_INT32:
return ncclInt;
case DT_INT64:
return ncclInt64;
default:
return ncclFloat;
}
}
void StringToNcclUniqueId(const std::string& str_id, ncclUniqueId* nccl_id) {
if (str_id.size() == NCCL_UNIQUE_ID_BYTES) {
memcpy(nccl_id->internal, str_id.data(), NCCL_UNIQUE_ID_BYTES);
}
}
} // namespace
// A `Collective` encapsulates state for a collective instance at one node.
// Typically, an instance in TensorFlow context would be defined by a collective
// group and the (step, frame iteration) for that execution.
//
// For each collective instance there will be one `Collective` object per node.
// For example, a NCCL collective that runs on a single node with 4 GPUs would
// have a single `Collective` per step. However, a collective that executes on
// 3 nodes with 4 GPUs each would have a `Collective` per node, each of which is
// tracking the 4 GPUs local to that node.
struct NcclManager::Collective : public core::RefCounted {
Collective(const std::string& collective_key_in, DataType data_type_in,
CollectiveType type_in, ncclRedOp_t reduction_op_in,
int num_local_devices_in, int num_global_devices_in,
const std::string& communicator_key_in)
: collective_key(collective_key_in),
data_type(data_type_in),
type(type_in),
reduction_op(reduction_op_in),
num_local_devices(num_local_devices_in),
num_global_devices(num_global_devices_in),
single_node(num_local_devices_in == num_global_devices_in),
communicator_key(communicator_key_in) {
participants.reserve(num_local_devices_in);
#if TENSORFLOW_USE_ROCM
// On ROCm platform, this allows caller to either use the singleton instance
// or to manage one non-singleton NcclManager instance.
// For example, the nccl_manager_test will use both paradigms in the same
// executable, but not running concurrently (which would hang otherwise).
if (NcclManager::instance_count > 1) {
status = errors::Internal(
"ROCm cannot use multi-node NCCL collectives on a single node");
}
#endif
}
const std::string collective_key; // A unique key for debugging.
const DataType data_type;
const CollectiveType type;
const ncclRedOp_t reduction_op; // applies when <type> is a reduction.
const int num_local_devices; // devices local to this node
const int num_global_devices; // devices across all nodes
const bool single_node; // true if all devices are at one node
const std::string communicator_key;
Communicator* communicator = nullptr;
// All collective participants.
//
// Adding values in this vector is guarded by the mutex of the containing
// NcclManager.
std::vector<std::unique_ptr<Participant>> participants;
// For collective types that have a root (e.g. the root of broadcast is the
// sender), this is the rank of the root.
int root_rank = -1;
// How many participants have been registered so far. The Collective is
// eligible for running with <available_participants> == num_local_devices.
//
// If this is a multi-node collective, we additionally have to synchronize
// across nodes. The caller would need to signal multi node readiness by
// calling NcclManager::SignalMultiNodeReady, which sets `multi_node_ready` to
// true.
//
// Guarded by the mutex of the containing Communicator.
int available_participants = 0;
bool multi_node_ready = false;
// trace_context is used by tracing system to associate collective
// scheduling and execution (cooperative kernel launch), which happen
// on different threads.
uint64_t trace_context = 0;
absl::Status status;
};
NcclManager::NcclManager() {
VLOG(2) << "New NcclManager " << this;
#if TENSORFLOW_USE_ROCM
++instance_count;
#endif
}
NcclManager::~NcclManager() {
VLOG(2) << "~NcclManager " << this;
#if TENSORFLOW_USE_ROCM
--instance_count;
#endif
for (auto& it : device_to_comm_streams_) {
for (NcclStream* nccl_stream : it.second) {
{
mutex_lock l(nccl_stream->mu);
nccl_stream->shutdown_requested = true;
nccl_stream->cv.notify_all();
}
nccl_stream->Unref();
}
}
}
NcclManager* NcclManager::instance() {
static NcclManager* instance = new NcclManager();
#if TENSORFLOW_USE_ROCM
// singleton does not count against total instances
// see comment above in Collective constructor concerning ROCm platform
static absl::once_flag once;
absl::call_once(once, [] { --NcclManager::instance_count; });
#endif
return instance;
}
std::string NcclManager::GenerateCommunicatorKey() {
ncclUniqueId nccl_id;
ncclGetUniqueId(&nccl_id);
return std::string(nccl_id.internal, NCCL_UNIQUE_ID_BYTES);
}
absl::Status NcclManager::GetCommunicator(
NcclManager::Collective* collective,
NcclManager::Communicator** communicator) {
// Sort by device ID, executor, and global rank to make ordering of
// participants deterministic.
std::sort(collective->participants.begin(), collective->participants.end(),
[](const std::unique_ptr<Participant>& a,
const std::unique_ptr<Participant>& b) {
if (a->gpu_device_id != b->gpu_device_id) {
return a->gpu_device_id < b->gpu_device_id;
}
if (a->executor != b->executor) {
return a->executor < b->executor;
}
return a->global_rank < b->global_rank;
});
mutex_lock l(mu_);
if (!status_.ok()) {
return status_;
}
if (collective->communicator_key.empty()) {
// For single-node collectives, when the caller does not specify a
// `communicator_key`, we identify a communicator uniquely by the set of
// devices participating in the collective. For example, if a collective is
// for GPUs 0, 1, and 2 then this will scan to find the communicator for
// GPUs 0, 1, and 2.
//
// Note that each executor identifies a context on one device, so this is
// the same as getting the communicator connecting the devices in the
// collective. A device can be in different communicators as well - for
// example, a communicator for GPUs 0 and 1 is separate from one for GPUs 0,
// 1, and 2.
//
// Since it's expected that a small number of distinct communicators will
// be needed, communicators_ is not garbage collected currently.
//
// Launching of kernels must be serialized so that, given collectives A and
// B, and an order of them (e.g., A before B), then for each comm_stream
// involved, the kernel for A is launched before the kernel for B. This is
// guaranteed currently by a global mutex controlling additions of the
// kernels to per-stream launch queues. The launch queues are processed by
// LoopKernelLaunches.
for (auto& comm : communicators_) {
if (comm->num_devices == collective->num_global_devices) {
int i;
for (i = 0; i < collective->num_local_devices; ++i) {
if (comm->members[i].nccl_stream->executor !=
collective->participants[i]->executor) {
break;
}
}
if (i == collective->num_local_devices) {
*communicator = comm.get();
return absl::OkStatus();
}
}
}
} else {
#if NCCL_MAJOR < 2
return errors::Internal(
"Cannot use multi-node NCCL collectives with NCCL 1.x");
#endif
if (collective->communicator_key.size() != NCCL_UNIQUE_ID_BYTES) {
return absl::InternalError(absl::StrCat(
"Expected communicator_key of size ", NCCL_UNIQUE_ID_BYTES,
" but found size ", collective->communicator_key.size()));
}
// This is an instance of multi-node collective. We have previously
// created a NCCL unique id and shared with all workers. Now we find the
// `Communicator` corresponding to this id.
for (auto& comm : communicators_) {
if (comm->key == collective->communicator_key) {
*communicator = comm.get();
return absl::OkStatus();
}
}
}
auto* env = Env::Default();
std::set<NcclStream*> used_streams;
// Create and initialize a new communicator.
// Note that this is done under the lock; performance is not expected to
// matter as this happens a very small number of times.
std::vector<CommunicatorMember> members(collective->num_local_devices);
std::vector<int> devices(collective->num_local_devices);
for (int i = 0; i < collective->num_local_devices; ++i) {
auto* executor = collective->participants[i]->executor;
// Find a communication stream to use for the device.
auto& streams = device_to_comm_streams_[executor];
NcclStream* nccl_stream = nullptr;
for (const auto& s : streams) {
if (used_streams.insert(s).second) {
nccl_stream = s;
break;
}
}
if (nccl_stream == nullptr) {
nccl_stream = new NcclStream();
nccl_stream->executor = executor;
#if TENSORFLOW_USE_ROCM
nccl_stream->stream = collective->participants[i]->context->nccl_stream();
#else
TF_ASSIGN_OR_RETURN(auto stream, executor->CreateStream());
nccl_stream->stream = std::move(stream);
#endif
streams.emplace_back(nccl_stream);
used_streams.insert(nccl_stream);
nccl_stream->Ref();
env->SchedClosure([this, nccl_stream]() {
LoopKernelLaunches(nccl_stream);
nccl_stream->Unref();
});
}
members[i].nccl_stream = nccl_stream;
devices[i] = collective->participants[i]->gpu_device_id;
}
std::vector<ncclComm_t> nccl_comms(collective->num_local_devices);
VLOG(2) << "Created nccl Communicator with "
<< "num_global_devices = " << collective->num_global_devices
<< " num_local_devices = " << collective->num_local_devices
<< " communicator_key ="
<< absl::StrJoin(
std::vector<int>{collective->communicator_key.begin(),
collective->communicator_key.end()},
" ");
#if NCCL_MAJOR >= 2
// For NCCL 2, we always initialize using ncclCommInitRank guarded by NCCL
// group primitives.
ncclUniqueId nccl_id;
if (collective->single_node) {
NCCL_RETURN_IF_ERROR(ncclGetUniqueId(&nccl_id));
} else {
StringToNcclUniqueId(collective->communicator_key, &nccl_id);
}
int saved_device = 0;
CUDA_RETURN_IF_ERROR(cudaGetDevice(&saved_device));
NCCL_RETURN_IF_ERROR(ncclGroupStart());
for (int i = 0; i < collective->num_local_devices; ++i) {
// Set rank to `participant->global_rank` if provided, else `i`.
const int rank = collective->participants[i]->global_rank >= 0
? collective->participants[i]->global_rank
: i;
CUDA_RETURN_IF_ERROR(cudaSetDevice(devices[i]));
NCCL_RETURN_IF_ERROR(ncclCommInitRank(
nccl_comms.data() + i, collective->num_global_devices, nccl_id, rank));
}
NCCL_RETURN_IF_ERROR(ncclGroupEnd());
CUDA_RETURN_IF_ERROR(cudaSetDevice(saved_device));
#else
// Since NCCL 1 is single node only, we use ncclCommInitAll. We could have
// used ncclCommInitRank with NCCL 1 as well, but then we would have to
// issue each init call from a different thread
// (https://docs.nvidia.com/deeplearning/sdk/nccl-developer-guide/docs/nccl1.html).
NCCL_RETURN_IF_ERROR(ncclCommInitAll(
nccl_comms.data(), collective->num_local_devices, devices.data()));
#endif
for (int i = 0; i < collective->num_local_devices; ++i) {
members[i].nccl_comm = nccl_comms[i];
}
communicators_.emplace_back(
new Communicator(std::move(members), collective->communicator_key));
*communicator = communicators_.back().get();
return absl::OkStatus();
}
void NcclManager::AddToAllReduce(std::unique_ptr<Participant> participant,
const Context& context,
ncclRedOp_t reduction_op) {
AddParticipant(std::move(participant), context, kAllReduce, reduction_op);
}
void NcclManager::AddToAllGather(std::unique_ptr<Participant> participant,
const Context& context) {
AddParticipant(std::move(participant), context, kAllGather,
ncclSum /* unused */);
}
void NcclManager::AddToReduceScatter(std::unique_ptr<Participant> participant,
const Context& context,
ncclRedOp_t reduction_op) {
AddParticipant(std::move(participant), context, kReduceScatter, reduction_op);
}
void NcclManager::AddToAllToAll(std::unique_ptr<Participant> participant,
const Context& context) {
AddParticipant(std::move(participant), context, kAllToAll,
ncclSum /* unused */);
}
void NcclManager::AddBroadcastSend(std::unique_ptr<Participant> participant,
const Context& context) {
participant->root = true;
AddParticipant(std::move(participant), context, kBroadcast,
ncclSum /* unused */);
}
void NcclManager::AddBroadcastRecv(std::unique_ptr<Participant> participant,
const Context& context) {
AddParticipant(std::move(participant), context, kBroadcast,
ncclSum /* unused */);
}
void NcclManager::AddReduceSend(std::unique_ptr<Participant> participant,
const Context& context,
ncclRedOp_t reduction_op) {
AddParticipant(std::move(participant), context, kReduce, reduction_op);
}
void NcclManager::AddReduceRecv(std::unique_ptr<Participant> participant,
const Context& context,
ncclRedOp_t reduction_op) {
participant->root = true;
AddParticipant(std::move(participant), context, kReduce, reduction_op);
}
void NcclManager::SignalMultiNodeReady(const std::string& collective_key) {
Collective* to_run = nullptr;
{
mutex_lock l(mu_);
auto collective_it = collectives_.find(collective_key);
if (collective_it != collectives_.end()) {
Collective* collective = collective_it->second;
collective->multi_node_ready = true;
if (CheckReady(collective_key, collective)) {
to_run = collective;
}
VLOG(2) << "SignalMultiNodeReady collective " << collective_key
<< " to_run " << to_run;
}
}
if (to_run != nullptr) RunCollective(to_run);
}
void NcclManager::AddParticipant(std::unique_ptr<Participant> participant,
const Context& context,
CollectiveType collective_type,
ncclRedOp_t reduction_op) {
Collective* to_run = nullptr;
DataType data_type;
absl::Status nccl_manager_status;
if (participant->input != nullptr) {
data_type = participant->input->dtype();
} else {
data_type = participant->output->dtype();
}
{
mutex_lock l(mu_);
nccl_manager_status = status_;
if (nccl_manager_status.ok()) {
auto collective_it = collectives_.find(context.collective_key);
Collective* collective = nullptr;
if (collective_it == collectives_.end()) {
collective = new Collective(
context.collective_key, data_type, collective_type, reduction_op,
context.num_local_devices, context.num_global_devices,
context.communicator_key);
collectives_.emplace(context.collective_key, collective);
} else {
collective = collective_it->second;
}
// Check `collective` is correct and consistent.
if (collective->status.ok() && !collective->single_node &&
collective->communicator_key.empty()) {
collective->status = absl::InternalError(absl::StrCat(
"Collective ", reduction_op,
" is multi node with num_local_devices=",
collective->num_local_devices,
" and num_global_devices=", collective->num_global_devices,
" but has an empty communicator_key"));
}
if (collective->status.ok() && collective->communicator_key.size() !=
context.communicator_key.size()) {
collective->status = absl::InternalError(
absl::StrCat("Collective ", reduction_op,
" mismatch in member communicator_key with size ",
collective->communicator_key.size(),
" and arg communicator_key with size ",
context.communicator_key.size()));
}
if (collective->status.ok() && collective->type != collective_type) {
collective->status = absl::InternalError(absl::StrCat(
"Collective ", reduction_op, " previously initialized with type ",
collective->type, " but now got type ", collective_type));
}
if (collective->status.ok() &&
collective->num_global_devices != context.num_global_devices) {
collective->status = absl::InternalError(
absl::StrCat("Collective ", reduction_op,
" previously initialized with num_global_devices ",
collective->num_global_devices, " but now got ",
context.num_global_devices));
}
if (collective->status.ok() &&
collective->num_local_devices != context.num_local_devices) {
collective->status = absl::InternalError(
absl::StrCat("Collective ", reduction_op,
"previously initialized with num_local_devices ",
collective->num_local_devices, " but now got ",
context.num_local_devices));
}
if (collective->status.ok() &&
collective->participants.size() >= collective->num_local_devices) {
collective->status = absl::InternalError(absl::StrCat(
"Collective ", reduction_op, " expected ",
collective->num_local_devices, " participants but now has ",
collective->participants.size(),
" with one more participant being added"));
}
if (collective->status.ok() && collective->root_rank >= 0 &&
context.source_rank >= 0 &&
collective->root_rank != context.source_rank) {
collective->status = absl::InternalError(absl::StrCat(
"Collective ", collective->collective_key,
" already has root_rank ", collective->root_rank,
" but new participant has root_rank ", context.source_rank));
}
if (collective->status.ok() &&
!kValidDataTypes.Contains(collective->data_type)) {
collective->status = absl::InternalError(absl::StrCat(
"Collective ", collective->collective_key,
" expected data types compatible with NCCL but instead got ",
DataTypeString(collective->data_type)));
}
if (context.source_rank >= 0) {
collective->root_rank = context.source_rank;
}
collective->participants.emplace_back(std::move(participant));
Participant* p = collective->participants.back().get();
if (collective->status.ok()) {
if (collective_type == kAllGather &&
p->output->NumElements() <
p->input->NumElements() * collective->num_global_devices) {
collective->status = absl::InternalError(absl::StrCat(
"Output tensor must be able to hold ",
p->input->NumElements() * collective->num_global_devices,
" elements for AllGather, but got ", p->output->NumElements()));
} else if (collective_type == kAllToAll &&
p->output->NumElements() < p->input->NumElements()) {
collective->status = absl::InternalError(absl::StrCat(
"Output tensor must be able to hold ", p->input->NumElements(),
" elements for AllToAll, but got ", p->output->NumElements()));
}
}
++collective->available_participants;
if (CheckReady(context.collective_key, collective)) {
to_run = collective;
}
}
}
if (!nccl_manager_status.ok()) {
participant->done_callback(nccl_manager_status);
return;
}
if (to_run != nullptr) RunCollective(to_run);
}
bool NcclManager::CheckReady(const std::string& collective_key,
Collective* collective) {
if (collective->available_participants == collective->num_local_devices) {
if (collective->num_global_devices == collective->num_local_devices ||
collective->multi_node_ready) {
// Ownership transferred to callee.
collectives_.erase(collective_key);
return true;
}
}
return false;
}
void NcclManager::RunCollective(Collective* collective) {
// For TraceMeConsumer in Connection::RPCDone().
tsl::profiler::TraceMeProducer traceme("Schedule Collective");
collective->trace_context = traceme.GetContextId();
static mutex collective_mu(LINKER_INITIALIZED);
absl::Status status = collective->status;
if (status.ok()) {
status = GetCommunicator(collective, &collective->communicator);
}
for (int i = 0; status.ok() && i < collective->num_local_devices; ++i) {
Participant* p = collective->participants[i].get();
NcclStream* nccl_stream = collective->communicator->members[i].nccl_stream;
CHECK(nccl_stream != nullptr);
const int rank = p->global_rank >= 0 ? p->global_rank : i;
if (p->input != nullptr) {
// Wait to ensure that the kernel that produces the data in the input
// tensor has finished running before the nccl kernel runs on the
// communication stream.
status = nccl_stream->stream->WaitFor(p->tensor_stream);
}
if (p->root) {
if (collective->root_rank == -1) {
collective->root_rank = rank;
} else if (collective->root_rank != rank) {
status = absl::InternalError(absl::StrCat(
"Inconsistent root rank ", collective->root_rank, " and GPU id ",
p->gpu_device_id, " rank ", rank, " also marked as root."));
}
}
VLOG(2) << "RunCollective rank " << rank << " global_rank "
<< p->global_rank << " root_rank " << collective->root_rank;
}
if (status.ok() && collective->type == kBroadcast &&
collective->root_rank < 0) {
status = absl::InternalError(absl::StrCat(
"Root rank not indicated for collective ", collective->collective_key));
}
if (!status.ok()) {
for (int i = 0; i < collective->num_local_devices; ++i) {
collective->participants[i]->done_callback(status);
}
collective->Unref();
return;
}
{
// Allow only one collective at a time to queue kernels for launching. This
// is to prevent collectives from deadlocking each other.
// Note that it would be possible to run multiple collectives at once, if
// they have non-intersecting sets of devices.
mutex_lock l(collective_mu);
for (int i = 0; i < collective->num_local_devices; ++i) {
NcclStream* nccl_stream =
collective->communicator->members[i].nccl_stream;
mutex_lock l(nccl_stream->mu);
nccl_stream->pending_launches_.push_front(std::make_pair(collective, i));
// Ownership is shared between LoopKernelLaunches for each stream in this
// collective.
collective->Ref();
nccl_stream->cv.notify_all();
}
}
collective->Unref();
}
namespace {
// For tracing purpose.
size_t ComputeBufferSize(const NcclManager::Participant* p,
DataType data_type) {
size_t num_elements = 0;
if (p->output) {
num_elements += p->output->NumElements();
} else if (p->input) {
num_elements += p->input->NumElements();
}
return num_elements * DataTypeSize(data_type);
}
} // namespace
void NcclManager::LoopKernelLaunches(NcclStream* nccl_stream) {
#if TENSORFLOW_USE_ROCM
se::Stream* comm_stream = nccl_stream->stream;
#else
se::Stream* comm_stream = nccl_stream->stream.get();
#endif
std::unique_ptr<se::ActivateContext> scoped_context =
nccl_stream->executor->Activate();
cudaStream_t cu_stream = reinterpret_cast<cudaStream_t>(
comm_stream->platform_specific_handle().stream);
while (true) {
// Find collective to run.
std::pair<Collective*, int> next_launch;
{
VLOG(3) << "Locking mutex nccl_stream " << nccl_stream;
mutex_lock l(nccl_stream->mu);
while (nccl_stream->pending_launches_.empty()) {
if (nccl_stream->shutdown_requested) {
// No work and shutdown requested, exit.
return;
}
nccl_stream->cv.wait(l);
}
next_launch = nccl_stream->pending_launches_.back();
nccl_stream->pending_launches_.pop_back();
}
// Launch the nccl kernel.
Collective* collective = next_launch.first;
tsl::profiler::TraceMeConsumer traceme("Run Collective",
collective->trace_context);
ncclDataType_t data_type = ToNcclType(collective->data_type);
int p_idx = next_launch.second;
Participant* p = collective->participants[p_idx].get();
auto nccl_comm = collective->communicator->members[p_idx].nccl_comm;
ncclResult_t nccl_result = ncclSuccess;
switch (collective->type) {
case kAllReduce: {
const void* sendbuff = p->input->tensor_data().data();
void* recvbuff = const_cast<char*>(p->output->tensor_data().data());
VLOG(2) << "call NcclAllReduce collective_key "
<< collective->collective_key << " participant " << p_idx
<< " num_participants " << collective->participants.size()
<< " sendbuff " << sendbuff << " recvbuff " << recvbuff
<< " nccl_comm " << nccl_comm << " comm_stream " << comm_stream
<< " cuda_stream " << cu_stream;
profiler::AnnotatedTraceMe traceme([&] {
return tsl::profiler::TraceMeEncode(
"ncclAllReduce",
{{"buffer_size", ComputeBufferSize(p, collective->data_type)},
{"collective_type", "all_reduce"}});
});
nccl_result = ncclAllReduce(sendbuff, recvbuff, p->input->NumElements(),
data_type, collective->reduction_op,
nccl_comm, cu_stream);
break;
}
case kBroadcast: {
const void* sendbuff = nullptr;
void* recvbuff = nullptr;
int num_elements = -1;
if (p->input) {
sendbuff = p->input->tensor_data().data();
num_elements = p->input->NumElements();
}
if (p->output) {
recvbuff = const_cast<char*>(p->output->tensor_data().data());
num_elements = p->output->NumElements();
} else {
// Operate in-place if no output (for the src node).
recvbuff = const_cast<void*>(sendbuff);
}
if (num_elements < 0) {
p->done_callback(absl::InternalError(
"Both input and output are null in ncclBroadcast"));
collective->Unref();
continue;
}
VLOG(2) << "call NcclBroadcast collective_key "
<< collective->collective_key << " participant " << p_idx
<< " sendbuff " << sendbuff << " recvbuff " << recvbuff
<< " nccl_comm " << nccl_comm << " comm_stream " << comm_stream
<< " cuda_stream " << cu_stream;
profiler::AnnotatedTraceMe traceme([&] {
return tsl::profiler::TraceMeEncode(
"ncclBroadcast",
{{"buffer_size", ComputeBufferSize(p, collective->data_type)},
{"collective_type", "broadcast"}});
});
nccl_result =
ncclBroadcast(sendbuff, recvbuff, num_elements, data_type,
collective->root_rank, nccl_comm, cu_stream);
break;
}
case kReduce: {
const void* sendbuff = p->input->tensor_data().data();
void* recvbuff =
p->output ? const_cast<char*>(p->output->tensor_data().data())
: nullptr;
profiler::AnnotatedTraceMe traceme([&] {
return tsl::profiler::TraceMeEncode(
"buffer_size",
{{"output_size", ComputeBufferSize(p, collective->data_type)},
{"collective_type", "reduce"}});
});
nccl_result = ncclReduce(sendbuff, recvbuff, p->input->NumElements(),
data_type, collective->reduction_op,
collective->root_rank, nccl_comm, cu_stream);
break;
}
case kAllGather: {
const void* sendbuff = p->input->tensor_data().data();
void* recvbuff = const_cast<char*>(p->output->tensor_data().data());
VLOG(2) << "call NcclAllGather collective_key "
<< collective->collective_key << " participant " << p_idx
<< " sendbuff " << sendbuff << " sendcount "
<< p->input->NumElements() << " recvbuff " << recvbuff
<< " recvcount " << p->output->NumElements() << " nccl_comm "
<< nccl_comm << " comm_stream " << comm_stream
<< " cuda_stream " << cu_stream;
profiler::AnnotatedTraceMe traceme([&] {
return tsl::profiler::TraceMeEncode(
"ncclAllGather",
{{"buffer_size", ComputeBufferSize(p, collective->data_type)},
{"collective_type", "all_gather"}});
});
nccl_result = ncclAllGather(sendbuff, recvbuff, p->input->NumElements(),
data_type, nccl_comm, cu_stream);
break;
}
case kReduceScatter: {
const void* sendbuff = p->input->tensor_data().data();
void* recvbuff = const_cast<char*>(p->output->tensor_data().data());
VLOG(2) << "call NcclReduceScatter collective_key "
<< collective->collective_key << " participant " << p_idx
<< " num_participants " << collective->participants.size()
<< " sendbuff " << sendbuff << " recvbuff " << recvbuff
<< " nccl_comm " << nccl_comm << " comm_stream " << comm_stream
<< " cuda_stream " << cu_stream;
profiler::AnnotatedTraceMe traceme([&] {
return tsl::profiler::TraceMeEncode(
"ncclReduceScatter",
{{"buffer_size", ComputeBufferSize(p, collective->data_type)},
{"collective_type", "reduce_scatter"}});
});
nccl_result = ncclReduceScatter(
sendbuff, recvbuff, p->output->NumElements(), data_type,
collective->reduction_op, nccl_comm, cu_stream);
break;
}
case kAllToAll: {
const char* sendbuff = p->input->tensor_data().data();
char* recvbuff = const_cast<char*>(p->output->tensor_data().data());
size_t count =
p->input->NumElements() / collective->participants.size();
size_t rank_offset = count * DataTypeSize(collective->data_type);
VLOG(2) << "call Nccl All to All collective_key "
<< collective->collective_key << " participant " << p_idx
<< " num_participants " << collective->participants.size()
<< " sendbuff " << static_cast<const void*>(sendbuff)
<< " recvbuff " << static_cast<void*>(recvbuff) << " nccl_comm "
<< nccl_comm << " comm_stream " << comm_stream
<< " cuda_stream " << cu_stream;
profiler::AnnotatedTraceMe traceme([&] {
return tsl::profiler::TraceMeEncode(
"ncclAllToAll",
{{"buffer_size", ComputeBufferSize(p, collective->data_type)},
{"collective_type", "all_to_all"}});
});
ncclGroupStart();
for (int i = 0; i < collective->participants.size(); ++i) {
ncclSend(sendbuff + i * rank_offset, count, data_type,
collective->participants[i]->global_rank, nccl_comm,
cu_stream);
ncclRecv(recvbuff + i * rank_offset, count, data_type,
collective->participants[i]->global_rank, nccl_comm,
cu_stream);
}
nccl_result = ncclGroupEnd();
break;
}
}
// Run the done_callback when the nccl kernel finishes running.
auto done_callback = [collective, p_idx, nccl_result]() {
VLOG(2) << "done Nccl kernel collective_key "
<< collective->collective_key << " participant " << p_idx
<< " ncclResult " << nccl_result;
if (nccl_result == ncclSuccess) {
collective->participants[p_idx]->done_callback(absl::OkStatus());
} else {
// Propagate the error, but note that if other members of the collective
// did launch their kernels, then they are hanging.
collective->participants[p_idx]->done_callback(
absl::UnknownError(absl::StrCat("Error invoking NCCL: ",
ncclGetErrorString(nccl_result))));
}
collective->Unref();
};
p->event_mgr->ThenExecute(comm_stream, done_callback);
}
}
void NcclManager::StartAbort(const absl::Status& s) {
absl::flat_hash_map<std::string, Collective*> collectives;
std::vector<std::unique_ptr<Communicator>> communicators;
{
mutex_lock l(mu_);
if (!status_.ok()) {
LOG(WARNING)
<< "NcclManager already aborted, ignoring subsequent StartAbort with "
<< s;
return;
}
status_ = s;
collectives.swap(collectives_);
communicators.swap(communicators_);
}
VLOG(2) << "Aborted NcclManager " << this << " with " << collectives.size()
<< " collectives and " << communicators.size()
<< " comms with status " << s;
// collectives_ contains pending launches that haven't been dispatched to
// kernel launch threads, so we can simply invoke the done callbacks of them.
for (const auto& item : collectives) {
for (const std::unique_ptr<Participant>& p : item.second->participants) {
p->done_callback(s);
}
item.second->Unref();
}
// Abort ncclComm. Note that there could be multiple ncclComm per device,
// and ncclCommAbort contains cuda calls that requires device
// synchronization. That is a collective on nccl_comm_0 can block
// ncclCommAbort(nccl_comm_1), so we need to abort all ncclComm in a
// concurrent fashion. This assumes that there's only one active NcclManager
// at a time.
UnboundedWorkQueue queue(Env::Default(), "nccl_abort");
int num_comms = 0;
for (std::unique_ptr<Communicator>& communicator : communicators) {
num_comms += communicator->members.size();
}
BlockingCounter pending(num_comms);
for (std::unique_ptr<Communicator>& communicator : communicators) {
for (CommunicatorMember& member : communicator->members) {
queue.Schedule([&member, &pending]() {
ncclCommAbort(member.nccl_comm);
member.nccl_comm = nullptr;
pending.DecrementCount();
});
}
}
pending.Wait();
}
void NcclManager::Reset() {
mutex_lock l(mu_);
status_ = absl::Status();
VLOG(2) << "Reset NcclManager " << this;
}
} // namespace tensorflow
#endif // GOOGLE_CUDA || TENSORFLOW_USE_ROCM
tensorflow/core/nccl/nccl_manager.h
+284
View File
@@ -0,0 +1,284 @@
/* Copyright 2016 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#ifndef TENSORFLOW_CORE_NCCL_NCCL_MANAGER_H_
#define TENSORFLOW_CORE_NCCL_NCCL_MANAGER_H_
#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
#include <vector>
// TODO(rmlarsen): Get rid of this workaround. "gpu_assert" is defined when
// setting EIGEN_USE_THREADS. But when defining EIGEN_USE_THREADS here,
// incAtomic and other CUDA specific symbols are no longer recognized.
#ifndef gpu_assert
#define gpu_assert(x)
#endif
#include "absl/container/flat_hash_map.h"
#if GOOGLE_CUDA
#include "third_party/nccl/nccl.h"
#elif TENSORFLOW_USE_ROCM
#include "rocm/rocm_config.h"
#if (TF_ROCM_VERSION >= 50200)
#include "rocm/include/rccl/rccl.h"
#else
#include "rocm/include/rccl.h"
#endif
#include "tensorflow/core/common_runtime/gpu_device_context.h"
#endif
#include "tensorflow/core/common_runtime/gpu/gpu_event_mgr.h"
#include "tensorflow/core/framework/device_base.h"
#include "tensorflow/core/framework/tensor.h"
#include "tensorflow/core/platform/mutex.h"
#include "tensorflow/core/platform/stream_executor.h"
namespace tensorflow {
// NCCL manager is used to make the asynchronous communicator calls and to
// manage the per-device streams used for communication.
//
// See nccl_ops.cc for example usage, including description of memory
// management and stream synchronization.
class NcclManager {
public:
typedef std::function<void(absl::Status)> DoneCallback;
NcclManager();
~NcclManager();
static NcclManager* instance();
#if TENSORFLOW_USE_ROCM
static int instance_count;
#endif
// Calls `ncclGetUniqueId` and returns the id as a string. The returned value
// may be shared with other participants on different nodes and passed in to
// multi-node collective invocations.
std::string GenerateCommunicatorKey();
// A participant in a Collective.
struct Participant {
Participant(se::StreamExecutor* executor, se::Stream* tensor_stream,
const DeviceBase::AcceleratorDeviceInfo* info,
const Tensor* input, Tensor* output, int global_rank,
DoneCallback done_callback)
: executor(executor),
tensor_stream(tensor_stream),
event_mgr(info->event_mgr),
gpu_device_id(info->gpu_id),
#if TENSORFLOW_USE_ROCM
context(static_cast<GPUDeviceContext*>(info->default_context)),
#endif
input(input),
output(output),
global_rank(global_rank),
done_callback(std::move(done_callback)),
root(false) {
DCHECK(executor != nullptr);
DCHECK(event_mgr != nullptr);
DCHECK(tensor_stream != nullptr);
}
// StreamExecutor for the device. Expected to be live for process lifetime.
se::StreamExecutor* const executor = nullptr;
// `tensor_stream` is the stream that should be waited on to ensure
// `input`'s data is available on the GPU for the communication stream to
// access. It is also the stream that will use the produced data;
// `done_callback` is not called until the next kernel launched on `stream`
// would see the data. Owned by the caller, who must keep it live until
// `done_callback` is called.
se::Stream* const tensor_stream;
// EventMgr which polls on executor.
// Owned by the caller, who must keep it live until `done_callback` is
// called.
EventMgr* const event_mgr;
const int gpu_device_id;
#if TENSORFLOW_USE_ROCM
GPUDeviceContext* const context;
#endif
// Owned by the caller, who must keep it live until `done_callback` is
// called. Is NULL for participants that only receive data.
const Tensor* input;
// Owned by the caller, who must keep it live until `done_callback` is
// called. Is NULL for participants that only send data.
Tensor* output;
// Rank across all devices and all nodes.
// `global_rank` is not required for single-node collectives.
const int global_rank;
// The callback which is called at the completion of the NCCL operation.
// When called, `output` has been set to the result of the operation. (note:
// the stream may not yet have been synced)
DoneCallback done_callback;
// True if this is the root of the collective, e.g. source of broadcast.
bool root;
};
// Data that provides context for the collective operation, including the
// operation key, number of participants, and communicator key.
struct Context {
Context(const std::string& collective_key, int num_local_devices,
int num_global_devices, const std::string& communicator_key,
int source_rank)
: collective_key(collective_key),
num_local_devices(num_local_devices),
num_global_devices(num_global_devices),
communicator_key(communicator_key),
source_rank(source_rank) {}
// Unique key for this collective instance
const std::string& collective_key;
// Devices local to this node
int num_local_devices;
// Devices across all nodes
int num_global_devices;
// In order to use NCCL across nodes, the callee first has to generate a
// `communicator_key` via `GenerateCommunicatorKey()` function and share
// this with all the other nodes. Each node should pass in this
// `communicator_key` to the `NcclManager` functions.
// `communicator_key` is not required for single-node collectives and can be
// empty.
const std::string& communicator_key;
// Rank of broadcast source.
int source_rank;
};
// Adds one participant to an all-reduce.
void AddToAllReduce(std::unique_ptr<Participant> participant,
const Context& context, ncclRedOp_t reduction_op);
// Adds one participant to an all-gather.
void AddToAllGather(std::unique_ptr<Participant> participant,
const Context& context);
// Adds one participant to a reduce-scatter.
void AddToReduceScatter(std::unique_ptr<Participant> participant,
const Context& context, ncclRedOp_t reduction_op);
// AddBroadcastSend and AddBroadcastRecv combine to send data from one sender
// to all receivers.
void AddBroadcastSend(std::unique_ptr<Participant> participant,
const Context& context);
void AddBroadcastRecv(std::unique_ptr<Participant> participant,
const Context& context);
// AddReduceSend and AddReduceRecv combine to send data from all senders
// to one receiver.
void AddReduceSend(std::unique_ptr<Participant> participant,
const Context& context, ncclRedOp_t reduction_op);
void AddReduceRecv(std::unique_ptr<Participant> participant,
const Context& context, ncclRedOp_t reduction_op);
// Adds one participant to an all-to-all.
void AddToAllToAll(std::unique_ptr<Participant> participant,
const Context& context);
// Signals that the `Collective` corresponding to `key` is ready to launch
// across all nodes participating in this multi-node collective operation.
//
// This should only be called for multi-node collectives; single-node
// collectives are implicitly ready when all participants have called Add*
// function.
void SignalMultiNodeReady(const std::string& collective_key);
// Aborts all collectives. After abortion, no further collectives can be
// launched with this NcclManager.
void StartAbort(const absl::Status& s);
// Resets a previously aborted NcclManager, making it available for future
// collectives.
void Reset();
private:
enum CollectiveType {
kAllReduce = 1,
kBroadcast = 2,
kReduce = 3,
kAllGather = 4,
kReduceScatter = 5,
kAllToAll = 6,
};
struct Collective;
struct Communicator;
struct CommunicatorMember;
struct NcclStream;
// Gets the `Communicator` object that will be used to enqueue NCCL kernels
// for `collective`, and returns it via `communicator`.
//
// This may involve creating CUDA streams and NCCL initialization. If a NCCL
// or CUDA error occurs in the process, this returns an INTERNAL error with
// the corresponding NCCL/CUDA error string.
absl::Status GetCommunicator(Collective* collective,
Communicator** communicator);
// Adds a participant device to the local `Collective` instance corresponding
// to `collective_key`. Launches the `Collective` if it is ready, which it
// checks by calling `CheckReady()`. Also performs consistency and sanity
// checks before launching.
void AddParticipant(std::unique_ptr<Participant> participant,
const Context& context, CollectiveType collective_type,
ncclRedOp_t reduction_op);
// If `collective` is ready to run, removes it from the `collectives_` map and
// returns true. Otherwise returns false.
// Assumes `collective_key` corresponds to `collective`.
//
// A collective is ready to run when all local participants have called Add*
// function, and the collective is signalled globally ready via
// `SetMultiNodeReady`.
bool CheckReady(const std::string& collective_key, Collective* collective)
TF_EXCLUSIVE_LOCKS_REQUIRED(mu_);
// Run <collective>. This calls takes ownership of <collective>.
void RunCollective(Collective* collective);
void LoopKernelLaunches(NcclStream* stream);
mutex mu_;
// Maps key to collectives currently being assembled or run.
absl::flat_hash_map<std::string, Collective*> collectives_ TF_GUARDED_BY(mu_);
// Maps a device to the communication streams that make up its collective.
// This is used to share the stream across different communicators that
// include the same device.
absl::flat_hash_map<se::StreamExecutor*, std::vector<NcclStream*>>
device_to_comm_streams_ TF_GUARDED_BY(mu_);
std::vector<std::unique_ptr<Communicator>> communicators_ TF_GUARDED_BY(mu_);
absl::Status status_ TF_GUARDED_BY(mu_);
NcclManager(const NcclManager&) = delete;
void operator=(const NcclManager&) = delete;
};
} // namespace tensorflow
#endif // GOOGLE_CUDA || TENSORFLOW_USE_ROCM
#endif // TENSORFLOW_CORE_NCCL_NCCL_MANAGER_H_
tensorflow/core/nccl/nccl_manager_test.cc
+972
View File
@@ -0,0 +1,972 @@
/* Copyright 2016 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include <atomic>
#include <cstdint>
#include <cstdlib>
#include <memory>
#include <string>
#include <utility>
#include <gtest/gtest.h>
#include "absl/log/check.h"
#include "absl/log/log.h"
#include "absl/status/status.h"
#include "absl/strings/str_cat.h"
#include "xla/tsl/protobuf/error_codes.pb.h"
#include "tensorflow/core/framework/types.pb.h"
#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
#include <algorithm>
#include <random>
#include <vector>
#include "tensorflow/core/common_runtime/device_factory.h"
#include "tensorflow/core/common_runtime/gpu/gpu_device.h"
#include "tensorflow/core/framework/tensor_testutil.h"
#include "tensorflow/core/lib/core/status_test_util.h"
#include "tensorflow/core/nccl/nccl_manager.h"
#include "tensorflow/core/platform/test.h"
#include "tensorflow/core/platform/unbounded_work_queue.h"
namespace tensorflow {
static std::vector<std::unique_ptr<BaseGPUDevice>> GetGPUDevices() {
std::vector<std::unique_ptr<Device>> devices;
TF_CHECK_OK(DeviceFactory::GetFactory(DEVICE_GPU)
->AddDevices(SessionOptions(), "", &devices));
std::vector<std::unique_ptr<BaseGPUDevice>> gpus;
for (std::unique_ptr<Device>& device : devices) {
if (device->device_type() == "GPU") {
// If `device_type()` is GPU, this `Device` is guaranteed to be a
// `BaseGPUDevice`, which is a subclass of `Device`.
gpus.emplace_back(static_cast<BaseGPUDevice*>(device.release()));
}
}
return gpus;
}
template <typename Scalar>
class NcclManagerTest : public ::testing::Test {
public:
// A single all-reduce to apply.
struct TestCase {
TestCase(int num_nodes, int num_ranks_per_node)
: num_nodes(num_nodes), num_ranks_per_node(num_ranks_per_node) {}
std::vector<Tensor> ins;
std::vector<Tensor> outs;
Tensor expected;
const int num_nodes;
const int num_ranks_per_node;
mutex mu;
absl::Status final_status;
int num_completed TF_GUARDED_BY(mu) = 0;
condition_variable done_cv;
};
static void SetUpTestSuite() {
setenv("NCCL_DEBUG", "INFO", 1 /* replace */);
setenv("NCCL_LAUNCH_MODE", "PARALLEL", 1 /* replace */);
devices_ = new std::vector<std::unique_ptr<BaseGPUDevice>>(GetGPUDevices());
VLOG(1) << "Running test with " << devices_->size() << " gpus";
if (devices_->size() <= 1) {
LOG(FATAL) << "Cannot run NCCL test without multiple GPUs";
}
work_queue_ = new UnboundedWorkQueue(Env::Default(), "nccl_manager_test");
}
void SetUp() override {
ASSERT_GT(devices_->size(), 0) << "No GPUs found";
ASSERT_NE(work_queue_, nullptr);
}
static int32_t NumGPUs() { return static_cast<int32_t>(devices_->size()); }
// Let N = #GPUs. When N is even, num_nodes=2 and num_ranks_per_node=N/2.
// When N is odd, num_nodes=2 and num_ranks_per_node=(N-1)/2.
static void PopulateMultiNodeParams(int* num_nodes, int* num_ranks_per_node) {
const auto num_gpus = NumGPUs();
CHECK_GT(num_gpus, 1);
*num_nodes = 2;
if (num_gpus % 2 == 0) {
*num_ranks_per_node = num_gpus / 2;
} else {
*num_ranks_per_node = (num_gpus - 1) / 2;
}
}
static void TearDownTestSuite() {
delete devices_;
delete work_queue_;
}
TestCase* MakeReductionTestCase(int num_nodes, int num_ranks_per_node,
ncclRedOp_t reduction_op, TensorShape shape,
float value_offset) {
TestCase* test_case = new TestCase(num_nodes, num_ranks_per_node);
test_case->expected = Tensor(data_type_, shape);
if (reduction_op == ncclProd) {
test::FillFn<Scalar>(&test_case->expected,
[](int) { return static_cast<Scalar>(1); });
} else if (reduction_op == ncclSum) {
test::FillFn<Scalar>(&test_case->expected,
[](int) { return static_cast<Scalar>(0); });
} else if (reduction_op == ncclMax) {
test::FillFn<Scalar>(&test_case->expected, [](int) { return -max_; });
} else if (reduction_op == ncclMin) {
test::FillFn<Scalar>(&test_case->expected, [](int) { return max_; });
} else {
LOG(FATAL) << "Invalid reduction_op " << reduction_op;
}
float value_scale = 0.01; // Small scale to avoid fp16 overflow.
for (int node = 0; node < num_nodes; ++node) {
for (int local_rank = 0; local_rank < num_ranks_per_node; ++local_rank) {
auto* device = GetDevice(num_ranks_per_node, node, local_rank);
auto* stream = device->tensorflow_accelerator_device_info()->stream;
Tensor in_cpu(data_type_, shape);
test::FillFn<Scalar>(&in_cpu, [&](int index) {
return static_cast<Scalar>((index + 1) * value_scale + value_offset);
});
for (int j = 0; j < shape.num_elements(); ++j) {
auto in_val = in_cpu.flat<Scalar>()(j);
auto out_expr = test_case->expected.template flat<Scalar>();
if (reduction_op == ncclProd) {
out_expr(j) = out_expr(j) * in_val;
} else if (reduction_op == ncclSum) {
out_expr(j) = out_expr(j) + in_val;
} else if (reduction_op == ncclMax) {
if (in_val > out_expr(j)) {
out_expr(j) = in_val;
}
} else if (reduction_op == ncclMin) {
if (in_val < out_expr(j)) {
out_expr(j) = in_val;
}
}
}
value_scale *= 10;
test_case->ins.emplace_back(GpuAllocator(device), data_type_, shape);
test_case->outs.emplace_back(GpuAllocator(device), data_type_, shape);
const Tensor& in_gpu = test_case->ins.back();
auto in_gpu_mem = AsDeviceMemory(in_gpu.flat<Scalar>().data());
TF_CHECK_OK(stream->Memcpy(&in_gpu_mem, in_cpu.flat<Scalar>().data(),
in_cpu.TotalBytes()));
}
}
return test_case;
}
TestCase* MakeGatherTestCase(int num_nodes, int num_ranks_per_node,
TensorShape in_shape, TensorShape out_shape) {
TestCase* test_case = new TestCase(num_nodes, num_ranks_per_node);
test_case->expected = Tensor(data_type_, out_shape);
test::FillFn<Scalar>(&test_case->expected,
[](int) { return static_cast<Scalar>(0); });
float value_scale = 0.01; // Small scale to avoid fp16 overflow.
for (int node = 0; node < num_nodes; ++node) {
for (int i = 0; i < num_ranks_per_node; ++i) {
auto* device = GetDevice(num_ranks_per_node, node, i);
auto* stream = device->tensorflow_accelerator_device_info()->stream;
Tensor in_cpu(data_type_, in_shape);
test::FillFn<Scalar>(&in_cpu, [&](int index) {
return static_cast<Scalar>((index + 1) * value_scale);
});
// Starting index for this rank's tensor in the all-gathered output.
int32_t gather_idx =
(node * num_ranks_per_node + i) * in_shape.num_elements();
for (int j = 0; j < in_shape.num_elements(); ++j) {
auto in_val = in_cpu.flat<Scalar>()(j);
auto out_expr = test_case->expected.template flat<Scalar>();
out_expr(gather_idx + j) = in_val;
}
value_scale *= 10;
test_case->ins.emplace_back(GpuAllocator(device), data_type_, in_shape);
test_case->outs.emplace_back(GpuAllocator(device), data_type_,
out_shape);
const Tensor& in_gpu = test_case->ins.back();
auto in_gpu_mem = AsDeviceMemory(in_gpu.flat<Scalar>().data());
TF_CHECK_OK(stream->Memcpy(&in_gpu_mem, in_cpu.flat<Scalar>().data(),
in_cpu.TotalBytes()));
}
}
return test_case;
}
// Make a broadcast test which broadcasts a tensor with shape `shape` from
// `src_node`, `src_rank` to all other ranks.
// If `in_place` is true, input and output are the same for the source,
// otherwise they are tensors backed by different buffers.
TestCase* MakeBroadcastTestCase(int num_nodes, int num_ranks_per_node,
TensorShape shape, int src_node, int src_rank,
bool in_place) {
TestCase* test_case = new TestCase(num_nodes, num_ranks_per_node);
test_case->expected = Tensor(data_type_, shape);
test::FillFn<Scalar>(&test_case->expected,
[](int) { return static_cast<Scalar>(1); });
for (int node = 0; node < num_nodes; ++node) {
for (int local_rank = 0; local_rank < num_ranks_per_node; ++local_rank) {
auto* device = GetDevice(num_ranks_per_node, node, local_rank);
if (node == src_node && local_rank == src_rank) {
test_case->ins.emplace_back(GpuAllocator(device), data_type_, shape);
if (in_place) {
test_case->outs.emplace_back(test_case->ins.back());
} else {
test_case->outs.emplace_back(GpuAllocator(device), data_type_,
shape);
}
Tensor in_cpu(data_type_, shape);
test::FillFn<Scalar>(&in_cpu,
[](int) { return static_cast<Scalar>(1); });
const Tensor& in_gpu = test_case->ins.back();
auto in_gpu_mem = AsDeviceMemory(in_gpu.flat<Scalar>().data());
auto* stream = device->tensorflow_accelerator_device_info()->stream;
TF_CHECK_OK(stream->Memcpy(&in_gpu_mem, in_cpu.flat<Scalar>().data(),
in_cpu.TotalBytes()));
} else {
test_case->ins.emplace_back(Tensor());
test_case->outs.emplace_back(GpuAllocator(device), data_type_, shape);
}
}
}
return test_case;
}
// Waits for the done callback to be called for each participant.
void WaitForTestCompletion(TestCase* test_case) {
mutex_lock l(test_case->mu);
while (test_case->num_completed != test_case->outs.size()) {
test_case->done_cv.wait(l);
}
}
void VerifyResults(TestCase* test_case) {
WaitForTestCompletion(test_case);
TF_ASSERT_OK(test_case->final_status);
// Copy memory to host and verify.
for (int node = 0; node < test_case->num_nodes; ++node) {
for (int local_rank = 0; local_rank < test_case->num_ranks_per_node;
++local_rank) {
auto* device =
GetDevice(test_case->num_ranks_per_node, node, local_rank);
auto* stream = device->tensorflow_accelerator_device_info()->stream;
const int global_rank =
GlobalRank(test_case->num_ranks_per_node, node, local_rank);
const Tensor& out_gpu = test_case->outs[global_rank];
Tensor out_cpu(data_type_, out_gpu.shape());
auto out_gpu_mem = AsDeviceMemory(out_gpu.flat<Scalar>().data());
TF_CHECK_OK(stream->Memcpy(out_cpu.flat<Scalar>().data(), out_gpu_mem,
out_cpu.TotalBytes()));
TF_ASSERT_OK(stream->BlockHostUntilDone());
VLOG(1) << "Verifying rank " << global_rank << " expected shape "
<< test_case->expected.shape() << " out shape "
<< out_cpu.shape();
test::ExpectClose(test_case->expected, out_cpu);
}
}
}
void VerifyError(TestCase* test_case) {
WaitForTestCompletion(test_case);
LOG(INFO) << test_case->final_status;
EXPECT_EQ(test_case->final_status.code(), error::INTERNAL);
}
NcclManager::DoneCallback CreateDoneCallback(TestCase* test_case) {
return [this, test_case](absl::Status s) {
mutex_lock l(test_case->mu);
test_case->final_status.Update(s);
if (++test_case->num_completed == test_case->outs.size()) {
test_case->done_cv.notify_one();
}
};
}
struct NodeState {
NcclManager nccl_manager;
std::atomic<int> launched{0};
};
void RunMultiNodeAllReduceTest(const int num_nodes,
const int num_ranks_per_node) {
std::vector<NodeState> node_states(num_nodes);
RunMultiNodeAllReduceTest(node_states, num_ranks_per_node);
}
void RunMultiNodeAllReduceTest(std::vector<NodeState>& node_states,
const int num_ranks_per_node) {
const int num_nodes = node_states.size();
const int num_global_ranks = num_nodes * num_ranks_per_node;
const std::string collective_key = "allreduce";
// The NcclManagers in this test synchronize in real-time, so we need to run
// each node's code in a separate thread.
// Specifically, the call to ncclGroupEnd() after calling ncclCommInitRank
// waits for all communicators before returning.
// First, initialize the communicator_key used for this collective.
const std::string communicator_key =
node_states[0].nccl_manager.GenerateCommunicatorKey();
for (int op = 0; op < 4; ++op) {
ncclRedOp_t reduction_op = static_cast<ncclRedOp_t>(op);
std::unique_ptr<TestCase> test_case(
this->MakeReductionTestCase(num_nodes, num_ranks_per_node,
reduction_op, TensorShape({2, 3}), 0.0f));
for (int node = 0; node < num_nodes; ++node) {
auto node_fn = [this, node, num_ranks_per_node, num_global_ranks,
&node_states, &communicator_key, &collective_key,
reduction_op, &test_case] {
for (int local_rank = 0; local_rank < num_ranks_per_node;
++local_rank) {
auto* device = GetDevice(num_ranks_per_node, node, local_rank);
auto* info = device->tensorflow_accelerator_device_info();
auto* stream = device->tensorflow_accelerator_device_info()->stream;
const int global_rank =
GlobalRank(num_ranks_per_node, node, local_rank);
auto participant = absl::make_unique<NcclManager::Participant>(
device->executor(), stream, info, &test_case->ins[global_rank],
&test_case->outs[global_rank], global_rank,
this->CreateDoneCallback(test_case.get()));
node_states[node].nccl_manager.AddToAllReduce(
std::move(participant),
{collective_key, num_ranks_per_node, num_global_ranks,
communicator_key, /*source_rank=*/-1},
reduction_op);
VLOG(1) << "AddToAllReduce node " << node << " global_rank "
<< global_rank;
}
// Signal collective ready to launch at this node.
node_states[node].nccl_manager.SignalMultiNodeReady(collective_key);
};
this->work_queue_->Schedule(node_fn);
}
VLOG(2) << "Verifying results";
this->VerifyResults(test_case.get());
}
}
void RunMultiNodeBroadcastTest(const int num_nodes,
const int num_ranks_per_node,
const int src_node, const int src_local_rank,
const bool in_place) {
const int num_global_ranks = num_nodes * num_ranks_per_node;
const int src_global_rank = src_node * num_ranks_per_node + src_local_rank;
const std::string collective_key = "broadcast";
std::vector<NodeState> node_states(num_nodes);
const std::string communicator_key =
node_states[0].nccl_manager.GenerateCommunicatorKey();
std::unique_ptr<TestCase> test_case(this->MakeBroadcastTestCase(
num_nodes, num_ranks_per_node, TensorShape({5, 6}), src_node,
src_local_rank, in_place));
for (int node = 0; node < num_nodes; ++node) {
for (int local_rank = 0; local_rank < num_ranks_per_node; ++local_rank) {
// Launch each rank in a separate thread to test concurrent,
// randomly-ordered calls into NcclManager.
auto rank_fn = [this, node, num_ranks_per_node, num_global_ranks,
src_global_rank, local_rank, &node_states,
&collective_key, &communicator_key, &test_case]() {
auto* device = GetDevice(num_ranks_per_node, node, local_rank);
auto* info = device->tensorflow_accelerator_device_info();
auto* stream = device->tensorflow_accelerator_device_info()->stream;
const int global_rank =
GlobalRank(num_ranks_per_node, node, local_rank);
auto* input = global_rank == src_global_rank
? &test_case->ins[global_rank]
: nullptr;
auto* output = test_case->outs[global_rank].NumElements() == 0
? nullptr
: &test_case->outs[global_rank];
auto participant = absl::make_unique<NcclManager::Participant>(
device->executor(), stream, info, input, output, global_rank,
this->CreateDoneCallback(test_case.get()));
if (global_rank == src_global_rank) {
node_states[node].nccl_manager.AddBroadcastSend(
std::move(participant),
{collective_key, num_ranks_per_node, num_global_ranks,
communicator_key, src_global_rank});
} else {
node_states[node].nccl_manager.AddBroadcastRecv(
std::move(participant),
{collective_key, num_ranks_per_node, num_global_ranks,
communicator_key, src_global_rank});
}
if (++node_states[node].launched == num_ranks_per_node) {
// Signal collective ready to launch at this node.
node_states[node].nccl_manager.SignalMultiNodeReady(collective_key);
}
};
this->work_queue_->Schedule(std::move(rank_fn));
}
}
VLOG(2) << "Verifying results";
this->VerifyResults(test_case.get());
}
static int GlobalRank(int num_ranks_per_node, int node, int local_rank) {
return node * num_ranks_per_node + local_rank;
}
static BaseGPUDevice* GetDevice(int num_ranks_per_node, int node,
int local_rank) {
const int device_idx = GlobalRank(num_ranks_per_node, node, local_rank);
CHECK_LT(device_idx, devices_->size());
return (*devices_)[device_idx].get();
}
static UnboundedWorkQueue* work_queue_;
private:
static Allocator* GpuAllocator(BaseGPUDevice* device) {
return device->GetAllocator(AllocatorAttributes());
}
static se::DeviceMemory<Scalar> AsDeviceMemory(const Scalar* cuda_memory) {
stream_executor::DeviceAddressBase wrapped(
const_cast<Scalar*>(cuda_memory));
se::DeviceMemory<Scalar> typed(wrapped);
return typed;
}
static std::vector<std::unique_ptr<BaseGPUDevice>>* devices_;
static const DataType data_type_;
static const Scalar max_;
};
template <typename Scalar>
std::vector<std::unique_ptr<BaseGPUDevice>>* NcclManagerTest<Scalar>::devices_ =
nullptr;
template <typename Scalar>
const DataType NcclManagerTest<Scalar>::data_type_ =
DataTypeToEnum<Scalar>::value;
template <typename Scalar>
const Scalar NcclManagerTest<Scalar>::max_ =
Eigen::NumTraits<Scalar>::highest();
template <typename Scalar>
UnboundedWorkQueue* NcclManagerTest<Scalar>::work_queue_ = nullptr;
// Instantiate tests for float and double.
using TypeList = ::testing::Types<float, double>;
TYPED_TEST_SUITE(NcclManagerTest, TypeList);
// Test basic sum reduction.
TYPED_TEST(NcclManagerTest, BasicSumReduction) {
const int num_ranks = this->NumGPUs();
for (int op = 0; op < 4; ++op) {
ncclRedOp_t reduction_op = static_cast<ncclRedOp_t>(op);
std::unique_ptr<typename TestFixture::TestCase> test_case(
this->MakeReductionTestCase(/*num_nodes=*/1, num_ranks, reduction_op,
TensorShape({2, 3}), 0.0f));
for (int rank = 0; rank < num_ranks; ++rank) {
auto* device = this->GetDevice(num_ranks, /*node=*/0, rank);
VLOG(2) << "rank " << rank << " device " << device->name();
auto* info = device->tensorflow_accelerator_device_info();
auto* stream = device->tensorflow_accelerator_device_info()->stream;
auto participant = absl::make_unique<NcclManager::Participant>(
device->executor(), stream, info, &test_case->ins[rank],
&test_case->outs[rank], /*global_rank=*/-1,
this->CreateDoneCallback(test_case.get()));
NcclManager::instance()->AddToAllReduce(
std::move(participant),
{"allreduce", /*num_local_devices=*/num_ranks,
/*num_global_devices=*/num_ranks, /*communicator_key=*/"",
/*source_rank=*/-1},
reduction_op);
}
LOG(INFO) << "Verifying results";
this->VerifyResults(test_case.get());
}
}
// Same as the Basic test, but with multiple threads launching parts of many
// reductions.
//
// To run test longer, increase num_ranks, num_collectives_per_iteration and
// time_limit_micros.
TYPED_TEST(NcclManagerTest, MultipleCallers) {
const int num_ranks = this->NumGPUs();
const int num_collectives_per_iteration = 10;
const int time_limit_micros = 1 * 1000 * 1000; // 1 second
int64_t start = Env::Default()->NowMicros();
srand(Env::Default()->NowMicros());
for (;;) {
std::vector<std::pair<int, int>> case_and_rank;
std::vector<std::unique_ptr<typename TestFixture::TestCase>> test_cases;
for (int i = 0; i < num_collectives_per_iteration; ++i) {
test_cases.emplace_back(this->MakeReductionTestCase(
/*num_nodes=*/1, num_ranks, ncclSum,
TensorShape({100, i % 5 + 1, i % 3 + 1}), 1.1f * i));
for (int j = 0; j < num_ranks; ++j) {
case_and_rank.emplace_back(i, j);
}
}
for (int rank = 0; rank < num_ranks; ++rank) {
auto* device = this->GetDevice(num_ranks, /*node=*/0, rank);
auto* stream = device->tensorflow_accelerator_device_info()->stream;
TF_ASSERT_OK(stream->BlockHostUntilDone());
}
std::shuffle(case_and_rank.begin(), case_and_rank.end(),
std::mt19937(std::random_device()()));
mutex mu; // guards case_and_rank.
const int to_schedule = case_and_rank.size();
for (int i = 0; i < to_schedule; ++i) {
auto fn = [&]() {
int rank;
int test_num;
{
mutex_lock l(mu);
test_num = case_and_rank.back().first;
rank = case_and_rank.back().second;
case_and_rank.pop_back();
}
auto* device = this->GetDevice(num_ranks, /*node=*/0, rank);
auto* info = device->tensorflow_accelerator_device_info();
auto* stream = device->tensorflow_accelerator_device_info()->stream;
typename TestFixture::TestCase* test_case = test_cases[test_num].get();
auto participant = absl::make_unique<NcclManager::Participant>(
device->executor(), stream, info, &test_case->ins[rank],
&test_case->outs[rank], /*global_rank=*/-1,
this->CreateDoneCallback(test_case));
NcclManager::instance()->AddToAllReduce(
std::move(participant),
{absl::StrCat("allreduce", test_num),
/*num_local_devices=*/num_ranks,
/*num_global_devices=*/num_ranks,
/*communicator_key=*/"", /*source_rank=*/-1},
ncclSum);
};
this->work_queue_->Schedule(fn);
}
VLOG(2) << "Verifying results for " << num_collectives_per_iteration
<< " collectives";
for (int i = 0; i < test_cases.size(); ++i) {
this->VerifyResults(test_cases[i].get());
}
int64_t delta = Env::Default()->NowMicros() - start;
if (delta > time_limit_micros) {
LOG(INFO) << "Ran for " << delta << " microsecs, now quitting";
break;
}
}
}
// Test basic all-gather.
TYPED_TEST(NcclManagerTest, BasicAllGather) {
const int num_ranks = this->NumGPUs();
for (int i = 0; i < num_ranks; ++i) {
std::unique_ptr<typename TestFixture::TestCase> test_case(
this->MakeGatherTestCase(/*num_nodes=*/1, num_ranks,
TensorShape({2, 3}),
TensorShape({2 * num_ranks, 3})));
for (int rank = 0; rank < num_ranks; ++rank) {
auto* device = this->GetDevice(num_ranks, /*node=*/0, rank);
VLOG(2) << "rank " << rank << " device " << device->name();
auto* info = device->tensorflow_accelerator_device_info();
auto* stream = device->tensorflow_accelerator_device_info()->stream;
auto participant = absl::make_unique<NcclManager::Participant>(
device->executor(), stream, info, &test_case->ins[rank],
&test_case->outs[rank], rank,
this->CreateDoneCallback(test_case.get()));
NcclManager::instance()->AddToAllGather(
std::move(participant),
{"allgather", /*num_local_devices=*/num_ranks,
/*num_global_devices=*/num_ranks, /*communicator_key=*/"",
/*source_rank=*/-1});
}
LOG(INFO) << "Verifying results";
this->VerifyResults(test_case.get());
}
}
// Test basic broadcast.
TYPED_TEST(NcclManagerTest, BasicBroadcast) {
this->RunMultiNodeBroadcastTest(/*num_nodes=*/1,
/*num_ranks_per_node=*/this->NumGPUs(),
/*src_node=*/0, /*src_local_rank=*/0,
/*in_place=*/false);
}
// Test in-place broadcast.
TYPED_TEST(NcclManagerTest, InPlaceBroadcast) {
this->RunMultiNodeBroadcastTest(/*num_nodes=*/1,
/*num_ranks_per_node=*/this->NumGPUs(),
/*src_node=*/0, /*src_local_rank=*/0,
/*in_place=*/true);
}
// Test broadcast with increasing ranks.
TYPED_TEST(NcclManagerTest, BroadcastWithDifferentRanks) {
for (int num_ranks = 1; num_ranks <= this->NumGPUs(); ++num_ranks) {
const int src_rank = static_cast<int>(random::New64() % num_ranks);
for (int in_place_idx = 0; in_place_idx <= 1; ++in_place_idx) {
const bool in_place = in_place_idx == 0;
this->RunMultiNodeBroadcastTest(/*num_nodes=*/1, num_ranks,
/*src_node=*/0, src_rank, in_place);
}
}
}
// Multi-node NCCL tests.
TEST(NcclManagerTest, CommunicatorKey) {
const std::string communicator_key =
NcclManager::instance()->GenerateCommunicatorKey();
EXPECT_EQ(communicator_key.size(), NCCL_UNIQUE_ID_BYTES);
}
#if !TENSORFLOW_USE_ROCM
// ROCm platform currently does not support simulating a multi-node
// environment, on a single node with multiple GPUS. So tests that rely
// upon such simulation need to be skipped on the ROCm platform
// This test creates `num_nodes` NcclManagers to simulate a multi-node
// environment. It works on a single node with multiple GPUs. It enqueues NCCL
// kernels on separate stream per rank.
TYPED_TEST(NcclManagerTest, MultiNode) {
int num_nodes;
int num_ranks_per_node;
this->PopulateMultiNodeParams(&num_nodes, &num_ranks_per_node);
VLOG(1) << "Calling RunMultiNodeAllReduceTest with num_nodes=" << num_nodes
<< " and num_ranks_per_node=" << num_ranks_per_node;
this->RunMultiNodeAllReduceTest(num_nodes, num_ranks_per_node);
}
#endif
// Tests that specifying `communicator_key` with a single node NCCL collective
// works well.
TYPED_TEST(NcclManagerTest, MultiNodeSingle) {
this->RunMultiNodeAllReduceTest(/*num_nodes=*/1,
/*num_ranks_per_node=*/this->NumGPUs());
}
#if !TENSORFLOW_USE_ROCM
// ROCm platform currently does not support simulating a multi-node
// environment, on a single node with multiple GPUS. So tests that rely
// upon such simulation need to be skipped on the ROCm platform
// Multi-node broadcast.
TYPED_TEST(NcclManagerTest, MultiNodeBroadcast) {
int num_nodes;
int num_ranks_per_node;
this->PopulateMultiNodeParams(&num_nodes, &num_ranks_per_node);
VLOG(1) << "Calling RunMultiNodeBroadcastTest with num_nodes=" << num_nodes
<< " and num_ranks_per_node=" << num_ranks_per_node;
this->RunMultiNodeBroadcastTest(num_nodes, num_ranks_per_node,
/*src_node=*/0, /*src_local_rank=*/0,
/*in_place=*/true);
}
#endif
// Checks that we return error status if a collective_key is used for different
// types of collectives, e.g.a reduction and a broadcast.
TYPED_TEST(NcclManagerTest, ConsistentCollectiveType) {
const int num_ranks = 2;
std::unique_ptr<typename TestFixture::TestCase> test_case(
this->MakeReductionTestCase(/*num_nodes=*/1, num_ranks, ncclSum,
TensorShape({2, 3}), 0.0f));
for (int rank = 0; rank < num_ranks; ++rank) {
auto* device = this->GetDevice(num_ranks, /*node=*/0, rank);
auto* info = device->tensorflow_accelerator_device_info();
auto* stream = device->tensorflow_accelerator_device_info()->stream;
auto participant = absl::make_unique<NcclManager::Participant>(
device->executor(), stream, info, &test_case->ins[rank],
&test_case->outs[rank], /*global_rank=*/-1,
this->CreateDoneCallback(test_case.get()));
if (rank == 0) {
NcclManager::instance()->AddToAllReduce(std::move(participant),
{"bad_coll_type",
/*num_local_devices=*/num_ranks,
/*num_global_devices=*/num_ranks,
/*communicator_key=*/"",
/*source_rank=*/-1},
ncclSum);
} else {
NcclManager::instance()->AddBroadcastSend(
std::move(participant),
{"bad_coll_type",
/*num_local_devices=*/num_ranks,
/*num_global_devices=*/num_ranks,
/*communicator_key=*/"", /*source_rank=*/-1});
}
}
this->VerifyError(test_case.get());
}
// Checks that we return error status if different communicator_key is passed to
// same collective.
TYPED_TEST(NcclManagerTest, ConsistentCommunicatorKey) {
const int num_ranks = 2;
std::unique_ptr<typename TestFixture::TestCase> test_case(
this->MakeReductionTestCase(/*num_nodes=*/1, num_ranks, ncclSum,
TensorShape({2, 3}), 0.0f));
for (int rank = 0; rank < num_ranks; ++rank) {
auto* device = this->GetDevice(num_ranks, /*node=*/0, rank);
auto* info = device->tensorflow_accelerator_device_info();
auto* stream = device->tensorflow_accelerator_device_info()->stream;
auto participant = absl::make_unique<NcclManager::Participant>(
device->executor(), stream, info, &test_case->ins[rank],
&test_case->outs[rank], /*global_rank=*/-1,
this->CreateDoneCallback(test_case.get()));
NcclManager::instance()->AddToAllReduce(
std::move(participant),
{"bad_coll_type",
/*num_local_devices=*/num_ranks,
/*num_global_devices=*/num_ranks,
rank == 0 ? "" : NcclManager::instance()->GenerateCommunicatorKey(),
/*source_rank=*/-1},
ncclSum);
}
this->VerifyError(test_case.get());
}
// Checks that we return error status if the number of devices is inconsistent
// across multiple participants of a collective.
TYPED_TEST(NcclManagerTest, ConsistentNumberOfDevices) {
const int num_ranks = 2;
std::unique_ptr<typename TestFixture::TestCase> test_case(
this->MakeReductionTestCase(/*num_nodes=*/1, num_ranks, ncclSum,
TensorShape({2, 3}), 0.0f));
for (int rank = 0; rank < num_ranks; ++rank) {
auto* device = this->GetDevice(num_ranks, /*node=*/0, rank);
auto* info = device->tensorflow_accelerator_device_info();
auto* stream = device->tensorflow_accelerator_device_info()->stream;
int num_devices = rank == 0 ? num_ranks : num_ranks + 1;
auto participant = absl::make_unique<NcclManager::Participant>(
device->executor(), stream, info, &test_case->ins[rank],
&test_case->outs[rank], /*global_rank=*/-1,
this->CreateDoneCallback(test_case.get()));
NcclManager::instance()->AddToAllReduce(std::move(participant),
{"bad_coll_type",
/*num_local_devices=*/num_devices,
/*num_global_devices=*/num_devices,
/*communicator_key=*/"",
/*source_rank=*/-1},
ncclSum);
}
this->VerifyError(test_case.get());
}
// Checks that we return error status if a broadcast does not have source.
TYPED_TEST(NcclManagerTest, BroadcastNoSource) {
const int num_ranks = 2;
std::unique_ptr<typename TestFixture::TestCase> test_case(
this->MakeBroadcastTestCase(/*num_nodes=*/1, num_ranks,
TensorShape({2, 3}), /*src_node=*/-1,
/*src_rank=*/-1, false));
for (int rank = 0; rank < num_ranks; ++rank) {
auto* device = this->GetDevice(num_ranks, /*node=*/0, rank);
auto* info = device->tensorflow_accelerator_device_info();
auto* stream = device->tensorflow_accelerator_device_info()->stream;
auto participant = absl::make_unique<NcclManager::Participant>(
device->executor(), stream, info, nullptr, &test_case->outs[rank], rank,
this->CreateDoneCallback(test_case.get()));
NcclManager::instance()->AddBroadcastRecv(std::move(participant),
{"bcast_no_send",
/*num_local_devices=*/num_ranks,
/*num_global_devices=*/num_ranks,
/*communicator_key=*/"",
/*source_rank=*/-1});
}
this->VerifyError(test_case.get());
}
// Checks that we return error status if a broadcast has multiple sends.
TYPED_TEST(NcclManagerTest, BroadcastMultipleSends) {
const int num_ranks = 2;
std::unique_ptr<typename TestFixture::TestCase> test_case(
this->MakeBroadcastTestCase(/*num_nodes=*/1, num_ranks,
TensorShape({2, 3}), /*src_node=*/-1,
/*src_rank=*/-1, false));
for (int rank = 0; rank < num_ranks; ++rank) {
auto* device = this->GetDevice(num_ranks, /*node=*/0, rank);
auto* info = device->tensorflow_accelerator_device_info();
auto* stream = device->tensorflow_accelerator_device_info()->stream;
auto participant = absl::make_unique<NcclManager::Participant>(
device->executor(), stream, info, &test_case->outs[rank],
&test_case->outs[rank], rank,
this->CreateDoneCallback(test_case.get()));
NcclManager::instance()->AddBroadcastSend(std::move(participant),
{"bcast_multiple_send",
/*num_local_devices=*/num_ranks,
/*num_global_devices=*/num_ranks,
/*communicator_key=*/"",
/*source_rank=*/-1});
}
this->VerifyError(test_case.get());
}
// Checks that we return error status if a broadcast has inconsistent source
// ranks.
TYPED_TEST(NcclManagerTest, BroadcastInconsistentSource) {
const int num_ranks = 2;
std::unique_ptr<typename TestFixture::TestCase> test_case(
this->MakeBroadcastTestCase(/*num_nodes=*/1, num_ranks,
TensorShape({2, 3}), /*src_node=*/-1,
/*src_rank=*/-1, false));
for (int rank = 0; rank < num_ranks; ++rank) {
auto* device = this->GetDevice(num_ranks, /*node=*/0, rank);
auto* info = device->tensorflow_accelerator_device_info();
auto* stream = device->tensorflow_accelerator_device_info()->stream;
auto participant = absl::make_unique<NcclManager::Participant>(
device->executor(), stream, info, &test_case->outs[rank],
&test_case->outs[rank], rank,
this->CreateDoneCallback(test_case.get()));
NcclManager::instance()->AddBroadcastRecv(std::move(participant),
{"bcast_inconsistent_source",
/*num_local_devices=*/num_ranks,
/*num_global_devices=*/num_ranks,
/*communicator_key=*/"",
/*source_rank=*/rank});
}
this->VerifyError(test_case.get());
}
#if !TENSORFLOW_USE_ROCM
// ROCm platform currently does not support simulating a multi-node
// environment, on a single node with multiple GPUS. So tests that rely
// upon such simulation need to be skipped on the ROCm platform
TYPED_TEST(NcclManagerTest, AbortThenReset) {
using NodeState = typename TestFixture::NodeState;
using TestCase = typename TestFixture::TestCase;
const int num_nodes = 2;
std::vector<NodeState> nodes(num_nodes);
// First do a normal all-reduce to simulate the case when there're
// multiple communicators.
this->RunMultiNodeAllReduceTest(nodes, /* num_ranks_per_node */ 1);
const std::string collective_key = "allreduce";
ncclRedOp_t reduction_op = static_cast<ncclRedOp_t>(0);
auto node_fn = [&](TestCase* test_case, int node,
const std::string& communicator_key) {
auto* device = this->GetDevice(/* num_ranks_per_node */ 1, node,
/* local_rank */ 0);
auto* info = device->tensorflow_accelerator_device_info();
auto* stream = device->tensorflow_accelerator_device_info()->stream;
auto participant = absl::make_unique<NcclManager::Participant>(
device->executor(), stream, info, &test_case->ins[node],
&test_case->outs[node], /* global_rank */ node,
this->CreateDoneCallback(test_case));
nodes[node].nccl_manager.AddToAllReduce(
std::move(participant),
{collective_key, /* num_local_devices */ 1,
/* num_global_devices */ num_nodes, communicator_key,
/*source_rank=*/-1},
reduction_op);
nodes[node].nccl_manager.SignalMultiNodeReady(collective_key);
};
// Use a new communicator_key, which uses a new set of ncclComm underneath.
std::string communicator_key =
nodes[0].nccl_manager.GenerateCommunicatorKey();
// Do a normal all-reduce with this communicator key to initialize ncclComm.
// This is because ncclCommInitRank waits for all ranks and is blocking.
{
std::unique_ptr<typename TestFixture::TestCase> test_case(
this->MakeReductionTestCase(
/* num_nodes */ num_nodes, /* num_ranks_per_node */ 1, reduction_op,
TensorShape({2, 3}), 0.0f));
for (int i = 0; i < num_nodes; ++i) {
this->work_queue_->Schedule(
[&node_fn, &test_case, i, communicator_key]() {
node_fn(test_case.get(), i, communicator_key);
});
}
this->VerifyResults(test_case.get());
}
// A hanging all-reduce.
ASSERT_GT(num_nodes, 1);
std::unique_ptr<typename TestFixture::TestCase> test_case(
this->MakeReductionTestCase(
/* num_nodes */ num_nodes, /* num_ranks_per_node */ 1, reduction_op,
TensorShape({2, 3}), 0.0f));
node_fn(test_case.get(), 0, communicator_key);
Env::Default()->SleepForMicroseconds(1000000);
for (auto& node : nodes) {
node.nccl_manager.StartAbort(absl::UnavailableError("peer down"));
}
{
mutex_lock l(test_case->mu);
while (test_case->num_completed != 1) {
test_case->done_cv.wait(l);
}
}
// Reset the aborted NcclManager and then run another all-reduce with the
// resetted NcclManagers.
for (auto& node : nodes) {
node.nccl_manager.Reset();
}
// Regenerate the communicator_key, because this is needed to create new
// communicators.
communicator_key = nodes[0].nccl_manager.GenerateCommunicatorKey();
{
std::unique_ptr<typename TestFixture::TestCase> test_case(
this->MakeReductionTestCase(
/* num_nodes */ num_nodes, /* num_ranks_per_node */ 1, reduction_op,
TensorShape({2, 3}), 0.0f));
for (int i = 0; i < num_nodes; ++i) {
this->work_queue_->Schedule(
[&node_fn, &test_case, i, communicator_key]() {
node_fn(test_case.get(), i, communicator_key);
});
}
this->VerifyResults(test_case.get());
}
}
#endif
} // namespace tensorflow
#endif // GOOGLE_CUDA || TENSORFLOW_USE_ROCM
tensorflow/core/nccl/nccl_rewrite.cc
+284
View File
@@ -0,0 +1,284 @@
/* Copyright 2017 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include <string>
#include "absl/status/status.h"
#include "absl/strings/match.h"
#include "absl/strings/str_cat.h"
#include "absl/strings/string_view.h"
#include "tensorflow/core/framework/types.pb.h"
#include "tensorflow/core/lib/strings/str_util.h"
#if GOOGLE_CUDA
#include <forward_list>
#include <vector>
#include "tensorflow/core/common_runtime/optimization_registry.h"
#include "tensorflow/core/framework/tensor.pb.h"
#include "tensorflow/core/graph/node_builder.h"
namespace tensorflow {
namespace {
// Replaces NcclReduce node with _NcclReduceRecv reusing one input of same
// device, adds one _NcclReduceSend for each other input.
absl::Status ReplaceReduce(Graph* graph, Node* node) {
std::string reduction;
TF_RETURN_IF_ERROR(GetNodeAttr(node->attrs(), "reduction", &reduction));
DataType dtype;
TF_RETURN_IF_ERROR(GetNodeAttr(node->attrs(), "T", &dtype));
int num_devices = node->num_inputs();
std::string shared_name = node->name();
auto make_builder = [&](absl::string_view op_name, absl::string_view suffix) {
return NodeBuilder(absl::StrCat(shared_name, suffix), op_name)
.Attr("reduction", reduction)
.Attr("num_devices", num_devices)
.Attr("shared_name", shared_name)
.Attr("T", dtype);
};
std::vector<Node*> control_inputs;
for (const auto& edge : node->in_edges()) {
if (edge->IsControlEdge()) {
control_inputs.push_back(edge->src());
}
}
std::vector<NodeBuilder::NodeOut> out_nodes;
for (const auto& edge : node->out_edges()) {
out_nodes.emplace_back(edge->dst(), edge->dst_input());
}
int recv_dev = node->assigned_device_name_index();
NodeBuilder recv_builder =
make_builder("_NcclReduceRecv", "Recv").ControlInputs(control_inputs);
bool recv_input_set = false;
int send_counter = 0;
for (const auto& edge : node->in_edges()) {
Node* src_node = edge->src();
if (edge->IsControlEdge()) {
continue;
}
int send_dev = src_node->assigned_device_name_index();
if (!recv_input_set && send_dev == recv_dev) {
recv_builder.Input(src_node);
recv_input_set = true;
continue;
}
auto send_builder =
make_builder("_NcclReduceSend", absl::StrCat("Send_", ++send_counter))
.Input(src_node)
.ControlInputs(control_inputs);
Node* send_node = nullptr;
TF_RETURN_IF_ERROR(send_builder.Finalize(graph, &send_node));
send_node->set_assigned_device_name_index(send_dev);
// Send nodes don't have any outputs and therefore have no data dependencies
// to the outputs of the graph. We add a control dependency to the receive
// node so that those 'dangling' nodes are run.
// TODO(b/67027412): Avoid these cross-device control edges.
for (const auto& out_node : out_nodes) {
graph->AddControlEdge(send_node, out_node.node);
}
}
if (!recv_input_set) {
return absl::InvalidArgumentError(
"No input tensor uses the same device as the NcclReduce op");
}
Node* recv_node = nullptr;
TF_RETURN_IF_ERROR(recv_builder.Finalize(graph, &recv_node));
recv_node->set_assigned_device_name_index(recv_dev);
graph->RemoveNode(node);
for (const auto& out_node : out_nodes) {
if (out_node.index == Graph::kControlSlot) {
graph->AddControlEdge(recv_node, out_node.node);
} else {
graph->AddEdge(recv_node, 0, out_node.node, out_node.index);
}
}
return absl::OkStatus();
}
TensorProto TensorFromShape(const TensorShapeProto& shape) {
TensorProto result;
result.set_dtype(DT_INT32);
for (const auto& dim : shape.dim()) {
result.add_int_val(dim.size());
}
result.mutable_tensor_shape()->add_dim()->set_size(shape.dim_size());
return result;
}
// Replaces NcclBroadcast node with _NcclBroadcastSend, connects the input to
// all outputs of same device, adds one _NcclBroadcastRecv for each other output
// device.
absl::Status ReplaceBroadcast(Graph* graph, Node* node) {
DataType dtype;
TF_RETURN_IF_ERROR(GetNodeAttr(node->attrs(), "T", &dtype));
int send_dev = node->assigned_device_name_index();
int num_devices = 0; // Number of distinct devices, incremented below.
std::vector<int> recv_index_map; // Map device name index to stable index.
// Map device name index to nodes that take the broadcast as input.
std::vector<std::forward_list<NodeBuilder::NodeOut>> out_nodes_map;
for (const auto& edge : node->out_edges()) {
int dst_dev = edge->IsControlEdge()
? send_dev
: edge->dst()->assigned_device_name_index();
if (out_nodes_map.size() <= dst_dev) {
out_nodes_map.resize(dst_dev + 1);
recv_index_map.resize(dst_dev + 1);
}
auto it = out_nodes_map.begin() + dst_dev;
if (it->empty()) {
recv_index_map[dst_dev] = num_devices;
++num_devices;
}
it->emplace_front(NodeBuilder::NodeOut(edge->dst(), edge->dst_input()));
}
if (num_devices <= 1) {
// Only one participating device, skip NCCL op.
const Edge* in_edge = nullptr;
TF_RETURN_IF_ERROR(node->input_edge(0, &in_edge));
Node* in_node = in_edge->src();
int in_index = in_edge->src_output();
graph->RemoveNode(node);
for (const auto& out_nodes : out_nodes_map) {
for (const auto& out_node : out_nodes) {
if (out_node.index == Graph::kControlSlot) {
graph->AddControlEdge(in_node, out_node.node);
} else {
graph->AddEdge(in_node, in_index, out_node.node, out_node.index);
}
}
}
return absl::OkStatus();
}
std::string shared_name = node->name();
auto make_builder = [&](absl::string_view op_name, absl::string_view suffix) {
return NodeBuilder(absl::StrCat(shared_name, suffix), op_name)
.Attr("num_devices", num_devices)
.Attr("shared_name", shared_name)
.Attr("T", dtype);
};
// Create broadcast send node and replace the original broadcast node.
NodeBuilder::NodeOut in_node;
NodeBuilder send_builder = make_builder("_NcclBroadcastSend", "Send");
for (const auto& edge : node->in_edges()) {
if (edge->IsControlEdge()) {
send_builder.ControlInput(edge->src());
} else {
in_node = NodeBuilder::NodeOut(edge->src(), edge->src_output());
send_builder.Input(in_node);
}
}
Node* send_node = nullptr;
TF_RETURN_IF_ERROR(send_builder.Finalize(graph, &send_node));
send_node->set_assigned_device_name_index(send_dev);
TensorShapeProto shape_proto;
TF_RETURN_IF_ERROR(GetNodeAttr(node->attrs(), "shape", &shape_proto));
// Delete the original node before reconnecting to outputs.
graph->RemoveNode(node);
// Connect all outputs on the device of broadcast send.
for (const auto& out_node : out_nodes_map[send_dev]) {
if (out_node.index == Graph::kControlSlot) {
graph->AddControlEdge(send_node, out_node.node);
} else {
graph->AddEdge(in_node.node, in_node.index, out_node.node,
out_node.index);
// Add control edge so send node is run.
graph->AddControlEdge(send_node, out_node.node);
}
}
out_nodes_map[send_dev].clear();
TensorProto tensor_proto = TensorFromShape(shape_proto);
bool is_fully_defined = TensorShape(shape_proto).IsFullyDefined();
std::string shape_name = absl::StrCat(in_node.node->name(), "/Shape");
Node* shape_node = nullptr;
if (!is_fully_defined) {
NodeBuilder shape_builder(shape_name, "Shape");
shape_builder.Input(in_node).Attr("out_type", DT_INT32).Attr("T", dtype);
TF_RETURN_IF_ERROR(shape_builder.Finalize(graph, &shape_node));
shape_node->set_assigned_device_name_index(send_dev);
}
// For all other devices, create a broadcast receive and connect outputs.
for (int recv_dev = 0; recv_dev < out_nodes_map.size(); ++recv_dev) {
if (out_nodes_map[recv_dev].empty()) {
continue;
}
int recv_index = recv_index_map[recv_dev];
if (is_fully_defined) {
// If the shape is fully defined, define one const node per device.
NodeBuilder shape_builder(absl::StrCat(shape_name, recv_index), "Const");
shape_builder.Attr("value", tensor_proto).Attr("dtype", DT_INT32);
TF_RETURN_IF_ERROR(shape_builder.Finalize(graph, &shape_node));
shape_node->set_assigned_device_name_index(recv_dev);
}
Node* recv_node;
TF_RETURN_IF_ERROR(
make_builder("_NcclBroadcastRecv", absl::StrCat("Recv_", recv_index))
.Input(shape_node)
.Finalize(graph, &recv_node));
recv_node->set_assigned_device_name_index(recv_dev);
for (const auto& out_node : out_nodes_map[recv_dev]) {
graph->AddEdge(recv_node, 0, out_node.node, out_node.index);
}
}
return absl::OkStatus();
}
// Replaces occurrences of Nccl{Reduce, Broadcast}Input/Output with their
// _Nccl...Send/Recv counterparts and removes data dependencies between them.
class NcclReplacePass : public GraphOptimizationPass {
public:
absl::Status Run(const GraphOptimizationPassOptions& options) override {
if (options.graph == nullptr) {
return absl::OkStatus();
}
Graph* graph = options.graph->get();
if (graph == nullptr) {
return absl::InternalError(
"NCCL replacement should happen before partitioning and a "
"graph should be available.");
}
// Find reduction and broadcast ops and replace them with Send/Recv ops.
for (Node* node : graph->op_nodes()) {
absl::string_view type = node->type_string();
if (!absl::StartsWith(type, "Nccl")) {
continue;
}
if (type == "NcclReduce") {
TF_RETURN_IF_ERROR(ReplaceReduce(graph, node));
}
if (type == "NcclBroadcast") {
TF_RETURN_IF_ERROR(ReplaceBroadcast(graph, node));
}
}
return absl::OkStatus();
}
};
REGISTER_OPTIMIZATION(OptimizationPassRegistry::POST_PLACEMENT, 0,
NcclReplacePass);
} // namespace
} // namespace tensorflow
#endif // GOOGLE_CUDA