2/N: CPU Comms Module

Summary:
This diff adds the basic building blocks for a zero overhead RecMetrics implementation. Follow up patches will contain integration with users of torchrec.

One of the main pain points of using RecMetricModule is that metric updates and computes are done synchronously. In training jobs, there has been cases where metric updates take +20% of a training iteration. Metric computations, although less frequent, can takes over a couple of seconds.

CPUOffloadedRecMetricModule aims to perform all metric updates/computes asynchronously, completely removing them from the critical path.

This patch adds:
- CPUCommsRecMetricModule: Submodule that all gathers, loads, and computes aggregated metric state tensors across ranks.

Differential Revision: D83773528

Author: jeffkbkim

torchrec/metrics/cpu_comms_metric_module.py

@@ -0,0 +1,155 @@
+#!/usr/bin/env python3
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+# pyre-strict
+import logging
+from typing import Any, cast, Dict
+
+from torch import nn
+
+from torch.profiler import record_function
+
+from torchrec.metrics.metric_module import RecMetricModule
+from torchrec.metrics.metric_state_snapshot import MetricStateSnapshot
+from torchrec.metrics.rec_metric import (
+    RecComputeMode,
+    RecMetric,
+    RecMetricComputation,
+    RecMetricList,
+)
+
+logger: logging.Logger = logging.getLogger(__name__)
+
+
+class CPUCommsRecMetricModule(RecMetricModule):
+    """
+    A submodule of CPUOffloadedRecMetricModule.
+
+    The comms module's main purposes are:
+    1. All gather metric state tensors
+    2. Load all gathered metric states
+    3. Compute metrics
+
+    This isolation allows CPUOffloadedRecMetricModule from having
+    to concern about aggregated states and instead focus solely
+    updating local state tensors and dumping snapshots to the comms module
+    for metric aggregations.
+    """
+
+    def __init__(
+        self,
+        *args: Any,
+        **kwargs: Any,
+    ) -> None:
+        """
+        All arguments are the same as RecMetricModule
+        """
+
+        super().__init__(*args, **kwargs)
+
+        rec_metrics_clone = self._clone_rec_metrics()
+        self.rec_metrics: RecMetricList = rec_metrics_clone
+
+        for metric in self.rec_metrics.rec_metrics:
+            # Disable automatic sync for all metrics - handled manually via
+            # RecMetricModule.get_pre_compute_states()
+            metric = cast(RecMetric, metric)
+            for computation in metric._metrics_computations:
+                computation = cast(RecMetricComputation, computation)
+                computation._to_sync = False
+
+    def load_local_metric_state_snapshot(
+        self, state_snapshot: MetricStateSnapshot
+    ) -> None:
+        """
+        Load local metric states before all gather.
+        MetricStateSnapshot provides already-reduced states.
+
+        Args:
+            state_snapshot (MetricStateSnapshot): a snapshot of metric states to load.
+        """
+
+        # Load states into comms module to be shared across ranks.
+
+        with record_function("## CPUCommsRecMetricModule: load_snapshot ##"):
+            for metric in self.rec_metrics.rec_metrics:
+                metric = cast(RecMetric, metric)
+                compute_mode = metric._compute_mode
+                if (
+                    compute_mode == RecComputeMode.FUSED_TASKS_COMPUTATION
+                    or compute_mode == RecComputeMode.FUSED_TASKS_AND_STATES_COMPUTATION
+                ):
+                    prefix = compute_mode.name
+                    computation = metric._metrics_computations[0]
+                    self._load_metric_states(
+                        prefix, computation, state_snapshot.metric_states
+                    )
+                for task, computation in zip(
+                    metric._tasks, metric._metrics_computations
+                ):
+                    self._load_metric_states(
+                        task.name, computation, state_snapshot.metric_states
+                    )
+
+            if state_snapshot.throughput_metric is not None:
+                self.throughput_metric = state_snapshot.throughput_metric
+
+    def _load_metric_states(
+        self, prefix: str, computation: nn.Module, metric_states: Dict[str, Any]
+    ) -> None:
+        """
+        Load metric states after all gather.
+        Uses aggregated states.
+        """
+
+        # All update() calls were done prior. Clear previous computed state.
+        # Otherwise, we get warnings that compute() was called before
+        # update() which is not the case.
+        computation = cast(RecMetricComputation, computation)
+        computation._update_called = True
+        computation._computed = None
+
+        computation_name = f"{prefix}_{computation.__class__.__name__}"
+        # Restore all cached states from reductions
+        for attr_name in computation._reductions:
+            cache_key = f"{computation_name}_{attr_name}"
+            if cache_key in metric_states:
+                cached_value = metric_states[cache_key]
+                setattr(computation, attr_name, cached_value)
+
+    def _clone_rec_metrics(self) -> RecMetricList:
+        """
+        Clone rec_metrics. We need to keep references to the original tasks
+        and computation to load the state tensors. More importantly, we need to
+        remove the references to the original metrics to prevent concurrent access
+        from the update and compute threads.
+        """
+
+        cloned_metrics = []
+        for metric in self.rec_metrics.rec_metrics:
+            metric = cast(RecMetric, metric)
+            cloned_metric = type(metric)(
+                world_size=metric._world_size,
+                my_rank=metric._my_rank,
+                batch_size=metric._batch_size,
+                tasks=metric._tasks,
+                compute_mode=metric._compute_mode,
+                # Standard initialization passes in the global window size. A RecMetric's
+                # window size is set as the local window size.
+                window_size=metric._window_size * metric._world_size,
+                fused_update_limit=metric._fused_update_limit,
+                compute_on_all_ranks=metric._metrics_computations[
+                    0
+                ]._compute_on_all_ranks,
+                should_validate_update=metric._should_validate_update,
+                # Process group should be none to prevent unwanted distributed syncs.
+                # This is handled manually via RecMetricModule.get_pre_compute_states()
+                process_group=None,
+            )
+            cloned_metrics.append(cloned_metric)
+
+        return RecMetricList(cloned_metrics)