123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244 |
- # Copyright (c) Microsoft Corporation.
- # SPDX-License-Identifier: Apache-2.0
- # DeepSpeed Team
- import abc
- from abc import ABC
- class DeepSpeedAccelerator(ABC):
- def __init__(self):
- self._name = None
- self._communication_backend_name = None
- @abc.abstractmethod
- def is_synchronized_device(self):
- ...
- # Device APIs
- @abc.abstractmethod
- def device_name(self, device_index):
- ...
- @abc.abstractmethod
- def device(self, device_index):
- ...
- @abc.abstractmethod
- def set_device(self, device_index):
- ...
- @abc.abstractmethod
- def current_device(self):
- ...
- @abc.abstractmethod
- def current_device_name(self):
- ...
- @abc.abstractmethod
- def device_count(self):
- ...
- @abc.abstractmethod
- def synchronize(self, device_index=None):
- ...
- # RNG APIs
- @abc.abstractmethod
- def random(self):
- ...
- @abc.abstractmethod
- def set_rng_state(self, new_state, device_index=None):
- ...
- @abc.abstractmethod
- def get_rng_state(self, device_index=None):
- ...
- @abc.abstractmethod
- def manual_seed(self, seed):
- ...
- @abc.abstractmethod
- def manual_seed_all(self, seed):
- ...
- @abc.abstractmethod
- def initial_seed(self, seed):
- ...
- @abc.abstractmethod
- def default_generator(self, device_index):
- ...
- # Streams/Events
- @property
- @abc.abstractmethod
- def Stream(self):
- ...
- @abc.abstractmethod
- def stream(self, stream):
- ...
- @abc.abstractmethod
- def current_stream(self, device_index=None):
- ...
- @abc.abstractmethod
- def default_stream(self, device_index=None):
- ...
- @property
- @abc.abstractmethod
- def Event(self):
- ...
- # Memory management
- @abc.abstractmethod
- def empty_cache(self):
- ...
- @abc.abstractmethod
- def memory_allocated(self, device_index=None):
- ...
- @abc.abstractmethod
- def max_memory_allocated(self, device_index=None):
- ...
- @abc.abstractmethod
- def reset_max_memory_allocated(self, device_index=None):
- ...
- @abc.abstractmethod
- def memory_cached(self, device_index=None):
- ...
- @abc.abstractmethod
- def max_memory_cached(self, device_index=None):
- ...
- @abc.abstractmethod
- def reset_max_memory_cached(self, device_index=None):
- ...
- @abc.abstractmethod
- def memory_stats(self, device_index=None):
- ...
- @abc.abstractmethod
- def reset_peak_memory_stats(self, device_index=None):
- ...
- @abc.abstractmethod
- def memory_reserved(self, device_index=None):
- ...
- @abc.abstractmethod
- def max_memory_reserved(self, device_index=None):
- ...
- @abc.abstractmethod
- def total_memory(self, device_index=None):
- ...
- # Data types
- @abc.abstractmethod
- def is_bf16_supported(self):
- ...
- @abc.abstractmethod
- def is_fp16_supported(self):
- ...
- # Misc
- @abc.abstractmethod
- def amp(self):
- ...
- @abc.abstractmethod
- def is_available(self):
- ...
- @abc.abstractmethod
- def range_push(self, msg):
- ...
- @abc.abstractmethod
- def range_pop(self):
- ...
- @abc.abstractmethod
- def lazy_call(self, callback):
- ...
- @abc.abstractmethod
- def communication_backend_name(self):
- ...
- # Tensor operations
- @property
- @abc.abstractmethod
- def BFloat16Tensor(self):
- ...
- @property
- @abc.abstractmethod
- def ByteTensor(self):
- ...
- @property
- @abc.abstractmethod
- def DoubleTensor(self):
- ...
- @property
- @abc.abstractmethod
- def FloatTensor(self):
- ...
- @property
- @abc.abstractmethod
- def HalfTensor(self):
- ...
- @property
- @abc.abstractmethod
- def IntTensor(self):
- ...
- @property
- @abc.abstractmethod
- def LongTensor(self):
- ...
- @abc.abstractmethod
- def pin_memory(self, tensor):
- ...
- @abc.abstractmethod
- def on_accelerator(self, tensor):
- ...
- @abc.abstractmethod
- def op_builder_dir(self):
- ...
- # create an instance of op builder, specified by class_name
- @abc.abstractmethod
- def create_op_builder(self, class_name):
- ...
- # return an op builder class, specified by class_name
- @abc.abstractmethod
- def get_op_builder(self, class_name):
- ...
- @abc.abstractmethod
- def build_extension(self):
- ...
|