Categorize TorchRec / FBGEMM TBE init and sharding as EMBEDDING (#21)

mosaic/libmosaic/utils/data_utils.py

@@ -163,6 +163,36 @@ def _no_frame_name_in(stack: List[Frame], names: frozenset) -> bool:
     "c10d::PythonCommHook::runHook",
 )
 
+# Frame names that mark an allocation as EMBEDDING (sparse params, distinct
+# from dense PARAMETER).
+_EMBEDDING_INIT_NAMES: frozenset = frozenset(
+    {
+        "_init_dmp",
+        "_shard_modules_impl",
+        "_create_shard_module",
+        "shard",
+        "apply_2d_emb_sharding",
+        "_create_lookups",
+        "create_lookup",
+        "_create_embedding_kernel",
+        "init_parameters",
+        "_init_fbgemm_regroup",
+        "_apply_split",
+        "_apply_cache_state",
+    }
+)
+
+# Filenames (substring match) that mark an allocation as EMBEDDING.
+_EMBEDDING_FILENAME_SUBSTRINGS: tuple = (
+    "torchrec/distributed/embedding_lookup.py",
+    "torchrec/distributed/embeddingbag.py",
+    "torchrec/distributed/batched_embedding_kernel.py",
+    "torchrec/distributed/embedding_sharding.py",
+    "torchrec/distributed/sharding/",
+    "fbgemm_gpu/split_table_batched_embeddings_ops_training.py",
+    "fbgemm_gpu/split_table_batched_embeddings_ops_training_common.py",
+)
+
 
 class AllocationType(enum.Enum):
     PARAMETER = 0
@@ -256,6 +286,18 @@ def _matches_custom_pattern(cls, frame_stack: List[Frame], pattern: str) -> bool
         lambda s: _any_filename_contains(s, "torchrec/distributed/comm_ops.py"),
         AllocationType.NET,
     ),
+    # EMBEDDING — must precede ACTIVATION/PARAMETER so `init_parameters`
+    # and init paths that trace through forward stay EMBEDDING.
+    (
+        lambda s: _any_frame_name_in(s, _EMBEDDING_INIT_NAMES),
+        AllocationType.EMBEDDING,
+    ),
+    (
+        lambda s: any(
+            sub in f.filename for f in s for sub in _EMBEDDING_FILENAME_SUBSTRINGS
+        ),
+        AllocationType.EMBEDDING,
+    ),
     # ACTIVATION — anything mentioning "forward" in its function name.
     (
         lambda s: _any_frame_name_contains(s, "forward"),

test/test_custom_profiling.py

@@ -782,6 +782,58 @@ def test_existing_init_group_still_works(self) -> None:
         )
 
 
+class TestEmbeddingCategorization(TestCase):
+    """Tests for EMBEDDING categorization rules."""
+
+    def test_init_dmp_is_embedding(self) -> None:
+        frames = [
+            Frame(
+                name="_init_dmp",
+                filename="torchrec/distributed/model_parallel.py",
+                line=300,
+            ),
+        ]
+        self.assertEqual(
+            AllocationType.from_frame_stack(frames), AllocationType.EMBEDDING
+        )
+
+    def test_apply_2d_emb_sharding_is_embedding(self) -> None:
+        frames = [
+            Frame(
+                name="apply_2d_emb_sharding",
+                filename="torchrec/distributed/sharding/cw_sharding.py",
+                line=120,
+            ),
+        ]
+        self.assertEqual(
+            AllocationType.from_frame_stack(frames), AllocationType.EMBEDDING
+        )
+
+    def test_fbgemm_tbe_init_is_embedding(self) -> None:
+        frames = [
+            Frame(
+                name="reset_uvm_cache_stats",
+                filename=("fbgemm_gpu/split_table_batched_embeddings_ops_training.py"),
+                line=900,
+            ),
+        ]
+        self.assertEqual(
+            AllocationType.from_frame_stack(frames), AllocationType.EMBEDDING
+        )
+
+    def test_torchrec_ebc_sharding_is_embedding(self) -> None:
+        frames = [
+            Frame(
+                name="ShardedEmbeddingBagCollection_init",
+                filename="torchrec/distributed/embeddingbag.py",
+                line=80,
+            ),
+        ]
+        self.assertEqual(
+            AllocationType.from_frame_stack(frames), AllocationType.EMBEDDING
+        )
+
+
 class TestOmegaConfIntegration(TestCase):
     """Tests for OmegaConf integration with custom profiling"""