Fuse "states" level tensors to reduce all gather during metrics compute (#2892)

Summary:
Pull Request resolved: #2892

# Contexts
During S503023, we found that the metrics `compute` time can take up to **30s** on a large OMNI FM model. (Note this is not the metrics `update` every iteration. This is `compute` which will create sync and refresh metrics on tensorboard. The frequency of `compute` can be set by `compute_interval_steps` (default=100).) If Zoomer happens to capture the iteration of metrics compute, it will even make the trace file too big and can't be opened.

The reason for metrics `compute` being so long was due to too many all_gather calls. See the screenshot of metrics `compute` from the above SEV. That single metrics `compute` takes 30s. And once zooming in, you will find it composed of hundreds of all_gather.
 {F1976937278}
 {F1976937301}

Therefore, this diff tried to fuse the tensors before all_gather during metrics `compute`. With this diff plus `FUSE_TASKS_COMPUTATION` to fuse task-level tensors, it can reduce "RecMetricModule compute" time from 1s 287 ms to 35 ms (36X reduction) on a shrunk OMNI FM model.

# What has been changed/added?
1. Add `FUSED_TASKS_AND_STATES_COMPUTATION` in RecComputeMode.
When turned on, this will both fuse tasks tensors and fuse state tensors.
2. Add `fuse_state_tensors` to [Metric](https://www.internalfb.com/code/fbsource/[0c89c01039abfadd62e8ec1b34eb24b249b99f3f]/fbcode/pytorch_lightning_deprecated/metrics/torchmetrics/metric.py?lines=43) class.
This will be turned on once `FUSED_TASKS_AND_STATES_COMPUTATION` is set from config. Then when a metric's `compute` is called, it will fuse (stack) state tensors before all_gather. Then reconstruct the output tensor to desired format to do reduction.
3. It is noted currently we only support fusing/stacking 1D state tensors. Therefore, for states with `List` (e.g. auc) or 2D tensor (e.g. multiclass_ne), the `FUSED_TASKS_AND_STATES_COMPUTATION` shouldn't be used or should at least fall back to either `FUSED_TASKS_COMPUTATION` or `UNFUSED_TASKS_COMPUTATION`.

Reviewed By: iamzainhuda

Differential Revision: D72010614

fbshipit-source-id: 5ac77088cf9737ad783ad7e1740daf6afc0224a8

Author: ge0405

torchrec/metrics/metrics_config.py

@@ -87,6 +87,7 @@ class RecComputeMode(Enum):
 
     FUSED_TASKS_COMPUTATION = 1
     UNFUSED_TASKS_COMPUTATION = 2
+    FUSED_TASKS_AND_STATES_COMPUTATION = 3
 
 
 _DEFAULT_WINDOW_SIZE = 10_000_000

torchrec/metrics/precision_session.py

@@ -196,7 +196,10 @@ def __init__(
         process_group: Optional[dist.ProcessGroup] = None,
         **kwargs: Any,
     ) -> None:
-        if compute_mode == RecComputeMode.FUSED_TASKS_COMPUTATION:
+        if compute_mode in [
+            RecComputeMode.FUSED_TASKS_COMPUTATION,
+            RecComputeMode.FUSED_TASKS_AND_STATES_COMPUTATION,
+        ]:
             raise RecMetricException(
                 "Fused computation is not supported for precision session-level metrics"
             )

torchrec/metrics/rec_metric.py

@@ -134,13 +134,19 @@ def __init__(
         window_size: int,
         compute_on_all_ranks: bool = False,
         should_validate_update: bool = False,
+        fuse_state_tensors: bool = False,
         process_group: Optional[dist.ProcessGroup] = None,
         fused_update_limit: int = 0,
         allow_missing_label_with_zero_weight: bool = False,
         *args: Any,
         **kwargs: Any,
     ) -> None:
-        super().__init__(process_group=process_group, *args, **kwargs)
+        super().__init__(
+            process_group=process_group,
+            fuse_state_tensors=fuse_state_tensors,
+            *args,
+            **kwargs,
+        )
 
         self._my_rank = my_rank
         self._n_tasks = n_tasks
@@ -341,7 +347,11 @@ def __init__(
         # TODO(stellaya): consider to inherit from TorchMetrics.Metric or
         # TorchMetrics.MetricCollection.
         if (
-            compute_mode == RecComputeMode.FUSED_TASKS_COMPUTATION
+            compute_mode
+            in [
+                RecComputeMode.FUSED_TASKS_COMPUTATION,
+                RecComputeMode.FUSED_TASKS_AND_STATES_COMPUTATION,
+            ]
             and fused_update_limit > 0
         ):
             raise ValueError(
@@ -376,7 +386,10 @@ def __init__(
                 f"Local window size must be larger than batch size. Got local window size {self._window_size} and batch size {self._batch_size}."
             )
 
-        if compute_mode == RecComputeMode.FUSED_TASKS_COMPUTATION:
+        if compute_mode in [
+            RecComputeMode.FUSED_TASKS_COMPUTATION,
+            RecComputeMode.FUSED_TASKS_AND_STATES_COMPUTATION,
+        ]:
             task_per_metric = len(self._tasks)
             self._tasks_iter = self._fused_tasks_iter
         else:
@@ -385,7 +398,11 @@ def __init__(
 
         for task_config in (
             [self._tasks]
-            if compute_mode == RecComputeMode.FUSED_TASKS_COMPUTATION
+            if compute_mode
+            in [
+                RecComputeMode.FUSED_TASKS_COMPUTATION,
+                RecComputeMode.FUSED_TASKS_AND_STATES_COMPUTATION,
+            ]
             else self._tasks
         ):
             # pyre-ignore
@@ -394,13 +411,16 @@ def __init__(
             # according to https://github.com/python/mypy/issues/3048.
             # pyre-fixme[45]: Cannot instantiate abstract class `RecMetricCoputation`.
             metric_computation = self._computation_class(
-                my_rank,
-                batch_size,
-                task_per_metric,
-                self._window_size,
-                compute_on_all_ranks,
-                self._should_validate_update,
-                process_group,
+                my_rank=my_rank,
+                batch_size=batch_size,
+                n_tasks=task_per_metric,
+                window_size=self._window_size,
+                compute_on_all_ranks=compute_on_all_ranks,
+                should_validate_update=self._should_validate_update,
+                fuse_state_tensors=(
+                    compute_mode == RecComputeMode.FUSED_TASKS_AND_STATES_COMPUTATION
+                ),
+                process_group=process_group,
                 **{**kwargs, **self._get_task_kwargs(task_config)},
             )
             required_inputs = self._get_task_required_inputs(task_config)
@@ -527,7 +547,10 @@ def _update(
         **kwargs: Dict[str, Any],
     ) -> None:
         with torch.no_grad():
-            if self._compute_mode == RecComputeMode.FUSED_TASKS_COMPUTATION:
+            if self._compute_mode in [
+                RecComputeMode.FUSED_TASKS_COMPUTATION,
+                RecComputeMode.FUSED_TASKS_AND_STATES_COMPUTATION,
+            ]:
                 task_names = [task.name for task in self._tasks]
 
                 if not isinstance(predictions, torch.Tensor):

torchrec/metrics/recall_session.py

@@ -235,7 +235,10 @@ def __init__(
         process_group: Optional[dist.ProcessGroup] = None,
         **kwargs: Any,
     ) -> None:
-        if compute_mode == RecComputeMode.FUSED_TASKS_COMPUTATION:
+        if compute_mode in [
+            RecComputeMode.FUSED_TASKS_COMPUTATION,
+            RecComputeMode.FUSED_TASKS_AND_STATES_COMPUTATION,
+        ]:
             raise RecMetricException(
                 "Fused computation is not supported for recall session-level metrics"
             )

torchrec/metrics/test_utils/__init__.py

@@ -291,7 +291,10 @@ def get_target_rec_metric_value(
             labels, predictions, weights, _ = parse_task_model_outputs(
                 tasks, model_outs[i]
             )
-            if target_compute_mode == RecComputeMode.FUSED_TASKS_COMPUTATION:
+            if target_compute_mode in [
+                RecComputeMode.FUSED_TASKS_COMPUTATION,
+                RecComputeMode.FUSED_TASKS_AND_STATES_COMPUTATION,
+            ]:
                 labels = torch.stack(list(labels.values()))
                 predictions = torch.stack(list(predictions.values()))
                 weights = torch.stack(list(weights.values()))

torchrec/metrics/tests/test_accuracy.py

@@ -52,7 +52,7 @@ class AccuracyMetricTest(unittest.TestCase):
     clazz: Type[RecMetric] = AccuracyMetric
     task_name: str = "accuracy"
 
-    def test_unfused_accuracy(self) -> None:
+    def test_accuracy_unfused(self) -> None:
         rec_metric_value_test_launcher(
             target_clazz=AccuracyMetric,
             target_compute_mode=RecComputeMode.UNFUSED_TASKS_COMPUTATION,
@@ -66,7 +66,7 @@ def test_unfused_accuracy(self) -> None:
             entry_point=metric_test_helper,
         )
 
-    def test_fused_accuracy(self) -> None:
+    def test_accuracy_fused_tasks(self) -> None:
         rec_metric_value_test_launcher(
             target_clazz=AccuracyMetric,
             target_compute_mode=RecComputeMode.FUSED_TASKS_COMPUTATION,
@@ -80,6 +80,20 @@ def test_fused_accuracy(self) -> None:
             entry_point=metric_test_helper,
         )
 
+    def test_accuracy_fused_tasks_and_states(self) -> None:
+        rec_metric_value_test_launcher(
+            target_clazz=AccuracyMetric,
+            target_compute_mode=RecComputeMode.FUSED_TASKS_AND_STATES_COMPUTATION,
+            test_clazz=TestAccuracyMetric,
+            metric_name=AccuracyMetricTest.task_name,
+            task_names=["t1", "t2", "t3"],
+            fused_update_limit=0,
+            compute_on_all_ranks=False,
+            should_validate_update=False,
+            world_size=WORLD_SIZE,
+            entry_point=metric_test_helper,
+        )
+
 
 class AccuracyMetricValueTest(unittest.TestCase):
     r"""This set of tests verify the computation logic of accuracy in several

torchrec/metrics/tests/test_cali_free_ne.py

@@ -89,7 +89,7 @@ def test_cali_free_ne_unfused(self) -> None:
             entry_point=metric_test_helper,
         )
 
-    def test_cali_free_ne_fused(self) -> None:
+    def test_cali_free_ne_fused_tasks(self) -> None:
         rec_metric_value_test_launcher(
             target_clazz=CaliFreeNEMetric,
             target_compute_mode=RecComputeMode.FUSED_TASKS_COMPUTATION,
@@ -103,7 +103,21 @@ def test_cali_free_ne_fused(self) -> None:
             entry_point=metric_test_helper,
         )
 
-    def test_cali_free_ne_update_fused(self) -> None:
+    def test_cali_free_ne_fused_tasks_and_states(self) -> None:
+        rec_metric_value_test_launcher(
+            target_clazz=CaliFreeNEMetric,
+            target_compute_mode=RecComputeMode.FUSED_TASKS_AND_STATES_COMPUTATION,
+            test_clazz=TestCaliFreeNEMetric,
+            metric_name=CaliFreeNEMetricTest.task_name,
+            task_names=["t1", "t2", "t3"],
+            fused_update_limit=0,
+            compute_on_all_ranks=False,
+            should_validate_update=False,
+            world_size=WORLD_SIZE,
+            entry_point=metric_test_helper,
+        )
+
+    def test_cali_free_ne_update_unfused(self) -> None:
         rec_metric_value_test_launcher(
             target_clazz=CaliFreeNEMetric,
             target_compute_mode=RecComputeMode.UNFUSED_TASKS_COMPUTATION,

torchrec/metrics/tests/test_calibration.py

@@ -52,7 +52,7 @@ class CalibrationMetricTest(unittest.TestCase):
     clazz: Type[RecMetric] = CalibrationMetric
     task_name: str = "calibration"
 
-    def test_unfused_calibration(self) -> None:
+    def test_calibration_unfused(self) -> None:
         rec_metric_value_test_launcher(
             target_clazz=CalibrationMetric,
             target_compute_mode=RecComputeMode.UNFUSED_TASKS_COMPUTATION,
@@ -66,7 +66,7 @@ def test_unfused_calibration(self) -> None:
             entry_point=metric_test_helper,
         )
 
-    def test_fused_calibration(self) -> None:
+    def test_calibration_fused_tasks(self) -> None:
         rec_metric_value_test_launcher(
             target_clazz=CalibrationMetric,
             target_compute_mode=RecComputeMode.FUSED_TASKS_COMPUTATION,
@@ -80,6 +80,20 @@ def test_fused_calibration(self) -> None:
             entry_point=metric_test_helper,
         )
 
+    def test_calibration_fused_tasks_and_states(self) -> None:
+        rec_metric_value_test_launcher(
+            target_clazz=CalibrationMetric,
+            target_compute_mode=RecComputeMode.FUSED_TASKS_AND_STATES_COMPUTATION,
+            test_clazz=TestCalibrationMetric,
+            metric_name=CalibrationMetricTest.task_name,
+            task_names=["t1", "t2", "t3"],
+            fused_update_limit=0,
+            compute_on_all_ranks=False,
+            should_validate_update=False,
+            world_size=WORLD_SIZE,
+            entry_point=metric_test_helper,
+        )
+
 
 class CalibrationGPUSyncTest(unittest.TestCase):
     clazz: Type[RecMetric] = CalibrationMetric

torchrec/metrics/tests/test_ctr.py

@@ -46,7 +46,7 @@ class CTRMetricTest(unittest.TestCase):
     clazz: Type[RecMetric] = CTRMetric
     task_name: str = "ctr"
 
-    def test_unfused_ctr(self) -> None:
+    def test_ctr_unfused(self) -> None:
         rec_metric_value_test_launcher(
             target_clazz=CTRMetric,
             target_compute_mode=RecComputeMode.UNFUSED_TASKS_COMPUTATION,
@@ -60,7 +60,7 @@ def test_unfused_ctr(self) -> None:
             entry_point=metric_test_helper,
         )
 
-    def test_fused_ctr(self) -> None:
+    def test_ctr_fused_tasks(self) -> None:
         rec_metric_value_test_launcher(
             target_clazz=CTRMetric,
             target_compute_mode=RecComputeMode.FUSED_TASKS_COMPUTATION,
@@ -74,6 +74,20 @@ def test_fused_ctr(self) -> None:
             entry_point=metric_test_helper,
         )
 
+    def test_ctr_fused_tasks_and_states(self) -> None:
+        rec_metric_value_test_launcher(
+            target_clazz=CTRMetric,
+            target_compute_mode=RecComputeMode.FUSED_TASKS_AND_STATES_COMPUTATION,
+            test_clazz=TestCTRMetric,
+            metric_name=CTRMetricTest.task_name,
+            task_names=["t1", "t2", "t3"],
+            fused_update_limit=0,
+            compute_on_all_ranks=False,
+            should_validate_update=False,
+            world_size=WORLD_SIZE,
+            entry_point=metric_test_helper,
+        )
+
 
 class CTRGPUSyncTest(unittest.TestCase):
     clazz: Type[RecMetric] = CTRMetric

torchrec/metrics/tests/test_hindsight_target_pr.py

@@ -126,7 +126,7 @@ class TestHindsightTargetPRMetricTest(unittest.TestCase):
     precision_task_name: str = "hindsight_target_precision"
     recall_task_name: str = "hindsight_target_recall"
 
-    def test_unfused_hindsight_target_precision(self) -> None:
+    def test_hindsight_target_precision_unfused(self) -> None:
         rec_metric_value_test_launcher(
             target_clazz=HindsightTargetPRMetric,
             target_compute_mode=RecComputeMode.UNFUSED_TASKS_COMPUTATION,
@@ -140,7 +140,7 @@ def test_unfused_hindsight_target_precision(self) -> None:
             entry_point=metric_test_helper,
         )
 
-    def test_unfused_hindsight_target_recall(self) -> None:
+    def test_hindsight_target_recall_unfused(self) -> None:
         rec_metric_value_test_launcher(
             target_clazz=HindsightTargetPRMetric,
             target_compute_mode=RecComputeMode.UNFUSED_TASKS_COMPUTATION,

torchrec/metrics/tests/test_mse.py

@@ -70,7 +70,7 @@ class MSEMetricTest(unittest.TestCase):
     task_name: str = "mse"
     rmse_task_name: str = "rmse"
 
-    def test_unfused_mse(self) -> None:
+    def test_mse_unfused(self) -> None:
         rec_metric_value_test_launcher(
             target_clazz=MSEMetric,
             target_compute_mode=RecComputeMode.UNFUSED_TASKS_COMPUTATION,
@@ -84,7 +84,7 @@ def test_unfused_mse(self) -> None:
             entry_point=metric_test_helper,
         )
 
-    def test_fused_mse(self) -> None:
+    def test_mse_fused_tasks(self) -> None:
         rec_metric_value_test_launcher(
             target_clazz=MSEMetric,
             target_compute_mode=RecComputeMode.FUSED_TASKS_COMPUTATION,
@@ -98,7 +98,21 @@ def test_fused_mse(self) -> None:
             entry_point=metric_test_helper,
         )
 
-    def test_unfused_rmse(self) -> None:
+    def test_mse_fused_tasks_and_states(self) -> None:
+        rec_metric_value_test_launcher(
+            target_clazz=MSEMetric,
+            target_compute_mode=RecComputeMode.FUSED_TASKS_AND_STATES_COMPUTATION,
+            test_clazz=TestMSEMetric,
+            metric_name=MSEMetricTest.task_name,
+            task_names=["t1", "t2", "t3"],
+            fused_update_limit=0,
+            compute_on_all_ranks=False,
+            should_validate_update=False,
+            world_size=WORLD_SIZE,
+            entry_point=metric_test_helper,
+        )
+
+    def test_rmse_unfused(self) -> None:
         rec_metric_value_test_launcher(
             target_clazz=MSEMetric,
             target_compute_mode=RecComputeMode.UNFUSED_TASKS_COMPUTATION,
@@ -112,7 +126,7 @@ def test_unfused_rmse(self) -> None:
             entry_point=metric_test_helper,
         )
 
-    def test_fused_rmse(self) -> None:
+    def test_rmse_fused_tasks(self) -> None:
         rec_metric_value_test_launcher(
             target_clazz=MSEMetric,
             target_compute_mode=RecComputeMode.FUSED_TASKS_COMPUTATION,
@@ -126,6 +140,20 @@ def test_fused_rmse(self) -> None:
             entry_point=metric_test_helper,
         )
 
+    def test_rmse_fused_tasks_and_states(self) -> None:
+        rec_metric_value_test_launcher(
+            target_clazz=MSEMetric,
+            target_compute_mode=RecComputeMode.FUSED_TASKS_AND_STATES_COMPUTATION,
+            test_clazz=TestRMSEMetric,
+            metric_name=MSEMetricTest.rmse_task_name,
+            task_names=["t1", "t2", "t3"],
+            fused_update_limit=0,
+            compute_on_all_ranks=False,
+            should_validate_update=False,
+            world_size=WORLD_SIZE,
+            entry_point=metric_test_helper,
+        )
+
 
 class MSEGPUSyncTest(unittest.TestCase):
     clazz: Type[RecMetric] = MSEMetric