Merge branch 'up-master' into gma/device-abstraction

deepspeed/accelerator/abstract_accelerator.py

@@ -6,13 +6,6 @@ class DeepSpeedAccelerator(ABC):
     def __init__(self):
         self._name = None
         self._communication_backend_name = None
-        self.BFloat16Tensor = None
-        self.ByteTensor = None
-        self.DoubleTensor = None
-        self.FloatTensor = None
-        self.HalfTensor = None
-        self.IntTensor = None
-        self.LongTensor = None
 
     # Device APIs
     @abc.abstractmethod
@@ -24,7 +17,7 @@ def device(self, device_index):
         ...
 
     @abc.abstractmethod
-    def set_device(self):
+    def set_device(self, device_index):
         ...
 
     @abc.abstractmethod
@@ -65,7 +58,7 @@ def manual_seed_all(self, seed):
         ...
 
     @abc.abstractmethod
-    def initial_seed(self):
+    def initial_seed(self, seed):
         ...
 
     @abc.abstractmethod
@@ -181,6 +174,41 @@ 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):
         ...
@@ -189,6 +217,10 @@ def pin_memory(self, tensor):
     def on_accelerator(self, tensor):
         ...
 
+    @abc.abstractmethod
+    def op_builder_dir(self):
+        ...
+
     @abc.abstractmethod
     def create_op_builder(self, class_name):
         ...

deepspeed/accelerator/cuda_accelerator.py

@@ -6,13 +6,6 @@ class CUDA_Accelerator(DeepSpeedAccelerator):
     def __init__(self):
         self._name = 'cuda'
         self._communication_backend_name = 'nccl'
-        self.DoubleTensor = torch.cuda.DoubleTensor
-        self.LongTensor = torch.cuda.LongTensor
-        self.FloatTensor = torch.cuda.FloatTensor
-        self.BFloat16Tensor = torch.cuda.BFloat16Tensor
-        self.HalfTensor = torch.cuda.HalfTensor
-        self.IntTensor = torch.cuda.IntTensor
-        self.ByteTensor = torch.cuda.ByteTensor
 
     # Device APIs
     def device_name(self, device_index=None):
@@ -161,6 +154,35 @@ def communication_backend_name(self):
         return self._communication_backend_name
 
     # Tensor operations
+
+    @property
+    def BFloat16Tensor(self):
+        return torch.cuda.BFloat16Tensor
+
+    @property
+    def ByteTensor(self):
+        return torch.cuda.ByteTensor
+
+    @property
+    def DoubleTensor(self):
+        return torch.cuda.DoubleTensor
+
+    @property
+    def FloatTensor(self):
+        return torch.cuda.FloatTensor
+
+    @property
+    def HalfTensor(self):
+        return torch.cuda.HalfTensor
+
+    @property
+    def IntTensor(self):
+        return torch.cuda.IntTensor
+
+    @property
+    def LongTensor(self):
+        return torch.cuda.LongTensor
+
     def pin_memory(self, tensor):
         return tensor.pin_memory()
 
@@ -171,6 +193,9 @@ def on_accelerator(self, tensor):
         else:
             return False
 
+    def op_builder_dir(self):
+        return "deepspeed.ops.op_builder"
+
     def create_op_builder(self, class_name):
         from deepspeed.ops.op_builder import AsyncIOBuilder, CPUAdagradBuilder, CPUAdamBuilder, FusedAdamBuilder, FusedLambBuilder, QuantizerBuilder, SparseAttnBuilder, StochasticTransformerBuilder, TransformerBuilder, InferenceBuilder, UtilsBuilder
         from deepspeed.ops.op_builder.builder_names import AsyncIOBuilder as AsyncIOBuilderName

deepspeed/inference/engine.py

@@ -82,7 +82,7 @@ def __init__(self, model, config):
         # This is a hack to remove the prepare_mask function on HF side for BLOOM architecture
         self.remove_mask_prepare_for_bloom()
 
