diff --git a/mosaic/libmosaic/utils/data_utils.py b/mosaic/libmosaic/utils/data_utils.py index fbacd18..63dce3f 100644 --- a/mosaic/libmosaic/utils/data_utils.py +++ b/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"), diff --git a/test/test_custom_profiling.py b/test/test_custom_profiling.py index b8dd7b4..239e060 100644 --- a/test/test_custom_profiling.py +++ b/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 = "") -> 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::Node*, torch::autograd::InputBuffer&, " + "std::shared_ptr 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 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>&&)" + ) + ), + 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>&&)" + ) + ), + AllocationType.BACKWARD, + ) + + def test_autograd_cppnode_apply_is_backward(self) -> None: + self.assertEqual( + AllocationType.from_frame_stack( + self._stack( + "torch::autograd::CppNode::apply(" + "std::vector>&&)" + ) + ), + 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>&&)" + ) + ), + 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> const&, " + "std::vector> 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"""