Better Tests for Better Engineering: metrics/throughput.py

Summary:
Wrote better tests for the checkpoint restoration between:
BSS job -> None-BSS job
None-BSS job ->BSS job
BSS job -> BSS job

To better simulate the exhibited code behavior:
- Create prev. modules, checkpoint it, create new module, restore from prev. checkpoint

Also caught a bug with the new tests:
Need to register the hooks regardless of the BSS being used for the cases when we run an online training job without BSS using an offline checkpoint

Reviewed By: burak-turk

Differential Revision: D72567969

fbshipit-source-id: 78b394596fb38a28ba8617ca1830503bab70199c

Author: ilyas409

torchrec/metrics/tests/test_metric_module.py

@@ -586,3 +586,20 @@ def test_save_and_load_state_dict(self) -> None:
         self.assertEqual(throughput_metric._num_batch, 100)
         # Make sure num_batch is correctly synchronized
         self.assertEqual(throughput_metric._num_batch, 100)
+
+        # Load the same checkpoint into a module that doesn't use BSS
+
+        no_bss_metric_module = generate_metric_module(
+            TestMetricModule,
+            metrics_config=DefaultMetricsConfig,
+            batch_size=128,
+            world_size=1,
+            my_rank=0,
+            state_metrics_mapping={},
+            device=torch.device("cpu"),
+            batch_size_stages=None,
+        )
+
+        no_bss_metric_module.load_state_dict(state_dict)
+        # Make sure num_batch wasn't created on the throughput module (and no exception was thrown above)
+        self.assertFalse(hasattr(no_bss_metric_module.throughput_metric, "_num_batch"))

torchrec/metrics/tests/test_throughput.py

@@ -310,6 +310,7 @@ def test_state_dict_hook_no_batch_size_stages(self) -> None:
         Verifies that the state_dict_hook does not add the 'num_batch' key when
         batch_size_stages is None.
         """
+        # Hook-only test
         throughput_metric = ThroughputMetric(
             batch_size=32,
             world_size=4,
@@ -321,41 +322,84 @@ def test_state_dict_hook_no_batch_size_stages(self) -> None:
         ThroughputMetric.state_dict_hook(throughput_metric, state_dict, prefix, {})
         self.assertNotIn(f"{prefix}num_batch", state_dict)
 
-    def test_load_state_dict_hook_restores_value(self) -> None:
+        # Lifecycle test
+
+        num_updates = 10
+        prev_job_throughput_metric = ThroughputMetric(
+            batch_size=32,
+            world_size=4,
+            window_seconds=100,
+            batch_size_stages=None,
+        )
+        for _ in range(num_updates):
+            prev_job_throughput_metric.update()
+        prev_state_dict = prev_job_throughput_metric.state_dict()
+
+        curr_job_throughput_metric = ThroughputMetric(
+            batch_size=32,
+            world_size=4,
+            window_seconds=100,
+            batch_size_stages=None,
+        )
+
+        curr_job_throughput_metric.load_state_dict(prev_state_dict)
+        # Make sure _num_batch is not present as an argument of the class
+        self.assertFalse(hasattr(curr_job_throughput_metric, "_num_batch"))
+
+    def test_load_state_dict_hook_resumes_from_checkpoint_with_bss_from_bss(
+        self,
+    ) -> None:
         """
         Checks that the load_state_dict_hook correctly restores the 'num_batch' value
         from the state_dict.
         """
-        throughput_metric = ThroughputMetric(
+        num_updates = 10
+        prev_job_throughput_metric = ThroughputMetric(
             batch_size=32,
             world_size=4,
             window_seconds=100,
             batch_size_stages=[BatchSizeStage(256, 1), BatchSizeStage(512, None)],
         )
-        state_dict: OrderedDict[str, torch.Tensor] = OrderedDict()
-        prefix: str = "test_prefix_"
-        state_dict[f"{prefix}num_batch"] = torch.tensor(10, dtype=torch.long)
-        throughput_metric.load_state_dict_hook(state_dict, prefix, {}, True, [], [], [])
-        self.assertEqual(throughput_metric._num_batch, 10)
+        for _ in range(num_updates):
+            prev_job_throughput_metric.update()
+        prev_state_dict = prev_job_throughput_metric.state_dict()
+
+        curr_job_throughput_metric = ThroughputMetric(
+            batch_size=32,
+            world_size=4,
+            window_seconds=100,
+            batch_size_stages=[BatchSizeStage(1024, 1), BatchSizeStage(2048, None)],
+        )
+
+        curr_job_throughput_metric.load_state_dict(prev_state_dict)
+        self.assertEqual(curr_job_throughput_metric._num_batch, num_updates)
 
     def test_load_state_dict_hook_resumes_from_checkpoint_without_bss(self) -> None:
         """
         Verifies that the load_state_dict_hook correctly handles the case where a
         previously checkpointed job used the batch_size_stages, but a subsequent job,
         restored from a checkpoint, isn't using them.
         """
-        throughput_metric = ThroughputMetric(
+
+        prev_job_throughput_metric = ThroughputMetric(
+            batch_size=32,
+            world_size=4,
+            window_seconds=100,
+            batch_size_stages=[BatchSizeStage(256, 1), BatchSizeStage(512, None)],
+        )
+
+        prev_state_dict = prev_job_throughput_metric.state_dict()
+
+        curr_job_throughput_metric = ThroughputMetric(
             batch_size=32,
             world_size=4,
             window_seconds=100,
             batch_size_stages=None,  # No batch_size_stages
         )
-        state_dict: OrderedDict[str, torch.Tensor] = OrderedDict()
-        prefix: str = "test_prefix_"
-        state_dict[f"{prefix}num_batch"] = torch.tensor(10, dtype=torch.long)
-        throughput_metric.load_state_dict_hook(state_dict, prefix, {}, True, [], [], [])
 
-        self.assertFalse(hasattr(throughput_metric, "_num_batch"))
+        curr_job_throughput_metric.load_state_dict(prev_state_dict)
+
+        self.assertFalse(hasattr(curr_job_throughput_metric, "_num_batch"))
 
     def test_load_state_dict_hook_resumes_from_checkpoint_with_bss_without_key(
         self,
@@ -365,15 +409,22 @@ def test_load_state_dict_hook_resumes_from_checkpoint_with_bss_without_key(
         previously checkpointed job didn't use batch_size_stages, but a subsequent job,
         restored from a checkpoint, is using them.
         """
