1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
|
#pragma once
#include <ATen/ATen.h>
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDACachingAllocator.h>
#include <c10/util/Optional.h>
#include <cstddef>
#include <vector>
// NCCL BFloat16 is enabled only for CUDA 11+ and NCCL versions 2.10+, or for
// HIP 3.1+
#if defined(__CUDA_BF16_TYPES_EXIST__)
#define HAS_NCCL_BF16_DATATYPE \
((NCCL_MAJOR > 2) || (NCCL_MAJOR == 2) && (NCCL_MINOR >= 10))
#elif defined(USE_ROCM) && (TORCH_HIP_VERSION >= 301)
#define HAS_NCCL_BF16_DATATYPE 1
#else
#define HAS_NCCL_BF16_DATATYPE 0
#endif
namespace torch {
namespace cuda {
namespace nccl {
/* The following are copied from <nccl.h> and redefined in torch::cuda::nccl
* namespace */
/* pytorch should only use the following definition within pytorch scope */
/* Opaque handle to communicator to ncclComm*, this will reinterpret as ncclComm
* in nccl.cpp */
typedef void* ncclComm_t;
/** redefine nccl unique ID in torch scope. this should be identical to native
* nccl impp. */
#define NCCL_UNIQUE_ID_BYTES 128
// NOLINTNEXTLINE(cppcoreguidelines-avoid-c-arrays,modernize-avoid-c-arrays)
typedef struct {
char internal[NCCL_UNIQUE_ID_BYTES];
} ncclUniqueId;
/* Error type */
enum class ncclResult {
Success = 0,
UnhandledCudaError = 1,
SystemError = 2,
InternalError = 3,
InvalidArgument = 4,
InvalidUsage = 5,
NumResults = 6
};
/* Reduction operation selector */
enum class ncclRedOp { Sum = 0, Prod = 1, Max = 2, Min = 3, NumOps = 4 };
/* Data types */
enum class ncclDataType {
Int8 = 0,
Char = 0,
Uint8 = 1,
Int32 = 2,
Int = 2,
Uint32 = 3,
Int64 = 4,
Uint64 = 5,
Float16 = 6,
Half = 6,
Float32 = 7,
Float = 7,
Float64 = 8,
Double = 8,
Bfloat16 = 9,
NumTypes = 10
};
// RAII helper class to manage NCCL group API and CUDA free mutex.
// The destructor is allowed to throw since this helper class only
// manages group and lock lifetimes.
struct AutoNcclGroup {
AutoNcclGroup();
~AutoNcclGroup() noexcept(false);
};
// NOTE: this is exposed only so that python_nccl.cpp can some of these helpers.
// Don't use them outside of these files.
namespace detail {
TORCH_CUDA_CPP_API void throw_nccl_error(ncclResult status);
static inline void NCCL_CHECK(ncclResult status) {
if (status != ncclResult::Success) {
throw_nccl_error(status);
}
}
TORCH_CUDA_CPP_API at::ArrayRef<ncclComm_t> get_communicators(
at::TensorList inputs);
TORCH_CUDA_CPP_API void check_inputs(
at::TensorList inputs,
at::TensorList outputs,
int input_multiplier,
int output_multiplier);
TORCH_CUDA_CPP_API void check_inputs(
at::TensorList inputs,
const at::Tensor& output,
int root,
int input_multiplier,
int output_multiplier);
} // namespace detail
using comm_list = std::vector<ncclComm_t>;
using stream_list = std::vector<c10::optional<at::cuda::CUDAStream>>;
TORCH_CUDA_CPP_API std::uint64_t version();
bool is_available(at::TensorList tensors);
TORCH_CUDA_CPP_API void get_unique_id(ncclUniqueId& id);
TORCH_CUDA_CPP_API ncclComm_t
comm_init_rank(int nranks, const ncclUniqueId& comm_id, int rank);
TORCH_CUDA_CPP_API void comm_destroy(ncclComm_t comm);
TORCH_CUDA_CPP_API void broadcast(
at::TensorList tensors,
const stream_list& streams = {},
const comm_list& user_comms = {});
size_t get_max_count();
TORCH_CUDA_CPP_API void reduce(
const std::vector<at::Tensor>& inputs,
at::Tensor& output,
int32_t root = 0,
int32_t op = static_cast<int>(ncclRedOp::Sum),
const stream_list& streams = {},
const comm_list& user_comms = {});
TORCH_CUDA_CPP_API void reduce(
std::vector<at::Tensor>& inputs,
int32_t root = 0,
int32_t op = static_cast<int>(ncclRedOp::Sum),
const stream_list& streams = {},
const comm_list& user_comms = {});
TORCH_CUDA_CPP_API void all_reduce(
const std::vector<at::Tensor>& inputs,
std::vector<at::Tensor>& outputs,
int32_t op = static_cast<int>(ncclRedOp::Sum),
const stream_list& streams = {},
const comm_list& user_comms = {});
TORCH_CUDA_CPP_API void reduce_scatter(
const std::vector<at::Tensor>& inputs,
std::vector<at::Tensor>& outputs,
int32_t op = static_cast<int>(ncclRedOp::Sum),
const stream_list& streams = {},
const comm_list& user_comms = {});
TORCH_CUDA_CPP_API void scatter(
const std::vector<at::Tensor>& inputs,
at::Tensor& outputs,
ncclComm_t comm,
at::cuda::CUDAStream& stream,
int32_t root = 0);
TORCH_CUDA_CPP_API void all_gather(
const std::vector<at::Tensor>& inputs,
std::vector<at::Tensor>& outputs,
const stream_list& streams = {},
const comm_list& user_comms = {});
TORCH_CUDA_CPP_API void gather(
const at::Tensor& inputs,
std::vector<at::Tensor>& outputs,
ncclComm_t comm,
at::cuda::CUDAStream& stream,
int32_t root = 0);
TORCH_CUDA_CPP_API void all2all_single_equal_split(
at::Tensor& input,
at::Tensor& output,
int size,
ncclComm_t comm,
at::cuda::CUDAStream& stream);
TORCH_CUDA_CPP_API void all2all_single_unequal_split(
void* sendbuff,
const size_t* sendcounts,
const size_t* senddispls,
void* recvbuff,
const size_t* recvcounts,
const size_t* recvdispls,
size_t size,
c10::ScalarType type,
ncclComm_t comm,
at::cuda::CUDAStream& stream);
TORCH_CUDA_CPP_API void all2all(
std::vector<at::Tensor>& outputTensors,
std::vector<at::Tensor>& inputTensors,
ncclComm_t _comm,
at::cuda::CUDAStream& stream);
TORCH_CUDA_CPP_API void send(
const at::Tensor& input,
ncclComm_t comm,
at::cuda::CUDAStream stream,
int dst);
TORCH_CUDA_CPP_API void recv(
at::Tensor& output,
ncclComm_t comm,
at::cuda::CUDAStream stream,
int src);
} // namespace nccl
} // namespace cuda
} // namespace torch
|
