Add throughput timer configuration

deepspeed/runtime/config.py

@@ -66,6 +66,8 @@
 from .data_pipeline.config import get_data_efficiency_enabled, get_data_efficiency_config, get_curriculum_enabled_legacy, get_curriculum_params_legacy
 from .data_pipeline.constants import *
 
+from ..utils.config import DeepSpeedThroughputTimerConfig
+
 TENSOR_CORE_ALIGN_SIZE = 8
 
 ADAGRAD_OPTIMIZER = 'adagrad'
@@ -911,6 +913,8 @@ def _initialize_params(self, param_dict):
 
         self.compile_config = get_compile_config(param_dict)
 
+        self.timers_config = DeepSpeedThroughputTimerConfig(param_dict)
+
     def _batch_assertion(self):
 
         train_batch = self.train_batch_size

deepspeed/runtime/engine.py

@@ -271,11 +271,10 @@ def __init__(self,
         # Configure wall clock timers
         self.timers = SynchronizedWallClockTimer()
         # Throughput timer
-        self.tput_timer = ThroughputTimer(
-            batch_size=self.train_batch_size(),
-            steps_per_output=self.steps_per_print(),
-            monitor_memory=False,
-        )
+        self.tput_timer = ThroughputTimer(self._config.timers_config,
+                                          batch_size=self.train_batch_size(),
+                                          steps_per_output=self.steps_per_print(),
+                                          monitor_memory=False)
 
         log_dist(f"DeepSpeed Flops Profiler Enabled: {self.flops_profiler_enabled()}", ranks=[0])
 

deepspeed/runtime/pipe/engine.py

@@ -119,7 +119,8 @@ def __init__(self, has_bool_tensors=False, *super_args, **super_kwargs):
 
         self._force_grad_boundary = False
 
-        self.batch_timer = ThroughputTimer(batch_size=self.train_batch_size(),
+        self.batch_timer = ThroughputTimer(self._config.timers_config,
+                                           batch_size=self.train_batch_size(),
                                            logging_fn=self.tput_log,
                                            monitor_memory=False,
                                            steps_per_output=self.steps_per_print())

deepspeed/utils/config.py

@@ -0,0 +1,28 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+from deepspeed.runtime.config_utils import get_scalar_param, DeepSpeedConfigObject
+from deepspeed.utils.constants import *
+
+
+class DeepSpeedThroughputTimerConfig(DeepSpeedConfigObject):
+
+    def __init__(self, param_dict):
+        super(DeepSpeedThroughputTimerConfig, self).__init__()
+
+        self.enabled = None
+        self.synchronized = None
+
+        timers_dict = {}
+        if param_dict and TIMERS in param_dict:
+            if TIMERS_THROUGHPUT in param_dict[TIMERS]:
+                timers_dict = param_dict[TIMERS][TIMERS_THROUGHPUT]
+
+        self._initialize(timers_dict)
+
+    def _initialize(self, param_dict):
+        self.enabled = get_scalar_param(param_dict, TIMERS_THROUGHPUT_ENABLED, TIMERS_THROUGHPUT_ENABLED_DEFAULT)
+        self.synchronized = get_scalar_param(param_dict, TIMERS_THROUGHPUT_SYNCHRONIZED,
+                                             TIMERS_THROUGHPUT_SYNCHRONIZED_DEFAULT)

deepspeed/utils/constants.py

@@ -0,0 +1,33 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+#########################################
+# Timers
+#########################################
+''' Elasticity Utility in DeepSpeed can be used to create highly elastic jobs compatible
+with a large number of GPUs. For elastic jobs, DeepSpeed will provide a batch size that
+can support a large number of GPUs based on the user specified parameters
+'''
+TIMERS_FORMAT = '''
+Timers should be enabled as:
+"timers": {
+  "throughput": {
+    "enabled": true,
+    "synchronized": true
+  }
+}
+'''
+
+TIMERS = "timers"
+TIMERS_THROUGHPUT = "throughput"
+
+TIMERS_THROUGHPUT_ENABLED = "enabled"
+TIMERS_THROUGHPUT_ENABLED_DEFAULT = True
+
+# Synchronizing a device is required to produce the most accurate timer measurements.
+# However, this comes at the expense of performance degradation. The CPU timer provides
+# sufficient accuracy in many cases.
+TIMERS_THROUGHPUT_SYNCHRONIZED = "synchronized"
+TIMERS_THROUGHPUT_SYNCHRONIZED_DEFAULT = True

deepspeed/utils/timer.py

@@ -197,15 +197,9 @@ def get_mean(self, names, normalizer=1.0, reset=True):
 
 class ThroughputTimer:
 
-    def __init__(
-        self,
-        batch_size,
-        start_step=2,
-        steps_per_output=50,
-        monitor_memory=False,
-        logging_fn=None,
-    ):
+    def __init__(self, config, batch_size, start_step=2, steps_per_output=50, monitor_memory=False, logging_fn=None):
         from deepspeed.utils import logger
+        self.config = config
         self.start_time = 0
         self.end_time = 0
         self.started = False
@@ -234,22 +228,26 @@ def _init_timer(self):
         self.initialized = True
 
     def start(self):
+        if not self.config.enabled:
+            return
         self._init_timer()
         self.started = True
         if self.global_step_count >= self.start_step:
-            get_accelerator().synchronize()
+            if self.config.synchronized:
+                get_accelerator().synchronize()
             self.start_time = time.time()
 
     def stop(self, global_step=False, report_speed=True):
-        if not self.started:
+        if not self.config.enabled or not self.started:
             return
         self.started = False
         self.micro_step_count += 1
         if global_step:
             self.global_step_count += 1
 
         if self.start_time > 0:
-            get_accelerator().synchronize()
+            if self.config.synchronized:
+                get_accelerator().synchronize()
             self.end_time = time.time()
             duration = self.end_time - self.start_time
             self.total_elapsed_time += duration