From cc97e93adda14bd74ab71493ccae8ddd9544ed09 Mon Sep 17 00:00:00 2001 From: Chenguang Zhu Date: Fri, 8 May 2026 23:52:15 -0700 Subject: [PATCH] Categorize DDP comm-hook bucket all-reduce and TorchRec comm ops as NET MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Summary: Add four new NET-domain rules to the categorization rule table: 1. `c10d::Reducer::all_reduce_bucket` — DDP bucket all-reduce. 2. `c10d::PythonCommHook::runHook` — Python comm-hook dispatch. 3. Filename contains `torch/distributed/algorithms/ddp_comm_hooks/` — any comm-hook implementation. 4. Filename contains `torchrec/distributed/comm_ops.py` — TorchRec distributed comm primitives (alltoall_pooled, etc.). Order-of-rules detail: these rules are inserted **after** the autograd-engine BACKWARD rules — so an autograd-bound bucket allocation still goes to BACKWARD — but **before** the broad `"backward" in frame.name` substring fallback, so a comm-hook frame in an otherwise generic backward stack still wins. Differential Revision: D103992608 --- mosaic/libmosaic/utils/data_utils.py | 24 +++++++ test/test_custom_profiling.py | 104 +++++++++++++++++++++++++++ 2 files changed, 128 insertions(+) diff --git a/mosaic/libmosaic/utils/data_utils.py b/mosaic/libmosaic/utils/data_utils.py index 63dce3f..a156005 100644 --- a/mosaic/libmosaic/utils/data_utils.py +++ b/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"), diff --git a/test/test_custom_profiling.py b/test/test_custom_profiling.py index 239e060..7c0a6e8 100644 --- a/test/test_custom_profiling.py +++ b/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="", + 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="", + 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="", + line=0, + ), + Frame( + name="c10d::Reducer::all_reduce_bucket(c10d::Reducer::Bucket&)", + filename="", + line=0, + ), + Frame( + name="torch::autograd::Engine::thread_main(" + "std::shared_ptr const&)", + filename="", + 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 const&)", + filename="", + line=0, + ), + ] + self.assertEqual( + AllocationType.from_frame_stack(frames), AllocationType.BACKWARD + ) + + class TestOmegaConfIntegration(TestCase): """Tests for OmegaConf integration with custom profiling"""