PT2 Compile Individual Metrics (#2517)

Summary:

We find that PT2 compiling the TorchMetrics helper functions that are called during  `.update()` and `._compute()` for `XMetricComputation` classes improves both the CPU  wall time (95.4ms -> 65.4ms) and GPU wall time (3.3ms -> 1.2ms)

We use `MetricsConfig.enable_pt2_config` to set the `enable_pt2_config` attribute within a class. The `pt2_compile` decorator compiles methods within a class based on whether the `pt2_compile` method is set to True.


Without PT2 compile
 {F1941555284}

With PT2 Compile
{F1941555918}

Reviewed By: iamzainhuda

Differential Revision: D64244129

torchrec/metrics/accuracy.py

@@ -17,6 +17,7 @@
     RecMetricComputation,
     RecMetricException,
 )
+from torchrec.pt2.utils import pt2_compile_callable
 
 
 THRESHOLD = "threshold"
@@ -84,6 +85,7 @@ def __init__(self, *args: Any, threshold: float = 0.5, **kwargs: Any) -> None:
         )
         self._threshold: float = threshold
 
+    @pt2_compile_callable
     def update(
         self,
         *,

torchrec/metrics/auc.py

@@ -21,6 +21,7 @@
     RecMetricComputation,
     RecMetricException,
 )
+from torchrec.pt2.utils import pt2_compile_callable
 
 PREDICTIONS = "predictions"
 LABELS = "labels"
@@ -243,6 +244,7 @@ def _init_states(self) -> None:
         if self._grouped_auc:
             getattr(self, GROUPING_KEYS).append(torch.tensor([-1], device=self.device))
 
+    @pt2_compile_callable
     def update(
         self,
         *,

torchrec/metrics/auprc.py

@@ -21,6 +21,7 @@
     RecMetricComputation,
     RecMetricException,
 )
+from torchrec.pt2.utils import pt2_compile_callable
 
 PREDICTIONS = "predictions"
 LABELS = "labels"
@@ -235,6 +236,7 @@ def _init_states(self) -> None:
         if self._grouped_auprc:
             getattr(self, GROUPING_KEYS).append(torch.tensor([-1], device=self.device))
 
+    @pt2_compile_callable
     def update(
         self,
         *,

torchrec/metrics/calibration.py

@@ -17,6 +17,7 @@
     RecMetricComputation,
     RecMetricException,
 )
+from torchrec.pt2.utils import pt2_compile_callable
 
 CALIBRATION_NUM = "calibration_num"
 CALIBRATION_DENOM = "calibration_denom"
@@ -65,6 +66,7 @@ def __init__(self, *args: Any, **kwargs: Any) -> None:
             persistent=True,
         )
 
+    @pt2_compile_callable
     def update(
         self,
         *,

torchrec/metrics/ctr.py

@@ -10,13 +10,15 @@
 from typing import Any, cast, Dict, List, Optional, Type
 
 import torch
+
 from torchrec.metrics.metrics_namespace import MetricName, MetricNamespace, MetricPrefix
 from torchrec.metrics.rec_metric import (
     MetricComputationReport,
     RecMetric,
     RecMetricComputation,
     RecMetricException,
 )
+from torchrec.pt2.utils import pt2_compile_callable
 
 CTR_NUM = "ctr_num"
 CTR_DENOM = "ctr_denom"
@@ -61,6 +63,7 @@ def __init__(self, *args: Any, **kwargs: Any) -> None:
             persistent=True,
         )
 
+    @pt2_compile_callable
     def update(
         self,
         *,

torchrec/metrics/mae.py

@@ -17,6 +17,7 @@
     RecMetricComputation,
     RecMetricException,
 )
+from torchrec.pt2.utils import pt2_compile_callable
 
 
 ERROR_SUM = "error_sum"
@@ -72,6 +73,7 @@ def __init__(self, *args: Any, **kwargs: Any) -> None:
             persistent=True,
         )
 
