From be9544cb48b66bce3f15a052c55857a0db5a15a6 Mon Sep 17 00:00:00 2001
From: Chenguang Zhu <cgzhu@meta.com>
Date: Tue, 19 May 2026 11:17:09 -0700
Subject: [PATCH] Categorize TorchRec / FBGEMM TBE init and sharding as
 EMBEDDING (#21)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Summary:

Adds two new rules:

1. **EMBEDDING name rule** — frame name in:
   `{_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}`.

2. **EMBEDDING filename rule** — filename contains any of:
   - 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

Note: `torchrec/distributed/model_parallel.py` is intentionally NOT in
the filename rule because it wraps the entire forward pass of every
sharded model — including all dense activations. Init paths through it
are already covered by the `_init_dmp` name rule.

Order-of-rules detail: inserted **before** the existing `"forward"`
substring rule (and therefore also ahead of all PARAMETER rules) so
that names like `init_parameters` (which contains "param") and any
TBE init that transitively traces forward stay categorized as
EMBEDDING.

Differential Revision: D103992615
---
 mosaic/libmosaic/utils/data_utils.py | 42 ++++++++++++++++++++++
 test/test_custom_profiling.py        | 52 ++++++++++++++++++++++++++++
 2 files changed, 94 insertions(+)

diff --git a/mosaic/libmosaic/utils/data_utils.py b/mosaic/libmosaic/utils/data_utils.py
index 3230515..2fc7e8d 100644
--- a/mosaic/libmosaic/utils/data_utils.py
+++ b/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"),
diff --git a/test/test_custom_profiling.py b/test/test_custom_profiling.py
index eab54bd..cbe9912 100644
--- a/test/test_custom_profiling.py
+++ b/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"""