feature: Record batch time

src/refiners/training_utils/clock.py

@@ -14,6 +14,26 @@
 from torch import Tensor
 
 
+# Ported from open-muse
+class AverageMeter(object):
+    """Computes and stores the average and current value"""
+
+    def __init__(self):
+        self.reset()
+
+    def reset(self):
+        self.val: float = 0
+        self.avg: float = 0
+        self.sum: float = 0
+        self.count: int = 0
+
+    def update(self, val: float):
+        self.val = val
+        self.sum += val
+        self.count += 1
+        self.avg = self.sum / self.count
+
+
 class ClockConfig(CallbackConfig):
     verbose: bool = True
 
@@ -49,6 +69,11 @@ def __init__(
         self.num_batches_processed = 0
         self.num_minibatches_processed = 0
         self.loss: Tensor | None = None
+        self.meter_start_time: float = 0
+        self.batch_time_meter = AverageMeter()
+        self.forward_time_meter = AverageMeter()
+        self.backprop_time_meter = AverageMeter()
+        self.data_time_meter = AverageMeter()
 
     @cached_property
     def unit_to_steps(self) -> dict[TimeUnit, int]:
@@ -172,6 +197,7 @@ def on_train_end(self, trainer: "Trainer[BaseConfig, Any]") -> None:
 
     def on_epoch_begin(self, trainer: "Trainer[BaseConfig, Any]") -> None:
         self.log(f"Epoch {trainer.clock.epoch} started.")
+        self.meter_start_time = time.time()
 
     def on_epoch_end(self, trainer: "Trainer[BaseConfig, Any]") -> None:
         self.log(f"Epoch {trainer.clock.epoch} ended.")
@@ -182,17 +208,41 @@ def on_batch_begin(self, trainer: "Trainer[BaseConfig, Any]") -> None:
         if self.num_minibatches_processed == 0:
             self.log(f"Iteration {trainer.clock.iteration} started.")
         self.log(f"Step {trainer.clock.step} started.")
+        self.data_time_meter.update(time.time() - self.meter_start_time)
+
+    def on_compute_loss_begin(self, trainer: Trainer[BaseConfig, Any]) -> None:
+        self.meter_start_time = time.time()
+
+    def on_compute_loss_end(self, trainer: Trainer[BaseConfig, Any]) -> None:
+        self.forward_time_meter.update(time.time() - self.meter_start_time)
+
+    def on_backward_begin(self, trainer: Trainer[BaseConfig, Any]) -> None:
+        self.meter_start_time = time.time()
 
     def on_backward_end(self, trainer: "Trainer[BaseConfig, Any]") -> None:
         self.log(f"Step {trainer.clock.step} ended.")
         trainer.clock.step += 1
         trainer.clock.num_batches_processed += 1
         trainer.clock.num_minibatches_processed += 1
+        if (not trainer.clock.is_optimizer_step) and (not trainer.clock.is_lr_scheduler_step):
+            self.backprop_time_meter.update(time.time() - self.meter_start_time)
 
     def on_optimizer_step_end(self, trainer: "Trainer[BaseConfig, Any]") -> None:
         self.log(f"Iteration {trainer.clock.iteration} ended.")
         trainer.clock.iteration += 1
         trainer.clock.num_minibatches_processed = 0
+        if not trainer.clock.is_lr_scheduler_step:
+            self.backprop_time_meter.update(time.time() - self.meter_start_time)
+
+    def on_lr_scheduler_step_end(self, trainer: Trainer[BaseConfig, Any]) -> None:
+        self.backprop_time_meter.update(time.time() - self.meter_start_time)
+
+    def on_batch_end(self, trainer: Trainer[BaseConfig, Any]) -> None:
+        data_time = self.data_time_meter.val
+        forward_time = self.forward_time_meter.val
+        backprop_time = self.backprop_time_meter.val
+        self.batch_time_meter.update(data_time + forward_time + backprop_time)
+        self.meter_start_time = time.time()
 
     def on_evaluate_begin(self, trainer: "Trainer[BaseConfig, Any]") -> None:
         self.log("Evaluation started.")

src/refiners/training_utils/wandb.py

@@ -105,6 +105,24 @@ def on_train_begin(self, trainer: "TrainerWithWandb") -> None:
         self.epoch_losses = []
         self.iteration_losses = []
 
+    def on_batch_end(self, trainer: "TrainerWithWandb") -> None:
+        batch_time, forward_time, backprop_time, data_time = (
+            trainer.clock.batch_time_meter.avg,
+            trainer.clock.forward_time_meter.avg,
+            trainer.clock.backprop_time_meter.avg,
+            trainer.clock.data_time_meter.avg,
+        )
+        if trainer.clock.is_evaluation_step:
+            effective_batch_size = trainer.clock.batch_size * trainer.clock.num_step_per_iteration
+            trainer.wandb_log(
+                data={
+                    "batch_time": batch_time / effective_batch_size,
+                    "forward_time": forward_time / effective_batch_size,
+                    "backprop_time": backprop_time / effective_batch_size,
+                    "data_time": data_time / effective_batch_size,
+                }
+            )
+
     def on_compute_loss_end(self, trainer: "TrainerWithWandb") -> None:
         loss_value = trainer.loss.detach().cpu().item()
         self.epoch_losses.append(loss_value)