Categorize DDP comm-hook bucket all-reduce and TorchRec comm ops as NET

mosaic/libmosaic/utils/data_utils.py

@@ -123,6 +123,13 @@ def _no_frame_name_in(stack: List[Frame], names: frozenset) -> bool:
 
 _GRAD_SCALE_HELPER_NAMES: frozenset = frozenset({"_instantiate_filtered_grads"})
 
+# DDP comm-hook bucket all-reduce frames. Prefix match handles
+# demangled C++ symbols with parameter signatures.
+_NET_COMM_HOOK_PREFIXES: tuple = (
+    "c10d::Reducer::all_reduce_bucket",
+    "c10d::PythonCommHook::runHook",
+)
+
 
 class AllocationType(enum.Enum):
     PARAMETER = 0
@@ -199,6 +206,23 @@ def _matches_custom_pattern(cls, frame_stack: List[Frame], pattern: str) -> bool
         lambda s: _any_filename_contains(s, "fully_sharded_data_parallel.py"),
         AllocationType.FSDP,
     ),
+    # NET — DDP comm-hook and TorchRec comm frames. Must precede the
+    # autograd BACKWARD and "forward" ACTIVATION rules below, since real
+    # comm stacks carry both an autograd parent and a forward ancestor.
+    (
+        lambda s: any(f.name.startswith(_NET_COMM_HOOK_PREFIXES) for f in s),
+        AllocationType.NET,
+    ),
+    (
+        lambda s: _any_filename_contains(
+            s, "torch/distributed/algorithms/ddp_comm_hooks/"
+        ),
+        AllocationType.NET,
+    ),
+    (
+        lambda s: _any_filename_contains(s, "torchrec/distributed/comm_ops.py"),
+        AllocationType.NET,
+    ),
     # ACTIVATION — anything mentioning "forward" in its function name.
     (
         lambda s: _any_frame_name_contains(s, "forward"),

test/test_custom_profiling.py

@@ -607,6 +607,110 @@ def test_existing_clip_grad_norm_still_works(self) -> None:
         )
 
 
+class TestNetCommHookCategorization(TestCase):
+    """NET rules for DDP comm-hook and TorchRec comm frames. C++ frame
+    names use the demangled-with-signature form snapshots actually emit."""
+
+    def test_ddp_all_reduce_bucket_is_net(self) -> None:
+        frames = [
+            Frame(
+                name="c10d::Reducer::all_reduce_bucket(c10d::Reducer::Bucket&)",
+                filename="<invalid>",
+                line=0,
+            ),
+        ]
+        self.assertEqual(AllocationType.from_frame_stack(frames), AllocationType.NET)
+
+    def test_python_comm_hook_run_hook_is_net(self) -> None:
+        frames = [
+            Frame(
+                name="c10d::PythonCommHook::runHook(c10d::GradBucket&)",
+                filename="<invalid>",
+                line=0,
+            ),
+        ]
+        self.assertEqual(AllocationType.from_frame_stack(frames), AllocationType.NET)
+
+    def test_ddp_comm_hooks_filename_is_net(self) -> None:
+        frames = [
+            Frame(
+                name="default_hooks_allreduce",
+                filename="torch/distributed/algorithms/ddp_comm_hooks/default_hooks.py",
+                line=20,
+            ),
+        ]
+        self.assertEqual(AllocationType.from_frame_stack(frames), AllocationType.NET)
+
+    def test_torchrec_comm_ops_filename_is_net(self) -> None:
+        frames = [
+            Frame(
+                name="alltoall_pooled",
+                filename="torchrec/distributed/comm_ops.py",
+                line=300,
+            ),
+        ]
+        self.assertEqual(AllocationType.from_frame_stack(frames), AllocationType.NET)
+
+    def test_comm_hook_under_autograd_engine_is_net(self) -> None:
+        # Locks NET-before-BACKWARD ordering: real comm-hook stacks
+        # carry an autograd parent frame.
+        frames = [
+            Frame(
+                name="_compress_hook",
+                filename="torch/distributed/algorithms/ddp_comm_hooks/default_hooks.py",
+                line=71,
+            ),
+            Frame(
+                name="c10d::PythonCommHook::runHook(c10d::GradBucket&)",
+                filename="<invalid>",
+                line=0,
+            ),
+            Frame(
+                name="c10d::Reducer::all_reduce_bucket(c10d::Reducer::Bucket&)",
+                filename="<invalid>",
+                line=0,
+            ),
+            Frame(
+                name="torch::autograd::Engine::thread_main("
+                "std::shared_ptr<torch::autograd::GraphTask> const&)",
+                filename="<invalid>",
+                line=0,
+            ),
+        ]
+        self.assertEqual(AllocationType.from_frame_stack(frames), AllocationType.NET)
+
+    def test_torchrec_comm_ops_under_forward_is_net(self) -> None:
+        # Locks NET-before-ACTIVATION ordering: TorchRec all-to-all
+        # runs inside EBC.forward.
+        frames = [
+            Frame(
+                name="alltoall_pooled",
+                filename="torchrec/distributed/comm_ops.py",
+                line=524,
+            ),
+            Frame(
+                name="forward",
+                filename="torchrec/distributed/embeddingbag.py",
+                line=1124,
+            ),
+        ]
+        self.assertEqual(AllocationType.from_frame_stack(frames), AllocationType.NET)
+
+    def test_pure_autograd_engine_remains_backward(self) -> None:
+        # Negative: bare autograd stack stays BACKWARD.
+        frames = [
+            Frame(
+                name="torch::autograd::Engine::thread_main("
+                "std::shared_ptr<torch::autograd::GraphTask> const&)",
+                filename="<invalid>",
+                line=0,
+            ),
+        ]
+        self.assertEqual(
+            AllocationType.from_frame_stack(frames), AllocationType.BACKWARD
+        )
+
+
 class TestOmegaConfIntegration(TestCase):
     """Tests for OmegaConf integration with custom profiling"""