Categorize autograd engine, DDP reducer, and gradient helpers as BACKWARD

mosaic/libmosaic/utils/data_utils.py

@@ -97,6 +97,32 @@ def _no_frame_name_in(stack: List[Frame], names: frozenset) -> bool:
 )
 _STATS_NAMES: frozenset = frozenset({"compute_grad_stats"})
 
+# Prefix match: real C++ symbols carry parameter signatures.
+_C_AUTOGRAD_FRAME_PREFIXES: tuple = (
+    "torch::autograd::Engine::",
+    "torch::autograd::python::PythonEngine::",
+    "torch::autograd::Node::operator()",
+    "torch::autograd::PyNode::",
+    "torch::autograd::CppNode<",
+    "torch::autograd::GraphTask::",
+    "torch::autograd::AccumulateGrad::",
+    "torch::autograd::deleteNode",
+    "torch::autograd::utils::LambdaPostHook",
+)
+
+# DDP reducer bookkeeping (non-comm). Prefix also covers `_dense` overloads.
+_DDP_REDUCER_BACKWARD_PREFIXES: tuple = (
+    "c10d::Reducer::mark_variable_ready",
+    "c10d::Reducer::mark_bucket_ready",
+    "c10d::Reducer::autograd_hook",
+    "c10d::Reducer::initialize_buckets",
+    "c10d::Reducer::initialize_local_used_map",
+    "c10d::Reducer::rebuild_buckets",
+    "c10d::Reducer::set_static_graph",
+)
+
+_GRAD_SCALE_HELPER_NAMES: frozenset = frozenset({"_instantiate_filtered_grads"})
+
 
 class AllocationType(enum.Enum):
     PARAMETER = 0
@@ -202,6 +228,21 @@ def _matches_custom_pattern(cls, frame_stack: List[Frame], pattern: str) -> bool
         ),
         AllocationType.PARAMETER,
     ),
+    # BACKWARD — C++ autograd engine / graph-task frames.
+    (
+        lambda s: any(f.name.startswith(_C_AUTOGRAD_FRAME_PREFIXES) for f in s),
+        AllocationType.BACKWARD,
+    ),
+    # BACKWARD — DDP reducer bookkeeping. Allreduce comm is categorized as NET.
+    (
+        lambda s: any(f.name.startswith(_DDP_REDUCER_BACKWARD_PREFIXES) for f in s),
+        AllocationType.BACKWARD,
+    ),
+    # BACKWARD — Python grad-scaling helpers.
+    (
+        lambda s: _any_frame_name_in(s, _GRAD_SCALE_HELPER_NAMES),
+        AllocationType.BACKWARD,
+    ),
     # BACKWARD — anything mentioning "backward" in its function name.
     (
         lambda s: _any_frame_name_contains(s, "backward"),

test/test_custom_profiling.py

@@ -453,6 +453,160 @@ def test_per_category_alloc_sum_accepts_embedding_and_compile(self) -> None:
         )
 
 