+    @pt2_compile_callable
     # pyre-fixme[14]: `update` overrides method defined in `RecMetricComputation`
     #  inconsistently.
     def update(

torchrec/metrics/metric_module.py

@@ -393,6 +393,8 @@ def _generate_rec_metrics(
         if metric_def and metric_def.arguments is not None:
             kwargs = metric_def.arguments
 
+        kwargs["enable_pt2_compile"] = metrics_config.enable_pt2_compile
+
         rec_tasks: List[RecTaskInfo] = []
         if metric_def.rec_tasks and metric_def.rec_task_indices:
             raise ValueError(
@@ -468,7 +470,7 @@ def generate_metric_module(
         metrics_config, world_size, my_rank, batch_size, process_group
     )
     """
-    Batch_size_stages currently only used by ThroughputMetric to ensure total_example correct so 
+    Batch_size_stages currently only used by ThroughputMetric to ensure total_example correct so
     different training jobs have aligned mertics.
     TODO: update metrics other than ThroughputMetric if it has dependency on batch_size
     """

torchrec/metrics/metrics_config.py

@@ -158,6 +158,7 @@ class MetricsConfig:
         should_validate_update (bool): whether to check the inputs of update() and skip
             update if the inputs are invalid. Invalid inputs include the case where all
             examples have 0 weights for a batch.
+        enable_pt2_compile (bool): whether to enable PT2 compilation for metrics.
     """
 
     rec_tasks: List[RecTaskInfo] = field(default_factory=list)
@@ -171,6 +172,7 @@ class MetricsConfig:
     max_compute_interval: float = float("inf")
     compute_on_all_ranks: bool = False
     should_validate_update: bool = False
+    enable_pt2_compile: bool = False
 
 
 DefaultTaskInfo = RecTaskInfo(

torchrec/metrics/mse.py

@@ -10,13 +10,15 @@
 from typing import Any, cast, Dict, List, Optional, Type
 
 import torch
+
 from torchrec.metrics.metrics_namespace import MetricName, MetricNamespace, MetricPrefix
 from torchrec.metrics.rec_metric import (
     MetricComputationReport,
     RecMetric,
     RecMetricComputation,
     RecMetricException,
 )
+from torchrec.pt2.utils import pt2_compile_callable
 
 
 ERROR_SUM = "error_sum"
@@ -80,6 +82,7 @@ def __init__(self, *args: Any, **kwargs: Any) -> None:
             persistent=True,
         )
 
+    @pt2_compile_callable
     def update(
         self,
         *,

torchrec/metrics/multiclass_recall.py

@@ -18,6 +18,7 @@
     RecMetricComputation,
     RecMetricException,
 )
+from torchrec.pt2.utils import pt2_compile_callable
 
 
 def compute_true_positives_at_k(
@@ -109,6 +110,7 @@ def __init__(self, *args: Any, **kwargs: Any) -> None:
             persistent=True,
         )
 
+    @pt2_compile_callable
     def update(
         self,
         *,

torchrec/metrics/ndcg.py

@@ -19,6 +19,7 @@
     RecMetricComputation,
     RecMetricException,
 )
+from torchrec.pt2.utils import pt2_compile_callable
 
 SUM_NDCG = "sum_ndcg"
 NUM_SESSIONS = "num_sessions"
@@ -331,6 +332,7 @@ def __init__(
             persistent=True,
         )
 
+    @pt2_compile_callable
     def update(
         self,
         *,

torchrec/metrics/ne.py

@@ -17,6 +17,7 @@
     RecMetricComputation,
     RecMetricException,
 )
+from torchrec.pt2.utils import pt2_compile_callable
 
 
 def compute_cross_entropy(
@@ -148,6 +149,7 @@ def __init__(
         )
         self.eta = 1e-12
 
+    @pt2_compile_callable
     def update(
         self,
         *,

torchrec/metrics/ne_positive.py

@@ -17,6 +17,7 @@
     RecMetricComputation,
     RecMetricException,
 )
+from torchrec.pt2.utils import pt2_compile_callable
 
 
 def compute_cross_entropy_positive(
@@ -130,6 +131,7 @@ def __init__(
         )
         self.eta = 1e-12
 
+    @pt2_compile_callable
     def update(
         self,
         *,

torchrec/metrics/output.py

@@ -21,6 +21,7 @@
     RecMetricException,
     RecTaskInfo,
 )
+from torchrec.pt2.utils import pt2_compile_callable
 
 
 class OutputMetricComputation(RecMetricComputation):
@@ -46,6 +47,7 @@ def __init__(self, *args: Any, **kwargs: Any) -> None:
             persistent=False,
         )
 
+    @pt2_compile_callable
     def update(
         self,
         *,

torchrec/metrics/precision.py

@@ -17,6 +17,7 @@
     RecMetricComputation,
     RecMetricException,
 )
+from torchrec.pt2.utils import pt2_compile_callable
 
 
 THRESHOLD = "threshold"
@@ -96,6 +97,7 @@ def __init__(self, *args: Any, threshold: float = 0.5, **kwargs: Any) -> None:
         )
         self._threshold: float = threshold
 
+    @pt2_compile_callable
     def update(
         self,
         *,

torchrec/metrics/rauc.py

@@ -21,6 +21,7 @@
     RecMetricComputation,
     RecMetricException,
 )
+from torchrec.pt2.utils import pt2_compile_callable
 
 PREDICTIONS = "predictions"
 LABELS = "labels"
@@ -287,6 +288,7 @@ def _init_states(self) -> None:
         if self._grouped_rauc:
             getattr(self, GROUPING_KEYS).append(torch.tensor([-1], device=self.device))
 
+    @pt2_compile_callable
     def update(
         self,
         *,

torchrec/metrics/rec_metric.py

@@ -46,6 +46,7 @@
     MetricNamespaceBase,
     MetricPrefix,
 )
+from torchrec.pt2.utils import pt2_compile_callable
 
 
 RecModelOutput = Union[torch.Tensor, Dict[str, torch.Tensor]]
@@ -136,6 +137,7 @@ def __init__(
         process_group: Optional[dist.ProcessGroup] = None,
         fused_update_limit: int = 0,
         allow_missing_label_with_zero_weight: bool = False,
+        enable_pt2_compile: bool = False,
         *args: Any,
         **kwargs: Any,
     ) -> None:
@@ -159,6 +161,7 @@ def __init__(
                 dist_reduce_fx=lambda x: torch.any(x, dim=0).byte(),
                 persistent=True,
             )
+        self.enable_pt2_compile = enable_pt2_compile
 
     @staticmethod
     def get_window_state_name(state_name: str) -> str:
@@ -244,6 +247,7 @@ def pre_compute(self) -> None:
         """
         return
 
+    @pt2_compile_callable
     def compute(self) -> List[MetricComputationReport]:
         with record_function(f"## {self.__class__.__name__}:compute ##"):
             if self._my_rank == 0 or self._compute_on_all_ranks:

torchrec/metrics/recall.py

@@ -17,6 +17,7 @@
     RecMetricComputation,
     RecMetricException,
 )
+from torchrec.pt2.utils import pt2_compile_callable
 
 
 THRESHOLD = "threshold"
@@ -96,6 +97,7 @@ def __init__(self, *args: Any, threshold: float = 0.5, **kwargs: Any) -> None:
         )
         self._threshold: float = threshold
 
+    @pt2_compile_callable
     def update(
         self,
         *,

torchrec/metrics/recall_session.py

@@ -21,6 +21,7 @@
     RecMetricComputation,
     RecMetricException,
 )
+from torchrec.pt2.utils import pt2_compile_callable
 
 logger: logging.Logger = logging.getLogger(__name__)
 
@@ -128,6 +129,7 @@ def __init__(
         self.run_ranking_of_labels: bool = session_metric_def.run_ranking_of_labels
         self.session_var_name: Optional[str] = session_metric_def.session_var_name
 
+    @pt2_compile_callable
     def update(
         self,
         *,

torchrec/metrics/scalar.py

@@ -17,6 +17,7 @@
     RecMetric,
     RecMetricComputation,
 )
+from torchrec.pt2.utils import pt2_compile_callable
 
 
 class ScalarMetricComputation(RecMetricComputation):
@@ -41,6 +42,7 @@ def __init__(self, *args: Any, **kwargs: Any) -> None:
             persistent=False,
         )
 
+    @pt2_compile_callable
     def update(
         self,
         *,

torchrec/metrics/segmented_ne.py

@@ -19,6 +19,7 @@
     RecMetricComputation,
     RecMetricException,
 )
+from torchrec.pt2.utils import pt2_compile_callable
 
 PREDICTIONS = "predictions"
 LABELS = "labels"
@@ -206,6 +207,7 @@ def __init__(
         )
         self.eta = 1e-12
 
+    @pt2_compile_callable
     def update(
         self,
         *,

torchrec/metrics/serving_calibration.py

@@ -18,6 +18,7 @@
     RecMetricComputation,
     RecMetricException,
 )
+from torchrec.pt2.utils import pt2_compile_callable
 
 CALIBRATION_NUM = "calibration_num"
 CALIBRATION_DENOM = "calibration_denom"
@@ -57,6 +58,7 @@ def __init__(self, *args: Any, **kwargs: Any) -> None:
             persistent=True,
         )
 
+    @pt2_compile_callable
     def update(
         self,
         *,