Add async AR_COL for 2d sync (#4530)

torchrec/distributed/model_parallel.py

@@ -1146,6 +1146,9 @@ def __init__(
         self._pg: dist.ProcessGroup = global_pg
         self._global_rank: int = dist.get_rank(global_pg)
         self._custom_all_reduce = custom_all_reduce
+        self._all_reduce_hook: Optional[
+            Callable[[dist.ProcessGroup, List[torch.Tensor]], None]
+        ] = None
 
         if sharders is None:
             sharders = get_default_sharders()
@@ -1562,7 +1565,7 @@ def _sync(
         self._restore_stashed_sync_tensors(ctx, include_optimizer_state)
 
         opts = None
-        if self._custom_all_reduce is None:
+        if self._custom_all_reduce is None and self._all_reduce_hook is None:
             opts = dist.AllreduceCoalescedOptions()
             opts.reduceOp = dist.ReduceOp.AVG
 
@@ -1614,8 +1617,15 @@ def _allreduce_tensors(
             )
             return
 
+        all_reduce_hook = self._all_reduce_hook
         custom_all_reduce = self._custom_all_reduce
-        if custom_all_reduce is not None:
+        if all_reduce_hook is not None:
+
+            def _all_reduce(tensors: List[torch.Tensor]) -> None:
+                with record_function(f"{annotation}_custom_hook"):
+                    all_reduce_hook(ctx.replica_pg, tensors)
+
+        elif custom_all_reduce is not None:
             # Custom all reduce hook
             def _all_reduce(tensors: List[torch.Tensor]) -> None:
                 with record_function(f"{annotation}_custom_hook"):
@@ -1632,23 +1642,24 @@ def _all_reduce(tensors: List[torch.Tensor]) -> None:
 
     def set_all_reduce_hook(
         self,
-        reduce_hook: Callable[[List[torch.Tensor]], None],
+        reduce_hook: Callable[[dist.ProcessGroup, List[torch.Tensor]], None],
     ) -> None:
         """
-        Replace default all reduce with custom callable. Users can alternatively
-        pass in the custom all reduce function through the constructor. The hook
-        expects the user to handle distributed communication call, associated
-        process group, and stream synchronization.
+        Replace the default all reduce with a process-group-aware callable.
+        The hook must handle the distributed communication call and stream
+        synchronization.
 
         Args:
-            reduce_hook (Callable[[List[torch.Tensor]], torch.Tensor]): The custom all reduce function to use for
-                embedding weights and optimizer states
+            reduce_hook: Custom all reduce function for embedding weights and
+                optimizer states. It receives the replication process group and
+                tensors for the current DMP collection context.
         """
-        if self._custom_all_reduce is not None:
+        if self._custom_all_reduce is not None or self._all_reduce_hook is not None:
             logger.warning(
                 "[TorchRec 2D Parallel] Custom all reduce function already defined, overriding with new callable"
             )
-        self._custom_all_reduce = reduce_hook
+        self._custom_all_reduce = None
+        self._all_reduce_hook = reduce_hook
 
     def ensure_reduce_scatter_complete(self) -> None:
         """

torchrec/distributed/tests/test_dmp_collection.py

@@ -11,7 +11,9 @@
 from unittest.mock import MagicMock
 
 import torch
+import torch.distributed as dist
 import torch.nn as nn
+from torchrec.distributed.model_parallel import DMPCollection
 from torchrec.distributed.types import (
     DMPCollectionConfig,
     DMPCollectionContext,
@@ -257,6 +259,53 @@ def test_repr_excludes_runtime_fields(self) -> None:
         self.assertNotIn("sharded_module=", repr_str)
 
 
+class TestDMPCollectionAllReduceHook(unittest.TestCase):
+    def _dmp(self) -> DMPCollection:
+        dmp = DMPCollection.__new__(DMPCollection)
+        dmp._custom_all_reduce = None
+        dmp._all_reduce_hook = None
+        dmp._use_sharded_relay = False
+        dmp._sharded_relay_state = None
+        return dmp
+
+    def test_set_hook_receives_process_group_and_tensors(self) -> None:
+        dmp = self._dmp()
+        process_group = MagicMock(spec=dist.ProcessGroup)
+        context = MagicMock(spec=DMPCollectionContext)
+        context.replica_pg = process_group
+        tensors = [torch.ones(2)]
+        hook = MagicMock()
+
+        dmp.set_all_reduce_hook(hook)
+        dmp._allreduce_tensors(
+            context,
+            {torch.float32: tensors},
+            "test_all_reduce",
+        )
+
+        hook.assert_called_once_with(process_group, tensors)
+
+    def test_set_hook_overrides_legacy_constructor_hook(self) -> None:
+        dmp = self._dmp()
+        legacy_hook = MagicMock()
+        setter_hook = MagicMock()
+        process_group = MagicMock(spec=dist.ProcessGroup)
+        context = MagicMock(spec=DMPCollectionContext)
+        context.replica_pg = process_group
+        tensors = [torch.ones(2)]
+        dmp._custom_all_reduce = legacy_hook
+
+        dmp.set_all_reduce_hook(setter_hook)
+        dmp._allreduce_tensors(
+            context,
+            {torch.float32: tensors},
+            "test_all_reduce",
+        )
+
+        legacy_hook.assert_not_called()
+        setter_hook.assert_called_once_with(process_group, tensors)
+
+
 class TestShardingStrategy(unittest.TestCase):
 
     def test_strategy_values(self) -> None: