DeepSpeed

1
# Copyright (c) Microsoft Corporation.
2
# SPDX-License-Identifier: Apache-2.0
3

4
# DeepSpeed Team
5

6
import abc
7
from abc import ABC
8

9

10
class DeepSpeedAccelerator(ABC):
11

12
    def __init__(self):
13
        self._name = None
14
        self._communication_backend_name = None
15
        self._compile_backend = None
16

17
    @abc.abstractmethod
18
    def is_synchronized_device(self):
19
        ...
20

21
    @abc.abstractmethod
22
    def use_host_timers(self):
23
        ...
24

25
    @abc.abstractmethod
26
    def resolves_data_dependency(self):
27
        ...
28

29
    @abc.abstractmethod
30
    def handles_memory_backpressure(self):
31
        ...
32

33
    # Device APIs
34
    @abc.abstractmethod
35
    def device_name(self, device_index):
36
        ...
37

38
    @abc.abstractmethod
39
    def device(self, device_index):
40
        ...
41

42
    @abc.abstractmethod
43
    def set_device(self, device_index):
44
        ...
45

46
    @abc.abstractmethod
47
    def current_device(self):
48
        ...
49

50
    @abc.abstractmethod
51
    def current_device_name(self):
52
        ...
53

54
    @abc.abstractmethod
55
    def device_count(self):
56
        ...
57

58
    @abc.abstractmethod
59
    def synchronize(self, device_index=None):
60
        ...
61

62
    # RNG APIs
63
    @abc.abstractmethod
64
    def random(self):
65
        ...
66

67
    @abc.abstractmethod
68
    def set_rng_state(self, new_state, device_index=None):
69
        ...
70

71
    @abc.abstractmethod
72
    def get_rng_state(self, device_index=None):
73
        ...
74

75
    @abc.abstractmethod
76
    def manual_seed(self, seed):
77
        ...
78

79
    @abc.abstractmethod
80
    def manual_seed_all(self, seed):
81
        ...
82

83
    @abc.abstractmethod
84
    def initial_seed(self):
85
        ...
86

87
    @abc.abstractmethod
88
    def default_generator(self, device_index):
89
        ...
90

91
    # Streams/Events
92
    @property
93
    @abc.abstractmethod
94
    def Stream(self):
95
        ...
96

97
    @abc.abstractmethod
98
    def stream(self, stream):
99
        ...
100

101
    @abc.abstractmethod
102
    def current_stream(self, device_index=None):
103
        ...
104

105
    @abc.abstractmethod
106
    def default_stream(self, device_index=None):
107
        ...
108

109
    @property
110
    @abc.abstractmethod
111
    def Event(self):
112
        ...
113

114
    # Memory management
115
    @abc.abstractmethod
116
    def empty_cache(self):
117
        ...
118

119
    @abc.abstractmethod
120
    def memory_allocated(self, device_index=None):
121
        ...
122

123
    @abc.abstractmethod
124
    def max_memory_allocated(self, device_index=None):
125
        ...
126

127
    @abc.abstractmethod
128
    def reset_max_memory_allocated(self, device_index=None):
129
        ...
130

131
    @abc.abstractmethod
132
    def memory_cached(self, device_index=None):
133
        ...
134

135
    @abc.abstractmethod
136
    def max_memory_cached(self, device_index=None):
137
        ...
138

139
    @abc.abstractmethod
140
    def reset_max_memory_cached(self, device_index=None):
141
        ...
142

143
    @abc.abstractmethod
144
    def memory_stats(self, device_index=None):
145
        ...
146

147
    @abc.abstractmethod
148
    def reset_peak_memory_stats(self, device_index=None):
149
        ...
150

151
    @abc.abstractmethod
152
    def memory_reserved(self, device_index=None):
153
        ...
154

155
    @abc.abstractmethod
156
    def max_memory_reserved(self, device_index=None):
157
        ...
158

159
    @abc.abstractmethod
160
    def total_memory(self, device_index=None):
161
        ...
162

163
    @abc.abstractmethod
164
    def available_memory(self, device_index=None):
165
        ...
166

167
    # Data types
168
    @abc.abstractmethod
169
    def is_bf16_supported(self):
170
        ...
171

172
    @abc.abstractmethod
173
    def is_fp16_supported(self):
174
        ...
175

176
    @abc.abstractmethod
177
    def supported_dtypes(self):
178
        ...
179

180
    # Misc
181
    @abc.abstractmethod
182
    def amp(self):
183
        ...
184

185
    @abc.abstractmethod
186
    def is_available(self):
187
        ...
188

189
    @abc.abstractmethod
190
    def range_push(self, msg):
191
        ...
192

193
    @abc.abstractmethod
194
    def range_pop(self):
195
        ...
196

197
    @abc.abstractmethod
198
    def lazy_call(self, callback):
199
        ...
200

201
    @abc.abstractmethod
202
    def communication_backend_name(self):
203
        ...
204

205
    @abc.abstractmethod
206
    def is_triton_supported(self):
207
        ...
208

209
    # Graph operations
210
    @abc.abstractmethod
211
    def create_graph(self):
212
        ...
213

214
    @abc.abstractmethod
215
    def capture_to_graph(self, graph, pool=None, stream=None):
216
        ...
217

218
    @abc.abstractmethod
219
    def replay_graph(self, graph):
220
        ...
221

222
    # Tensor operations
223
    @property
224
    @abc.abstractmethod
225
    def BFloat16Tensor(self):
226
        ...
227

228
    @property
229
    @abc.abstractmethod
230
    def ByteTensor(self):
231
        ...
232

233
    @property
234
    @abc.abstractmethod
235
    def DoubleTensor(self):
236
        ...
237

238
    @property
239
    @abc.abstractmethod
240
    def FloatTensor(self):
241
        ...
242

243
    @property
244
    @abc.abstractmethod
245
    def HalfTensor(self):
246
        ...
247

248
    @property
249
    @abc.abstractmethod
250
    def IntTensor(self):
251
        ...
252

253
    @property
254
    @abc.abstractmethod
255
    def LongTensor(self):
256
        ...
257

258
    @abc.abstractmethod
259
    def pin_memory(self, tensor, align_bytes=1):
260
        ...
261

262
    @abc.abstractmethod
263
    def is_pinned(self, tensor):
264
        ...
265

266
    @abc.abstractmethod
267
    def on_accelerator(self, tensor):
268
        ...
269

270
    @abc.abstractmethod
271
    def op_builder_dir(self):
272
        ...
273

274
    # create an instance of op builder, specified by class_name
275
    @abc.abstractmethod
276
    def create_op_builder(self, class_name):
277
        ...
278

279
    # return an op builder class, specified by class_name
280
    @abc.abstractmethod
281
    def get_op_builder(self, class_name):
282
        ...
283

284
    @abc.abstractmethod
285
    def build_extension(self):
286
        ...
287

288
    @abc.abstractmethod
289
    def export_envs(self):
290
        ...
291

292
    @abc.abstractmethod
293
    def visible_devices_envs(self):
294
        ...
295

296
    @abc.abstractmethod
297
    def set_visible_devices_envs(self, current_env, local_accelerator_ids):
298
        ...
299

300
    @abc.abstractmethod
301
    def get_compile_backend(self):
302
        ...
303

304
    @abc.abstractmethod
305
    def set_compile_backend(self, backend):
306
        ...
307