Categorize optimizer step and construction frames as OPTIMIZER

mosaic/libmosaic/utils/data_utils.py

@@ -14,6 +14,12 @@
 from dataclasses import dataclass, field, fields as dataclass_fields
 from typing import Any, Callable, List, Optional, Union
 
+# Extension hook: empty by default; downstream builds may override.
+from mosaic.libmosaic.utils.internal_optimizer_rules import (
+    INTERNAL_OPTIMIZER_FILENAME_SUBSTRINGS as _INTERNAL_OPTIMIZER_FILENAME_SUBSTRINGS,
+    INTERNAL_OPTIMIZER_NAMES as _INTERNAL_OPTIMIZER_NAMES,
+)
+
 NUM_GPUS_PER_HOST = 8
 
 logger: logging.Logger = logging.getLogger(__name__)
@@ -91,7 +97,34 @@ def _no_frame_name_in(stack: List[Frame], names: frozenset) -> bool:
 _BACKWARD_GRAD_HELPER_NAMES: frozenset = frozenset(
     {"clip_grad_norm_", "calc_grad_norm"}
 )
-_OPTIMIZER_VANILLA_NAMES: frozenset = frozenset({"custom_adamw", "_init_group"})
+_OPTIMIZER_VANILLA_NAMES: frozenset = (
+    frozenset(
+        {
+            "custom_adamw",
+            "_init_group",
+            # Optimizer step entry points.
+            "step",
+            "_per_group_step_impl",
+            # Optimizer construction.
+            "create_optimizer",
+            "_create_optimizer",
+            "create_keyed_optimizer",
+            # Distributed Shampoo preconditioner construction.
+            "_instantiate_shampoo_preconditioner_list",
+            "_create_kronecker_factors_state",
+            "_create_kronecker_factors_state_for_block",
+            "_create_base_kronecker_factors",
+        }
+    )
+    | _INTERNAL_OPTIMIZER_NAMES
+)
+
+# Filenames (substring match) that mark an allocation as belonging to an
+# optimizer implementation.
+_OPTIMIZER_FILENAME_SUBSTRINGS: tuple = (
+    "torch/optim/optimizer.py",
+    "torchrec/optim/keyed.py",
+) + _INTERNAL_OPTIMIZER_FILENAME_SUBSTRINGS
 _NET_AGGREGATE_NAMES: frozenset = frozenset(
     {"dist_max", "dist_mean", "dist_sum", "_aggregate"}
 )
@@ -282,6 +315,13 @@ def _matches_custom_pattern(cls, frame_stack: List[Frame], pattern: str) -> bool
         lambda s: _any_frame_name_in(s, _OPTIMIZER_VANILLA_NAMES),
         AllocationType.OPTIMIZER,
     ),
+    # OPTIMIZER — files implementing optimizer step / construction.
+    (
+        lambda s: any(
+            sub in f.filename for f in s for sub in _OPTIMIZER_FILENAME_SUBSTRINGS
+        ),
+        AllocationType.OPTIMIZER,
+    ),
     # NET — distributed aggregation primitives.
     (
         lambda s: _any_frame_name_in(s, _NET_AGGREGATE_NAMES),

mosaic/libmosaic/utils/internal_optimizer_rules.py

@@ -0,0 +1,17 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+# pyre-strict
+
+"""Extension hook for additional OPTIMIZER categorization rules.
+
+The constants exported here are empty by default. Downstream builds may
+override this module to inject extra optimizer filename substrings or
+frame names without forking ``data_utils.py``.
+"""
+
+INTERNAL_OPTIMIZER_FILENAME_SUBSTRINGS: tuple = ()
+INTERNAL_OPTIMIZER_NAMES: frozenset = frozenset()

test/test_custom_profiling.py

@@ -711,6 +711,77 @@ def test_pure_autograd_engine_remains_backward(self) -> None:
         )
 
 
+class TestOptimizerCategorization(TestCase):
+    """Tests for OPTIMIZER-domain rules: optimizer step / construction
+    frames and Distributed Shampoo / TorchRec keyed-optimizer files."""
+
+    def test_distributed_shampoo_step_is_optimizer(self) -> None:
+        # Shampoo's step is a `step` frame inside the shampoo file.
+        frames = [
+            Frame(
+                name="step",
+                filename="distributed_shampoo/shampoo.py",
+                line=500,
+            ),
+        ]
+        self.assertEqual(
+            AllocationType.from_frame_stack(frames), AllocationType.OPTIMIZER
+        )
+
+    def test_vanilla_torch_optim_step_is_optimizer(self) -> None:
+        # A generic torch.optim Optimizer.step() frame.
+        frames = [
+            Frame(
+                name="step",
+                filename="torch/optim/optimizer.py",
+                line=180,
+            ),
+        ]
+        self.assertEqual(
+            AllocationType.from_frame_stack(frames), AllocationType.OPTIMIZER
+        )
+
+    def test_torchrec_keyed_optimizer_step_is_optimizer(self) -> None:
+        # TorchRec keyed optimizer step (matches via filename rule).
+        frames = [
+            Frame(
+                name="some_internal_helper",
+                filename="torchrec/optim/keyed.py",
+                line=100,
+            ),
+        ]
+        self.assertEqual(
+            AllocationType.from_frame_stack(frames), AllocationType.OPTIMIZER
+        )
+
+    def test_shampoo_preconditioner_construction_is_optimizer(self) -> None:
+        # Shampoo preconditioner construction frame.
+        frames = [
+            Frame(
+                name="_instantiate_shampoo_preconditioner_list",
+                filename="distributed_shampoo/preconditioner.py",
+                line=42,
+            ),
+        ]
+        self.assertEqual(
+            AllocationType.from_frame_stack(frames), AllocationType.OPTIMIZER
+        )
+
+    def test_existing_init_group_still_works(self) -> None:
+        # Regression check: the pre-existing rule for _init_group must
+        # continue to map to OPTIMIZER even with the expanded name set.
+        frames = [
+            Frame(
+                name="_init_group",
+                filename="torch/optim/adam.py",
+                line=60,
+            ),
+        ]
+        self.assertEqual(
+            AllocationType.from_frame_stack(frames), AllocationType.OPTIMIZER
+        )
+
+
 class TestOmegaConfIntegration(TestCase):
     """Tests for OmegaConf integration with custom profiling"""