fix RecMetrics loading (make trained_batches a buffer) (#3534)

Summary:

This diff addresses the following task: T209753398

Currently `trained_batches` is not stored in state_dict, requiring us to manually sync this variable upon checkpoint loading. We make this variable a buffer so that it can now be captured with model state dict.

Differential Revision: D86697665

Author: prajjwal1

torchrec/metrics/metric_module.py

@@ -202,7 +202,10 @@ def __init__(
         self.rec_metrics = rec_metrics if rec_metrics else RecMetricList([])
         self.throughput_metric = throughput_metric
         self.state_metrics = state_metrics if state_metrics else {}
-        self.trained_batches: int = 0
+
+        trained_batches = torch.tensor(0, dtype=torch.int64)
+        self.register_buffer("_trained_batches", trained_batches, persistent=True)
+
         self.batch_size = batch_size
         self.world_size = world_size
         self.oom_count = 0
@@ -228,6 +231,15 @@ def __init__(
         )
         self.last_compute_time = -1.0
 
+    @property
+    def trained_batches(self) -> int:
+        # .trained_batches should return an int
+        return int(self._trained_batches.item())
+
+    @trained_batches.setter
+    def trained_batches(self, value: int) -> None:
+        self._trained_batches.fill_(int(value))
+
     def _update_rec_metrics(
         self, model_out: Dict[str, torch.Tensor], **kwargs: Any
     ) -> None:
@@ -260,7 +272,7 @@ def update(self, model_out: Dict[str, torch.Tensor], **kwargs: Any) -> None:
             self._update_rec_metrics(model_out, **kwargs)
             if self.throughput_metric:
                 self.throughput_metric.update()
-            self.trained_batches += 1
+            self._trained_batches.add_(1)
 
     def _adjust_compute_interval(self) -> None:
         """

torchrec/metrics/tests/test_cpu_offloaded_metric_module.py

@@ -341,6 +341,7 @@ def test_state_dict_save_load(self) -> None:
                 "rec_metrics.rec_metrics.0._metrics_computations.0.state_3": torch.tensor(
                     [6.0]
                 ),
+                "_trained_batches": torch.tensor([0], dtype=torch.int64),
             },
         )