Report the batch size used for the qps/throughput calculation (#2788)

Summary:

If using batch size stages, also report the current batch size that was used for the throughput calculation

Reviewed By: kausv

Differential Revision: D70892994

torchrec/metrics/metrics_namespace.py

@@ -46,6 +46,7 @@ class MetricName(MetricNameBase):
     THROUGHPUT = "throughput"
     TOTAL_EXAMPLES = "total_examples"
     ATTEMPT_EXAMPLES = "attempt_examples"
+    BATCH_SIZE = "batch_size"
     CTR = "ctr"
     CALIBRATION = "calibration"
     MSE = "mse"

torchrec/metrics/tests/test_throughput.py

@@ -199,6 +199,7 @@ def test_batch_size_schedule(self, time_mock: Mock) -> None:
             {
                 "throughput-throughput|total_examples": total_examples,
                 "throughput-throughput|attempt_examples": total_examples,
+                "throughput-throughput|batch_size": 256,
             },
         )
 
@@ -209,5 +210,6 @@ def test_batch_size_schedule(self, time_mock: Mock) -> None:
             {
                 "throughput-throughput|total_examples": total_examples,
                 "throughput-throughput|attempt_examples": total_examples,
+                "throughput-throughput|batch_size": 512,
             },
         )

torchrec/metrics/throughput.py

@@ -74,6 +74,7 @@ class ThroughputMetric(nn.Module):
     _attempt_throughput_key: str
     _total_examples_key: str
     _attempt_examples_key: str
+    _batch_size_key: str
     _steps: int
 
     def __init__(
@@ -167,6 +168,11 @@ def __init__(
             str(self._namespace),
             MetricName.ATTEMPT_EXAMPLES,
         )
+        self._batch_size_key = compose_metric_key(
+            self._namespace,
+            str(self._namespace),
+            MetricName.BATCH_SIZE,
+        )
         self._steps = 0
 
     def _get_batch_size(self) -> int:
@@ -258,4 +264,15 @@ def compute(self) -> Dict[str, torch.Tensor]:
                         self._attempt_throughput_key: attempt_throughput.clone().detach(),
                     }
                 )
+        # If using batch_size_stages, also report the current batch size
+        # that was used for the throughput calculation
+        if self._batch_size_stages is not None:
+            ret.update(
+                {
+                    self._batch_size_key: torch.tensor(
+                        self._get_batch_size(), dtype=torch.int32
+                    ),
+                }
+            )
+
         return ret