From 4a9d1acc92ef5924f325c67a5616ffa404ac09f1 Mon Sep 17 00:00:00 2001
From: Chenguang Zhu <cgzhu@meta.com>
Date: Thu, 7 May 2026 17:50:13 -0700
Subject: [PATCH] Add EMBEDDING and COMPILE AllocationType enum values
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Summary:
Enum-only diff. Adds two new categories to `AllocationType`:

- `EMBEDDING = 10` — sparse embedding parameters (TBE tables,
  sharded EBC init).
- `COMPILE = 11` — graph-compile-time allocations (Inductor /
  Dynamo / AOTAutograd).

No rules emit these values yet, so this diff is behavior-neutral.
`MemoryUsage.per_category_alloc_sum` is a `defaultdict(float)`
keyed by `AllocationType`, so it accepts the new keys naturally
without any additional wiring.

Differential Revision: D103992612
---
 mosaic/libmosaic/utils/data_utils.py |  4 ++++
 test/test_custom_profiling.py        | 35 ++++++++++++++++++++++++++++
 2 files changed, 39 insertions(+)

diff --git a/mosaic/libmosaic/utils/data_utils.py b/mosaic/libmosaic/utils/data_utils.py
index 4a37529..fbacd18 100644
--- a/mosaic/libmosaic/utils/data_utils.py
+++ b/mosaic/libmosaic/utils/data_utils.py
@@ -109,6 +109,10 @@ class AllocationType(enum.Enum):
     STATS = 7
     FSDP = 8
     CUSTOM = 9
+    # Sparse embedding parameters (TBE tables, sharded EBC init).
+    EMBEDDING = 10
+    # Graph-compile-time allocations (Inductor / Dynamo / AOTAutograd).
+    COMPILE = 11
 
     @classmethod
     def from_frame_stack(cls, frame_stack: List[Frame]) -> "AllocationType":
diff --git a/test/test_custom_profiling.py b/test/test_custom_profiling.py
index c4f4d57..b8dd7b4 100644
--- a/test/test_custom_profiling.py
+++ b/test/test_custom_profiling.py
@@ -418,6 +418,41 @@ def test_no_frame_name_in(self) -> None:
         self.assertTrue(_no_frame_name_in([], frozenset({"forward"})))
 
 
+class TestAllocationTypeEnumExtensions(TestCase):
+    """Tests for the EMBEDDING and COMPILE enum members."""
+
+    def test_embedding_member_exists_with_expected_value(self) -> None:
+        self.assertEqual(AllocationType.EMBEDDING.value, 10)
+        self.assertEqual(AllocationType.EMBEDDING.name, "EMBEDDING")
+
+    def test_compile_member_exists_with_expected_value(self) -> None:
+        self.assertEqual(AllocationType.COMPILE.value, 11)
+        self.assertEqual(AllocationType.COMPILE.name, "COMPILE")
+
+    def test_per_category_alloc_sum_accepts_embedding_and_compile(self) -> None:
+        """The MemoryUsage per-category dict is keyed by AllocationType, so
+        new enum values must be usable as keys without any extra wiring."""
+        memory_usage = MemoryUsage(save_profile=True)
+
+        for cat in (AllocationType.EMBEDDING, AllocationType.COMPILE):
+            evt = TraceEvent(
+                action="alloc",
+                addr=hash(cat),
+                size=4096,
+                stream=0,
+                time_us=0,
+                classification=cat,
+            )
+            memory_usage.update(evt, ["categories"])
+
+        self.assertEqual(
+            memory_usage.per_category_alloc_sum[AllocationType.EMBEDDING], 4096
+        )
+        self.assertEqual(
+            memory_usage.per_category_alloc_sum[AllocationType.COMPILE], 4096
+        )
+
+
 class TestOmegaConfIntegration(TestCase):
     """Tests for OmegaConf integration with custom profiling"""