Merge pull request #151 from AI-Hypercomputer/yujiedeng/throughput

add throughput related metrics

Author: vlad-karp

recml/core/training/jax_trainer.py

@@ -20,6 +20,7 @@
 import os
 import pprint
 from typing import Any, Generic, Protocol, Self, TypeVar
+import time
 
 from absl import logging
 from clu import data as clu_data
@@ -557,28 +558,52 @@ def _write_marker_file(self):
     ) as f:
       f.write("COMPLETED")
 
-  def _train_n_steps(
-      self,
-      train_iter: Iterator[PyTree],
-      train_step: partitioning.StepFn,
-      state: State,
-      start_step: int,
-      num_steps: int,
-      summary_writer: metrics_tools.AsyncMultiWriter,
-  ) -> tuple[State, Mapping[str, Any]]:
-    """Performs a training loop and returns the updated state and metrics."""
+  def _train_n_steps(self, train_iter, train_step, state, start_step, num_steps, summary_writer):
     metrics_accum = metrics_tools.MetricAccumulator(summary_writer)
+  
+    warmup_steps = 3
+    total_examples_in_loop = 0
+    valid_steps_in_loop = 0
+    loop_start_time = time.time()
+
     for step in range(start_step, start_step + num_steps):
       with jax.profiler.StepTraceAnnotation("train", step_num=step):
+        if step == warmup_steps:
+          loop_start_time = time.time()
         train_batch = next(train_iter)
         inputs = self._partitioner.shard_inputs(train_batch)
+      
         state, metrics_update = train_step(inputs, state)
+        if step >= warmup_steps:
+          if 'common/batch_size' in metrics_update:
+             total_examples_in_loop += metrics_update['common/batch_size'].compute()
+             valid_steps_in_loop += 1
+
         metrics_accum.accumulate(metrics_update, step)
         self.report_progress(step)
+        
         if (step != start_step + num_steps - 1) and self._enable_checkpointing:
           self._maybe_save_checkpoint(step, state)
 
+    duration = time.time() - loop_start_time
+    
     metrics = metrics_accum.compute_and_log_scalars(start_step + num_steps - 1)
+    
+    # Calculate and inject overall loop performance
+    if valid_steps_in_loop > 0 and duration > 0:
+        throughput = total_examples_in_loop / duration
+        ms_per_step = (duration / valid_steps_in_loop) * 1000
+        
+        metrics.update({
+            'perf/loop_throughput_ex_per_sec': throughput,
+            'perf/loop_ms_per_step': ms_per_step,
+        })
+        
+        summary_writer.write_scalars(start_step + num_steps - 1, {
+            'perf/loop_throughput_ex_per_sec': throughput,
+            'perf/loop_ms_per_step': ms_per_step,
+        })
+
     return state, metrics
 
   def _evaluate_n_steps(

recml/examples/dlrm_experiment.py

@@ -276,6 +276,7 @@ def _loss_fn(params: jt.PyTree) -> tuple[jt.Scalar, jt.Array]:
       loss = jnp.mean(optax.sigmoid_binary_cross_entropy(logits, label), axis=0)
       return loss, logits
 
+    global_batch_size = self.train_data.global_batch_size
     grad_fn = jax.value_and_grad(_loss_fn, has_aux=True, allow_int=True)
     (loss, logits), grads = grad_fn(state.params)
     state = state.update(grads=grads)
@@ -287,6 +288,7 @@ def _loss_fn(params: jt.PyTree) -> tuple[jt.Scalar, jt.Array]:
         'aucroc': recml.metrics.aucroc(label, logits, from_logits=True),
         'label/mean': recml.metrics.mean(label),
         'prediction/mean': recml.metrics.mean(jax.nn.sigmoid(logits)),
+        "common/batch_size": recml.metrics.scalar(global_batch_size),
     }
     return state, metrics