Skip shape exchange for fixed metric states (#4535)

Summary:

Reduce distributed RecMetric synchronization overhead when metric state shapes are fixed across ranks. Metrics with variable state shapes and custom synchronization retain their existing behavior.

Differential Revision: D115732360

Author: jeffkbkim

torchrec/metrics/rec_metric.py

@@ -10,11 +10,13 @@
 import abc
 import inspect
 import itertools
+import logging
 import math
 import weakref
 from collections import defaultdict, deque
 from dataclasses import dataclass
 from enum import Enum
+from functools import partial
 from typing import (
     Any,
     Callable,
@@ -38,6 +40,7 @@
 import torch.nn as nn
 from torch.profiler import record_function
 from torchmetrics import Metric
+from torchmetrics.utilities.distributed import gather_all_tensors
 
 try:
     from torchrec.distributed.logging_handlers import (
@@ -72,6 +75,8 @@ def n_batch_log_event(*args: object, **kwargs: object) -> None:
 
 RecModelOutput = Union[torch.Tensor, Dict[str, torch.Tensor]]
 
+logger: logging.Logger = logging.getLogger(__name__)
+
 
 @dataclass(frozen=True)
 class MetricComputationReport:
@@ -87,6 +92,8 @@ class MetricComputationReport:
 ]
 
 MAX_BUFFER_COUNT = 1000
+_FIXED_SHAPE_SYNC_JK = "pytorch/torchrec:enable_fixed_shape_metric_sync"
+_FIXED_SHAPE_GATHER = partial(gather_all_tensors, assume_same_shape=True)
 _WINDOW_BUFFER_REGISTRY: weakref.WeakValueDictionary[int, Any] = (
     torch.__dict__.setdefault(
         "_torchrec_window_buffer_registry", weakref.WeakValueDictionary()
@@ -112,6 +119,13 @@ def _window_buffer_aggregate_state(
     )
 
 
+def _supports_fixed_shape_sync(dist_sync_fn: Callable) -> bool:
+    try:
+        return "assume_same_shape" in inspect.signature(dist_sync_fn).parameters
+    except (TypeError, ValueError):
+        return False
+
+
 class RecMetricException(Exception):
     pass
 
@@ -229,6 +243,54 @@ def __init__(
                 persistent=True,
             )
         self._compute_mode: RecComputeMode = compute_mode
+        self._logged_dist_sync_paths: Set[str] = set()
+
+    def _state_shapes_match_defaults(self) -> bool:
+        defaults: Optional[Mapping[str, Any]] = getattr(self, "_defaults", None)
+        if not defaults:
+            return False
+        for name, default in defaults.items():
+            state = getattr(self, name)
+            if not isinstance(default, torch.Tensor) or not isinstance(
+                state, torch.Tensor
+            ):
+                return False
+            if state.shape != default.shape or state.dtype != default.dtype:
+                return False
+        return True
+
+    def _sync_dist(
+        self,
+        dist_sync_fn: Callable = gather_all_tensors,
+        process_group: Optional[Any] = None,
+    ) -> None:
+        uses_default_sync = dist_sync_fn is gather_all_tensors
+        uses_fixed_shape_sync = (
+            uses_default_sync
+            and _supports_fixed_shape_sync(dist_sync_fn)
+            and self._state_shapes_match_defaults()
+            and torch._utils_internal.justknobs_check(
+                _FIXED_SHAPE_SYNC_JK, default=True
+            )
+        )
+        if uses_fixed_shape_sync:
+            dist_sync_fn = _FIXED_SHAPE_GATHER
+            sync_path = "fixed_shape"
+        elif uses_default_sync:
+            sync_path = "variable_shape"
+        else:
+            sync_path = "custom"
+        if sync_path not in self._logged_dist_sync_paths:
+            if self._my_rank == 0:
+                logger.info(
+                    "RecMetric distributed sync path: metric=%s path=%s states=%d",
+                    type(self).__name__,
+                    sync_path,
+                    len(getattr(self, "_defaults", {})),
+                )
+            self._logged_dist_sync_paths.add(sync_path)
+
+        super()._sync_dist(dist_sync_fn, process_group)
 
     @staticmethod
     def get_window_state_name(state_name: str) -> str:

torchrec/metrics/tests/test_auc.py

@@ -8,7 +8,7 @@
 # pyre-strict
 
 import unittest
-from typing import Dict, Iterable, List, Optional, Type, Union
+from typing import cast, Dict, Iterable, List, Optional, Type, Union
 from unittest.mock import patch
 
 import torch
@@ -18,6 +18,7 @@
 from torchrec.metrics.rec_metric import (
     RecComputeMode,
     RecMetric,
+    RecMetricComputation,
     RecMetricException,
     RecTaskInfo,
 )
@@ -133,9 +134,10 @@ def test_auc_compile(self) -> None:
         # module attributes as dynamic to isolate the AUC list-length recompiles,
         # then assert the recompile limit is never hit so this stays a regression
         # guard.
-        with patch.object(
-            torch._dynamo.config, "allow_unspec_int_on_nn_module", True
-        ), patch.object(torch._dynamo.config, "fail_on_recompile_limit_hit", True):
+        with (
+            patch.object(torch._dynamo.config, "allow_unspec_int_on_nn_module", True),
+            patch.object(torch._dynamo.config, "fail_on_recompile_limit_hit", True),
+        ):
             for _ in range(10):
                 compiled_update(
                     predictions={DefaultTaskInfo.name: model_output["predictions"][0]},
@@ -167,6 +169,24 @@ class AUCMetricTest(unittest.TestCase):
     clazz: Type[RecMetric] = AUCMetric
     task_name: str = "auc"
 
+    def test_variable_shape_sync_path(self) -> None:
+        auc = AUCMetric(
+            world_size=1,
+            my_rank=0,
+            batch_size=1,
+            tasks=[DefaultTaskInfo],
+        )
+        computation = cast(RecMetricComputation, auc._metrics_computations[0])
+
+        with (
+            patch.object(torch.distributed, "get_world_size", return_value=1),
+            patch.object(torch.distributed, "barrier") as barrier,
+            patch.object(torch.distributed, "all_gather"),
+        ):
+            computation.sync(distributed_available=lambda: True)
+
+        barrier.assert_called()
+
     def test_unfused_auc(self) -> None:
         rec_metric_value_test_launcher(
             target_clazz=AUCMetric,

torchrec/metrics/tests/test_ne.py

@@ -9,16 +9,19 @@
 
 import unittest
 from functools import partial, update_wrapper
-from typing import Callable, Dict, Optional, Type
+from typing import Any, Callable, cast, Dict, List, Optional, Type
+from unittest.mock import patch
 
 import torch
+import torchrec.metrics.rec_metric as rec_metric
+from torchrec.metrics.metrics_config import DefaultTaskInfo
 from torchrec.metrics.ne import (
     compute_cross_entropy,
     compute_logloss,
     compute_ne,
     NEMetric,
 )
-from torchrec.metrics.rec_metric import RecComputeMode, RecMetric
+from torchrec.metrics.rec_metric import RecComputeMode, RecMetric, RecMetricComputation
 from torchrec.metrics.test_utils import (
     metric_test_helper,
     rec_metric_gpu_sync_test_launcher,
@@ -115,6 +118,149 @@ class NEMetricTest(unittest.TestCase):
     target_compute_mode: RecComputeMode = RecComputeMode.UNFUSED_TASKS_COMPUTATION
     task_name: str = "ne"
 
+    @staticmethod
+    def _set_distinct_sync_states(
+        computation: RecMetricComputation,
+    ) -> Dict[str, torch.Tensor]:
+        defaults = cast(Dict[str, Any], cast(Any, computation)._defaults)
+        expected: Dict[str, torch.Tensor] = {}
+        for value, name in enumerate(defaults, start=1):
+            state = cast(torch.Tensor, getattr(computation, name))
+            state.fill_(value)
+            expected[name] = state.clone()
+        return expected
+
+    def test_fixed_shape_sync_path(self) -> None:
+        ne = NEMetric(
+            world_size=1,
+            my_rank=0,
+            batch_size=1,
+            tasks=[DefaultTaskInfo],
+        )
+        computation = cast(RecMetricComputation, ne._metrics_computations[0])
+
+        with (
+            patch.object(torch._utils_internal, "justknobs_check", return_value=True),
+            patch.object(torch.distributed, "get_world_size", return_value=1),
+            patch.object(torch.distributed, "barrier") as barrier,
+            patch.object(torch.distributed, "all_gather") as all_gather,
+        ):
+            computation.sync(distributed_available=lambda: True)
+
+        barrier.assert_not_called()
+        self.assertGreater(all_gather.call_count, 0)
+
+    def test_fixed_shape_sync_killswitch_uses_variable_shape_path(self) -> None:
+        ne = NEMetric(
+            world_size=1,
+            my_rank=0,
+            batch_size=1,
+            tasks=[DefaultTaskInfo],
+        )
+        computation = cast(RecMetricComputation, ne._metrics_computations[0])
+
+        with (
+            patch.object(torch._utils_internal, "justknobs_check", return_value=False),
+            patch.object(torch.distributed, "get_world_size", return_value=1),
+            patch.object(torch.distributed, "barrier") as barrier,
+            patch.object(torch.distributed, "all_gather"),
+        ):
+            computation.sync(distributed_available=lambda: True)
+
+        barrier.assert_called()
+
+    def test_runtime_shape_change_uses_variable_shape_path(self) -> None:
+        ne = NEMetric(
+            world_size=1,
+            my_rank=0,
+            batch_size=1,
+            tasks=[DefaultTaskInfo],
+        )
+        computation = cast(RecMetricComputation, ne._metrics_computations[0])
+        cast(Any, computation).cross_entropy_sum = torch.zeros(2, dtype=torch.double)
+
+        with (
+            patch.object(torch._utils_internal, "justknobs_check", return_value=True),
+            patch.object(torch.distributed, "get_world_size", return_value=1),
+            patch.object(torch.distributed, "barrier") as barrier,
+            patch.object(torch.distributed, "all_gather"),
+        ):
+            computation.sync(distributed_available=lambda: True)
+
+        barrier.assert_called()
+
+    def test_missing_state_defaults_disables_fixed_shape_sync(self) -> None:
+        ne = NEMetric(
+            world_size=1,
+            my_rank=0,
+            batch_size=1,
+            tasks=[DefaultTaskInfo],
+        )
+        computation = cast(RecMetricComputation, ne._metrics_computations[0])
+        defaults = cast(Any, computation)._defaults
+        delattr(computation, "_defaults")
+        try:
+            self.assertFalse(computation._state_shapes_match_defaults())
+        finally:
+            cast(Any, computation)._defaults = defaults
+
+    def test_upstream_gather_without_fixed_shape_support_is_unchanged(self) -> None:
+        ne = NEMetric(
+            world_size=1,
+            my_rank=0,
+            batch_size=1,
+            tasks=[DefaultTaskInfo],
+        )
+        computation = cast(RecMetricComputation, ne._metrics_computations[0])
+        expected = self._set_distinct_sync_states(computation)
+        synced_tensors: List[torch.Tensor] = []
+
+        def upstream_gather(
+            tensor: torch.Tensor, group: Optional[Any] = None
+        ) -> List[torch.Tensor]:
+            synced_tensors.append(tensor)
+            return [tensor]
+
+        with patch.object(rec_metric, "gather_all_tensors", upstream_gather):
+            computation.sync(
+                dist_sync_fn=upstream_gather,
+                distributed_available=lambda: True,
+            )
+
+        self.assertEqual(len(synced_tensors), len(expected))
+        for actual, expected_tensor in zip(synced_tensors, expected.values()):
+            torch.testing.assert_close(actual, expected_tensor)
+        for name, expected_tensor in expected.items():
+            torch.testing.assert_close(getattr(computation, name), expected_tensor)
+
+    def test_custom_sync_path_is_unchanged(self) -> None:
+        ne = NEMetric(
+            world_size=1,
+            my_rank=0,
+            batch_size=1,
+            tasks=[DefaultTaskInfo],
+        )
+        computation = cast(RecMetricComputation, ne._metrics_computations[0])
+        expected = self._set_distinct_sync_states(computation)
+        synced_tensors: List[torch.Tensor] = []
+
+        def custom_sync(
+            tensor: torch.Tensor, group: Optional[Any] = None
+        ) -> List[torch.Tensor]:
+            synced_tensors.append(tensor)
+            return [tensor]
+
+        computation.sync(
+            dist_sync_fn=custom_sync,
+            distributed_available=lambda: True,
+        )
+
+        self.assertEqual(len(synced_tensors), len(expected))
+        for actual, expected_tensor in zip(synced_tensors, expected.values()):
+            torch.testing.assert_close(actual, expected_tensor)
+        for name, expected_tensor in expected.items():
+            torch.testing.assert_close(getattr(computation, name), expected_tensor)
+
     def test_ne_unfused(self) -> None:
         rec_metric_value_test_launcher(
             target_clazz=NEMetric,