Repository: microsoft/DeepSpeed

Issue:
[BUG] Incorrect batch_size used for ThroughputTimer
**Describe the bug**
`ThroughputTimer` is started/stopped for model steps, but in `runtime/engine` it uses the micro_batch_size only:
[here](https://github.com/microsoft/DeepSpeed/blob/ee39187d8f07f7efc615f64affa8cc60ccf41eb5/deepspeed/runtime/engine.py#L297)

Compare the above line to the following that uses the right value in `runtime/pipe/engine`: [here](https://github.com/microsoft/DeepSpeed/blob/b8fb9c3f1a8a8e3e574a7d53e83ce2b72d471aa3/deepspeed/runtime/pipe/engine.py#L107)

**To Reproduce**
1. Enable `wall_clock_breakdown`

**Expected behavior**
This timer should use  `self.train_micro_batch_size_per_gpu() * self.gradient_accumulation_steps()`.

**ds_report output**
N/A

**Screenshots**
N/A

**System info (please complete the following information):**
 - OS: `Amazon Linux 2`
 - GPU count and types: `8 x A100`
 - Interconnects (if applicable): `4 x 100 Gbps`
 - Python version: `3.8`
 - Any other relevant info about your setup: 

**Launcher context**
Launching the experiment with the `deepspeed` launcher.

**Docker context**
Can't share

**Additional context**
DeepSpeed v0.7.3



Discussion:
Hi @ShijieZZZZ, can you please check my PR? There's a weird environment issue for one of the checks, while others run fine.
Hi @clumsy, checked tests are passed. 
Are you using [avg_sample_per_sec()](https://github.com/microsoft/DeepSpeed/blob/ee39187d8f07f7efc615f64affa8cc60ccf41eb5/deepspeed/utils/timer.py#L218) to compute throughput?
Hi @ShijieZZZZ, I'm using both `RunningAvg` and `Curr`, and both were askew because they use the same inaccurate `batch_size`.
On the second thought, I'm a bit confused now.

The current throughput number _could_ be consistent if the steps it refers to are really __micro steps__.

In the docs [here](https://www.deepspeed.ai/docs/config-json/) `steps_per_print` are referred to as "training steps". So my assumption is that it should be printed at gradient accumulation boundaries, after full step processing `train_batch_size` number samples.

If it's intended to be printed at micro steps I wonder how useful it is when the value fluctuates between intermediate micro steps and micro step at the gradient boundary where additional operations are executed slowing down the throughput.

I'll adjust my PR to report for full steps just in case this is how it was supposed to be.
So is my understanding correct, @ShijieZZZZ? Is the throughput timer supposed to report per full step, not micro step as it does currently? If yes, please check out the recent version of my PR? The unit tests get 5xx from HuggingFace though.
Hi @clumsy, I agree that throughput timer should report every gradient_accumulation_steps, I will take a look your PR.
@clumsy, please check this. https://github.com/microsoft/DeepSpeed/pull/2553 
Closing. Please re-open if needed. 
Sorry @tjruwase and @ShijieZZZZ, I was just checking the suggested fix and I believe the problem is still there, in part. While I believe the timer should now report throughput at model steps, `curr_samples_sec` is still not accurate. It is still effectively `(self.train_micro_batch_size_per_gpu() * self.num_workers) / duration` (does not take gradient accumulation (micro) steps into account [here](https://github.com/microsoft/DeepSpeed/blob/master/deepspeed/runtime/engine.py#L300)) which was the main goal of this issue.
Also we seem to ignore `steps_per_output` after the latest change. I found one problem with my original PR but I'll fix it and keep the tests.
I think the root of all issues here is the fact that `steps` inside the timer are really micro steps. So if we pass gradient accumulation number we can use it to both multiply micro batch size per GPU and mod the number of start/stop invocations to detect model steps. And once `steps_per_output` is a factor of the number of micro steps divided by gradient accumulation number - we can report the value.
@clumsy, thanks for the analysis. I agree that there are still issues to resolve. However, one issue that I notice with your proposal is that it will introduce the concept of gradient accumulation into the timer, which I think it unnecessary and makes the timer less generic. I would like a design where the timer remains simple and the client (i.e., the engine) incurs the complexity of using the timer to meet its needs. Also, for simplicity timer `start()` and `stop()` are called on micro steps, and so it is better for timer to be initialized with `micro_batch_size` so that the state is always consistent, e.g., for debugging. 

So, I am thinking of an alternative approach:

1. Rename timer fields for clarity, `local_step_count` as `micro_step_count`, and `total_step_count` as `global_step_count`.
2. Redefine timer stop as `def stop(self, global_step=False, report_speed=False)`
3. On gradient accumulation boundaries, client calls `timer.stop()` with `global_step=True`, and `report_speed=True` if rank 0, otherwise call both flags are False
4. timer stop() updates `micro_step_count` always, updates `global_step_count` if `global_step=True` and prints if `report_speed=True`

What do you think?

Resolving patch:
diff --git a/deepspeed/runtime/engine.py b/deepspeed/runtime/engine.py
--- a/deepspeed/runtime/engine.py
+++ b/deepspeed/runtime/engine.py
@@ -308,8 +308,7 @@ def __init__(
         self.timers = SynchronizedWallClockTimer()
         # Throughput timer
         self.tput_timer = ThroughputTimer(
-            batch_size=self.train_micro_batch_size_per_gpu(),
-            num_workers=self.dp_world_size,
+            batch_size=self.train_batch_size(),
             steps_per_output=self.steps_per_print(),
             monitor_memory=False,
         )
@@ -2168,7 +2167,8 @@ def step(self, lr_kwargs=None):
 
             report_progress = self.global_rank == 0 if self.global_rank else True
 
-        self.tput_timer.stop(report_progress)
+        self.tput_timer.stop(global_step=self.is_gradient_accumulation_boundary(),
+                             report_speed=report_progress)
 
         self._stop_timers(self.engine_timers.step_timers)
 
diff --git a/deepspeed/runtime/pipe/engine.py b/deepspeed/runtime/pipe/engine.py
--- a/deepspeed/runtime/pipe/engine.py
+++ b/deepspeed/runtime/pipe/engine.py
@@ -104,9 +104,7 @@ def __init__(self, has_bool_tensors=False, *super_args, **super_kwargs):
 
         self._force_grad_boundary = False
 
-        self.batch_timer = ThroughputTimer(batch_size=self.micro_batch_size *
-                                           self.micro_batches,
-                                           num_workers=self.dp_world_size,
+        self.batch_timer = ThroughputTimer(batch_size=self.train_batch_size(),
                                            logging_fn=self.tput_log,
                                            monitor_memory=False,
                                            steps_per_output=self.steps_per_print())
diff --git a/deepspeed/utils/timer.py b/deepspeed/utils/timer.py
--- a/deepspeed/utils/timer.py
+++ b/deepspeed/utils/timer.py
@@ -136,7 +136,6 @@ class ThroughputTimer:
     def __init__(
         self,
         batch_size,
-        num_workers,
         start_step=2,
         steps_per_output=50,
         monitor_memory=False,
@@ -146,15 +145,13 @@ def __init__(
         self.start_time = 0
         self.end_time = 0
         self.started = False
-        self.batch_size = batch_size
-        if batch_size is None:
-            self.batch_size = 1
-        self.num_workers = num_workers
+        self.batch_size = 1 if batch_size is None else batch_size
         self.start_step = start_step
         self.epoch_count = 0
-        self.local_step_count = 0
-        self.total_step_count = 0
+        self.micro_step_count = 0
+        self.global_step_count = 0
         self.total_elapsed_time = 0
+        self.step_elapsed_time = 0
         self.steps_per_output = steps_per_output
         self.monitor_memory = monitor_memory
         self.logging = logging_fn
@@ -167,7 +164,7 @@ def __init__(
 
     def update_epoch_count(self):
         self.epoch_count += 1
-        self.local_step_count = 0
+        self.micro_step_count = 0
 
     def _init_timer(self):
         self.initialized = True
@@ -175,54 +172,60 @@ def _init_timer(self):
     def start(self):
         self._init_timer()
         self.started = True
-        if self.total_step_count >= self.start_step:
+        if self.global_step_count >= self.start_step:
             torch.cuda.synchronize()
             self.start_time = time.time()
 
-    def stop(self, report_speed=True):
+    def stop(self, global_step=False, report_speed=True):
         if not self.started:
             return
         self.started = False
-        self.total_step_count += 1
-        self.local_step_count += 1
-        if self.total_step_count > self.start_step:
+        self.micro_step_count += 1
+        if global_step:
+            self.global_step_count += 1
+
+        if self.start_time > 0:
             torch.cuda.synchronize()
             self.end_time = time.time()
             duration = self.end_time - self.start_time
             self.total_elapsed_time += duration
-
-            curr_samples_sec = (self.batch_size * self.num_workers) / duration
-
-            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,
-                        round(torch.cuda.memory_allocated() / 1024**3,
-                              2),
-                        round(torch.cuda.max_memory_allocated() / 1024**3,
-                              2),
-                    ))
-                if self.monitor_memory:
-                    virt_mem = psutil.virtual_memory()
-                    swap = psutil.swap_memory()
-                    self.logging("{}/{}, vm percent: {}, swap percent: {}".format(
-                        self.epoch_count,
-                        self.local_step_count,
-                        virt_mem.percent,
-                        swap.percent,
-                    ))
+            self.step_elapsed_time += duration
+
+            if global_step:
+                if report_speed and self.global_step_count % self.steps_per_output == 0:
+                    self.logging(
+                        "epoch={}/micro_step={}/global_step={}, RunningAvgSamplesPerSec={}, CurrSamplesPerSec={}, "
+                        "MemAllocated={}GB, MaxMemAllocated={}GB".format(
+                            self.epoch_count,
+                            self.micro_step_count,
+                            self.global_step_count,
+                            self.avg_samples_per_sec(),
+                            self.batch_size / self.step_elapsed_time,
+                            round(torch.cuda.memory_allocated() / 1024**3,
+                                  2),
+                            round(torch.cuda.max_memory_allocated() / 1024**3,
+                                  2),
+                        ))
+                    if self.monitor_memory:
+                        virt_mem = psutil.virtual_memory()
+                        swap = psutil.swap_memory()
+                        self.logging(
+                            "epoch={}/micro_step={}/global_step={}, vm %: {}, swap %: {}"
+                            .format(
+                                self.epoch_count,
+                                self.micro_step_count,
+                                self.global_step_count,
+                                virt_mem.percent,
+                                swap.percent,
+                            ))
+                self.step_elapsed_time = 0
 
     def avg_samples_per_sec(self):
-        if self.total_step_count > 0:
-            samples_per_step = self.batch_size * self.num_workers
-            total_step_offset = self.total_step_count - self.start_step
+        if self.global_step_count > 0:
+            total_step_offset = self.global_step_count - self.start_step
             avg_time_per_step = self.total_elapsed_time / total_step_offset
             # training samples per second
-            return samples_per_step / avg_time_per_step
+            return self.batch_size / avg_time_per_step
         return float("-inf")
 
 


Test patch:
diff --git a/tests/unit/runtime/zero/test_zero_context.py b/tests/unit/runtime/zero/test_zero_context.py
--- a/tests/unit/runtime/zero/test_zero_context.py
+++ b/tests/unit/runtime/zero/test_zero_context.py
@@ -9,6 +9,7 @@
 import deepspeed.comm as dist
 
 from unit.common import DistributedTest, get_master_port
+from unit.simple_model import SimpleModel
 
 
 def setup_serial_env():
@@ -89,6 +90,95 @@ def test(self):
 }
 
 
+def test_throughput_calculation():
+    setup_serial_env()
+
+    train_micro_batch_size_per_gpu = 7
+    gradient_accumulation_steps = 6
+    config_dict = {
+        "train_micro_batch_size_per_gpu": train_micro_batch_size_per_gpu,
+        "gradient_accumulation_steps": gradient_accumulation_steps,
+        "optimizer": {
+            "type": "Adam",
+            "params": {
+                "lr": 0.001,
+            }
+        },
+        "zero_optimization": {
+            "stage": 0
+        },
+    }
+
+    args = SimpleNamespace(local_rank=0)
+    net = SimpleModel(hidden_dim=4)
+    engine, _, _, _ = deepspeed.initialize(args=args,
+                                           config=config_dict,
+                                           model=net,
+                                           model_parameters=net.parameters())
+    assert engine.tput_timer.batch_size == train_micro_batch_size_per_gpu * gradient_accumulation_steps
+
+    assert not engine.tput_timer.initialized
+    assert not engine.tput_timer.started
+    assert engine.tput_timer.start_step == 2
+    assert engine.tput_timer.start_time == 0
+    assert engine.tput_timer.micro_step_count == 0
+    assert engine.tput_timer.global_step_count == 0
+    assert engine.tput_timer.total_elapsed_time == 0
+
+    # calling stop() while uninitialized - has no effect
+    engine.tput_timer.stop()
+    assert not engine.tput_timer.initialized
+    assert not engine.tput_timer.started
+    assert engine.tput_timer.start_time == 0
+    assert engine.tput_timer.micro_step_count == 0
+    assert engine.tput_timer.global_step_count == 0
+    assert engine.tput_timer.total_elapsed_time == 0
+
+    # any call to start() (from dataloader or not) initializes the timer
+    engine.tput_timer.start()
+    assert engine.tput_timer.initialized
+    assert engine.tput_timer.started
+    assert engine.tput_timer.start_time == 0
+    assert engine.tput_timer.micro_step_count == 0
+    assert engine.tput_timer.global_step_count == 0
+    assert engine.tput_timer.total_elapsed_time == 0
+
+    # calling stop() after initialized - increments the local micro step counter
+    engine.tput_timer.stop()
+    assert engine.tput_timer.initialized
+    assert not engine.tput_timer.started
+    assert engine.tput_timer.start_time == 0
+    assert engine.tput_timer.micro_step_count == 1
+    assert engine.tput_timer.global_step_count == 0
+    assert engine.tput_timer.total_elapsed_time == 0
+
+    # calling start()/stop() to increment the step counter until start_step
+    while engine.tput_timer.micro_step_count < (gradient_accumulation_steps *
+                                                engine.tput_timer.start_step):
+        engine.tput_timer.start()
+        global_step = (engine.tput_timer.micro_step_count +
+                       1) % gradient_accumulation_steps == 0
+        engine.tput_timer.stop(global_step=global_step)
+    assert engine.tput_timer.global_step_count == engine.tput_timer.start_step
+    assert engine.tput_timer.total_elapsed_time == 0
+
+    # calling start()/stop() accumulates duration during gradient accumulation
+    while engine.tput_timer.global_step_count == engine.tput_timer.start_step:
+        engine.tput_timer.start()
+        current_duration = engine.tput_timer.step_elapsed_time
+        total_duration = engine.tput_timer.total_elapsed_time
+
+        global_step = (engine.tput_timer.micro_step_count +
+                       1) % gradient_accumulation_steps == 0
+        engine.tput_timer.stop(global_step=global_step)
+        duration = engine.tput_timer.end_time - engine.tput_timer.start_time
+        # step elapsed time is reset after gradient accumulation steps
+        assert engine.tput_timer.step_elapsed_time == (
+            0 if engine.tput_timer.global_step_count != engine.tput_timer.start_step else
+            current_duration + duration)
+        assert engine.tput_timer.total_elapsed_time == total_duration + duration
+
+
 def test_ext_param_getattr():
     setup_serial_env()