-        if get_accelerator().device_name() == 'cuda' and enable_cuda_graph:
+        if get_accelerator().device_name() == 'cuda' and config.enable_cuda_graph:
             assert pkg_version.parse(torch.__version__) >= pkg_version.parse("1.10"), \
                 "If you want to use cuda graph, please upgrade torch to at least v1.10"
 
@@ -127,7 +127,7 @@ def __init__(self, model, config):
         device = get_accelerator().current_device_name()
         self.module.to(device)
 
-        if self.mp_world_size > 1:
+        if config.tensor_parallel.tp_size > 1:
             _rng_state = get_accelerator().get_rng_state().to(
                 get_accelerator().current_device_name())
             dist.broadcast(_rng_state, 0)
@@ -514,11 +514,11 @@ def forward(self, *inputs, **kwargs):
         """
         start = None
         if self.model_profile_enabled and get_accelerator().device_name(
-        ) == 'cuda' and self.enable_cuda_graph:
+        ) == 'cuda' and self._config.enable_cuda_graph:
             get_accelerator().synchronize()
             start = time.time()
 
-        if get_accelerator().device_name() == 'cuda' and self.enable_cuda_graph:
+        if get_accelerator().device_name() == 'cuda' and self._config.enable_cuda_graph:
             if self.cuda_graph_created:
                 outputs = self._graph_replay(*inputs, **kwargs)
             else:
@@ -527,7 +527,7 @@ def forward(self, *inputs, **kwargs):
         else:
             outputs = self.module(*inputs, **kwargs)
 
-        if self.model_profile_enabled and self.enable_cuda_graph:
+        if self.model_profile_enabled and self._config.enable_cuda_graph:
             get_accelerator().synchronize()
             duration = time.time() - start
             self._model_times.append(duration)

deepspeed/module_inject/replace_module.py

@@ -608,7 +608,7 @@ def _transpose(x):
                     get_accelerator().current_device_name())
                 new_module.attn_nb.data = attn_nb.to(
                     get_accelerator().current_device_name())
-                if moe_type == 'residual':
+                if config.moe_type == 'residual':
                     new_module.res_mlp.inter_w.data = _res_h4h_w.to(
                         get_accelerator().current_device_name())
                     new_module.res_mlp.inter_b.data = _res_h4h_b.to(

deepspeed/runtime/engine.py

@@ -81,6 +81,8 @@
 from deepspeed.accelerator import get_accelerator
 from deepspeed.ops.op_builder.builder_names import UtilsBuilder
 
+from deepspeed.inference.config import DtypeEnum
+
 # Set to torch's distributed package or deepspeed.comm based inside DeepSpeedEngine init
 dist = None
 

deepspeed/utils/timer.py

@@ -195,12 +195,10 @@ def stop(self, report_speed=True):
 
             curr_samples_sec = (self.batch_size * self.num_workers) / duration
 
-            if self.local_step_count % self.steps_per_output == 0:
-                if report_speed:
-                    self.logging(
-                        "{}/{}, RunningAvgSamplesPerSec={}, CurrSamplesPerSec={}, MemAllocated={}GB, MaxMemAllocated={}GB"
-                        .format(
-                            self.epoch_count,
+            if report_speed:
+                self.logging(
+                    "{}/{}, RunningAvgSamplesPerSec={}, CurrSamplesPerSec={}, MemAllocated={}GB, MaxMemAllocated={}GB"
+                    .format(self.epoch_count,
                             self.local_step_count,
                             self.avg_samples_per_sec(),
                             curr_samples_sec,

op_builder/all_ops.py

@@ -1,23 +1,26 @@
 """
 Copyright 2020 The Microsoft DeepSpeed Team
 """
+import os
+import pkgutil
+import importlib
 from deepspeed.accelerator import get_accelerator
-from deepspeed.ops.op_builder.builder_names import CPUAdamBuilder, CPUAdagradBuilder, FusedAdamBuilder, FusedLambBuilder, SparseAttnBuilder, TransformerBuilder, StochasticTransformerBuilder, AsyncIOBuilder, UtilsBuilder, QuantizerBuilder, InferenceBuilder, SpatialInferenceBuilder
 
-# TODO: infer this list instead of hard coded
 # List of all available ops
-__op_builders__ = [
-    get_accelerator().create_op_builder(CPUAdamBuilder),
-    get_accelerator().create_op_builder(CPUAdagradBuilder),
-    get_accelerator().create_op_builder(FusedAdamBuilder),
-    get_accelerator().create_op_builder(FusedLambBuilder),
-    get_accelerator().create_op_builder(SparseAttnBuilder),
-    get_accelerator().create_op_builder(TransformerBuilder),
-    get_accelerator().create_op_builder(StochasticTransformerBuilder),
-    get_accelerator().create_op_builder(AsyncIOBuilder),
-    get_accelerator().create_op_builder(UtilsBuilder),
-    get_accelerator().create_op_builder(QuantizerBuilder),
-    get_accelerator().create_op_builder(InferenceBuilder),
-    get_accelerator().create_op_builder(SpatialInferenceBuilder)
-]
+
+# reflect all builder names into __op_builders__
+op_builder_dir = get_accelerator().op_builder_dir()
+op_builder_module = importlib.import_module(op_builder_dir)
+__op_builders__ = []
+
+for _, module_name, _ in pkgutil.iter_modules([os.path.dirname(op_builder_module.__file__)]):
+    # avoid self references
+    if module_name != 'all_ops' and module_name != 'builder' and module_name != 'builder_names':
+        module = importlib.import_module("{}.{}".format(op_builder_dir, module_name))
+        for member_name in module.__dir__():
+            if member_name.endswith('Builder'):
+                # append builder to __op_builders__ list
+                builder = get_accelerator().create_op_builder(member_name)
+                __op_builders__.append(builder)
+
 ALL_OPS = {op.name: op for op in __op_builders__ if op is not None}

op_builder/builder_names.py

@@ -1,13 +1,25 @@
-CPUAdamBuilder = "CPUAdamBuilder"
-CPUAdagradBuilder = "CPUAdagradBuilder"
-FusedAdamBuilder = "FusedAdamBuilder"
-FusedLambBuilder = "FusedLambBuilder"
-SparseAttnBuilder = "SparseAttnBuilder"
-TransformerBuilder = "TransformerBuilder"
-StochasticTransformerBuilder = "StochasticTransformerBuilder"
-AsyncIOBuilder = "AsyncIOBuilder"
-UtilsBuilder = "UtilsBuilder"
-QuantizerBuilder = "QuantizerBuilder"
-InferenceBuilder = "InferenceBuilder"
-InferenceSpecializedBuilder = "InferenceSpecializedBuilder"
-SpatialInferenceBuilder = "SpatialInferenceBuilder"
+import sys
+import os
+import pkgutil
+import importlib
+
+# List of all available op builders from deepspeed op_builder
+
+op_builder_dir = "deepspeed.ops.op_builder"
+op_builder_module = importlib.import_module(op_builder_dir)
+__op_builders__ = []
+
+this_module = sys.modules[__name__]
+
+# reflect all builder names into variable definition such as 'TransformerBuilder = "TransformerBuilder"'
+for _, module_name, _ in pkgutil.iter_modules([os.path.dirname(op_builder_module.__file__)]):
+    # avoid self references
+    if module_name != 'all_ops' and module_name != 'builder' and module_name != 'builder_names':
+        module = importlib.import_module("{}.{}".format(op_builder_dir, module_name))
+        for member_name in module.__dir__():
+            if member_name.endswith(
+                    'Builder'
+            ) and member_name != "OpBuilder" and member_name != "CUDAOpBuilder":
+                # assign builder name to variable with same name
+                # the following is equivalent to i.e. TransformerBuilder = "TransformerBuilder"
+                this_module.__dict__[member_name] = member_name