-        throughput_metric = ThroughputMetric(
+        prev_job_throughput_metric = ThroughputMetric(
+            batch_size=32,
+            world_size=4,
+            window_seconds=100,
+            batch_size_stages=None,  # No batch_size_stages
+        )
+        prev_state_dict = prev_job_throughput_metric.state_dict()
+
+        curr_job_throughput_metric = ThroughputMetric(
             batch_size=32,
             world_size=4,
             window_seconds=100,
             batch_size_stages=[BatchSizeStage(256, 1), BatchSizeStage(512, None)],
         )
-        # Empty state_dict
-        state_dict: OrderedDict[str, torch.Tensor] = OrderedDict()
-        prefix: str = "test_prefix_"
-        throughput_metric.load_state_dict_hook(state_dict, prefix, {}, True, [], [], [])
+
+        curr_job_throughput_metric.load_state_dict(prev_state_dict)
+
         # Expecting 0
-        self.assertEqual(throughput_metric._num_batch, 0)
+        self.assertEqual(curr_job_throughput_metric._num_batch, 0)

torchrec/metrics/throughput.py

@@ -115,8 +115,9 @@ def __init__(
         if self._batch_size_stages is not None:
             # Keep track of the number of batches if using batch_size_stages
             self._num_batch: int = 0
-            self._register_load_state_dict_pre_hook(self.load_state_dict_hook)
-            self.register_state_dict_post_hook(self.state_dict_hook)
+
+        self._register_load_state_dict_pre_hook(self.load_state_dict_hook)
+        self.register_state_dict_post_hook(self.state_dict_hook)
 
         self.register_buffer("total_examples", torch.tensor(0, dtype=torch.long))
         self.register_buffer("warmup_examples", torch.tensor(0, dtype=torch.long))