+class TestBackwardCategorization(TestCase):
+    """BACKWARD-domain rules: C++ autograd, DDP reducer, grad-scale helpers.
+
+    Frame names use the demangled C++ symbol form with parameter signatures,
+    so prefix-vs-exact-match bugs cannot hide.
+    """
+
+    def _stack(self, name: str, filename: str = "<unknown>") -> list[Frame]:
+        return [Frame(name=name, filename=filename, line=0)]
+
+    def test_autograd_engine_evaluate_function_is_backward(self) -> None:
+        self.assertEqual(
+            AllocationType.from_frame_stack(
+                self._stack(
+                    "torch::autograd::Engine::evaluate_function("
+                    "std::shared_ptr<torch::autograd::GraphTask>&, "
+                    "torch::autograd::Node*, torch::autograd::InputBuffer&, "
+                    "std::shared_ptr<torch::autograd::ReadyQueue> const&)"
+                )
+            ),
+            AllocationType.BACKWARD,
+        )
+
+    def test_autograd_python_engine_thread_init_is_backward(self) -> None:
+        self.assertEqual(
+            AllocationType.from_frame_stack(
+                self._stack(
+                    "torch::autograd::python::PythonEngine::thread_init(int, "
+                    "std::shared_ptr<torch::autograd::ReadyQueue> const&, bool)"
+                )
+            ),
+            AllocationType.BACKWARD,
+        )
+
+    def test_autograd_node_operator_call_is_backward(self) -> None:
+        self.assertEqual(
+            AllocationType.from_frame_stack(
+                self._stack(
+                    "torch::autograd::Node::operator()("
+                    "std::vector<at::Tensor, std::allocator<at::Tensor>>&&)"
+                )
+            ),
+            AllocationType.BACKWARD,
+        )
+
+    def test_autograd_pynode_apply_is_backward(self) -> None:
+        self.assertEqual(
+            AllocationType.from_frame_stack(
+                self._stack(
+                    "torch::autograd::PyNode::apply("
+                    "std::vector<at::Tensor, std::allocator<at::Tensor>>&&)"
+                )
+            ),
+            AllocationType.BACKWARD,
+        )
+
+    def test_autograd_cppnode_apply_is_backward(self) -> None:
+        self.assertEqual(
+            AllocationType.from_frame_stack(
+                self._stack(
+                    "torch::autograd::CppNode<SomeBackwardFn>::apply("
+                    "std::vector<at::Tensor, std::allocator<at::Tensor>>&&)"
+                )
+            ),
+            AllocationType.BACKWARD,
+        )
+
+    def test_autograd_graph_task_post_processing_is_backward(self) -> None:
+        self.assertEqual(
+            AllocationType.from_frame_stack(
+                self._stack("torch::autograd::GraphTask::exec_post_processing()")
+            ),
+            AllocationType.BACKWARD,
+        )
+
+    def test_autograd_accumulate_grad_apply_is_backward(self) -> None:
+        self.assertEqual(
+            AllocationType.from_frame_stack(
+                self._stack(
+                    "torch::autograd::AccumulateGrad::apply("
+                    "std::vector<at::Tensor, std::allocator<at::Tensor>>&&)"
+                )
+            ),
+            AllocationType.BACKWARD,
+        )
+
+    def test_autograd_delete_node_is_backward(self) -> None:
+        self.assertEqual(
+            AllocationType.from_frame_stack(
+                self._stack("torch::autograd::deleteNode(torch::autograd::Node*)")
+            ),
+            AllocationType.BACKWARD,
+        )
+
+    def test_autograd_lambda_post_hook_is_backward(self) -> None:
+        self.assertEqual(
+            AllocationType.from_frame_stack(
+                self._stack(
+                    "torch::autograd::utils::LambdaPostHook::operator()("
+                    "std::vector<at::Tensor, std::allocator<at::Tensor>> const&, "
+                    "std::vector<at::Tensor, std::allocator<at::Tensor>> const&)"
+                )
+            ),
+            AllocationType.BACKWARD,
+        )
+
+    def test_ddp_reducer_mark_variable_ready_is_backward(self) -> None:
+        self.assertEqual(
+            AllocationType.from_frame_stack(
+                self._stack("c10d::Reducer::mark_variable_ready(unsigned long)")
+            ),
+            AllocationType.BACKWARD,
+        )
+
+    def test_ddp_reducer_mark_variable_ready_dense_is_backward(self) -> None:
+        # Overload variant exact-name matching would miss.
+        self.assertEqual(
+            AllocationType.from_frame_stack(
+                self._stack("c10d::Reducer::mark_variable_ready_dense(unsigned long)")
+            ),
+            AllocationType.BACKWARD,
+        )
+
+    def test_ddp_reducer_autograd_hook_is_backward(self) -> None:
+        self.assertEqual(
+            AllocationType.from_frame_stack(
+                self._stack("c10d::Reducer::autograd_hook(unsigned long)")
+            ),
+            AllocationType.BACKWARD,
+        )
+
+    def test_ddp_reducer_rebuild_buckets_is_backward(self) -> None:
+        self.assertEqual(
+            AllocationType.from_frame_stack(
+                self._stack("c10d::Reducer::rebuild_buckets()")
+            ),
+            AllocationType.BACKWARD,
+        )
+
+    def test_python_grad_scale_helper_is_backward(self) -> None:
+        self.assertEqual(
+            AllocationType.from_frame_stack(self._stack("_instantiate_filtered_grads")),
+            AllocationType.BACKWARD,
+        )
+
+    def test_existing_clip_grad_norm_still_works(self) -> None:
+        # Regression: pre-existing rule must still match after new rules are
+        # inserted ahead of it.
+        self.assertEqual(
+            AllocationType.from_frame_stack(self._stack("clip_grad_norm_")),
+            AllocationType.BACKWARD,
+        )
+
+
 class TestOmegaConfIntegration(TestCase):
     """Tests for OmegaConf integration with custom profiling"""