Enabling DMP to init raw_id_tracker if configured (#3502)

Summary:

This diff introduces init_raw_id_tracker to initialize RawIdTracker if enabled through ModelTrackerConfig

Key Changes
-----------

1.  **Added `TrackerType` enum** in [`**types.py**`](command:code-compose.open?%5B%22%2Fdata%2Fusers%2Fmaliafzal%2Ffbsource%2Ffbcode%2Ftorchrec%2Fdistributed%2Fmodel_tracker%2Ftypes.py%22%2Cnull%5D "/data/users/maliafzal/fbsource/fbcode/torchrec/distributed/model_tracker/types.py")
    
    *   Defines three tracker types: `NONE`, `TREC` (for EC/EBC), and `RAW_ID` (for MPZCH)
    *   Extended `ModelTrackerConfig` dataclass with `tracker_type` field (defaults to `TrackerType.NONE`)
2.  **Enhanced DMP initialization** in [`**model_parallel.py**`](command:code-compose.open?%5B%22%2Fdata%2Fusers%2Fmaliafzal%2Ffbsource%2Ffbcode%2Ftorchrec%2Fdistributed%2Fmodel_parallel.py%22%2Cnull%5D "/data/users/maliafzal/fbsource/fbcode/torchrec/distributed/model_parallel.py")
    
    *   Added `init_raw_id_tracker()` method to create `RawIdTracker` instances
    *   Modified constructor to conditionally initialize trackers based on `tracker_type` configuration

internal 
This is needed to support MPZCH modules for Raw embedding streaming. 
Mode details : 
https://docs.google.com/document/d/1KEHwiXKLgXwRIdDFBYopjX3OiP3mRLM24Qkbiiu-TgE/edit?tab=t.0#bookmark=id.lhhgee2cs6ld

Reviewed By: chouxi

Differential Revision: D84920233

Author: aliafzal

torchrec/distributed/model_parallel.py

@@ -29,8 +29,16 @@
 from torch.nn.modules.module import _IncompatibleKeys
 from torch.nn.parallel import DistributedDataParallel
 from torchrec.distributed.comm import get_local_size
-from torchrec.distributed.model_tracker.model_delta_tracker import ModelDeltaTrackerTrec
-from torchrec.distributed.model_tracker.types import ModelTrackerConfig, UniqueRows
+from torchrec.distributed.model_tracker.model_delta_tracker import (
+    ModelDeltaTracker,
+    ModelDeltaTrackerTrec,
+)
+from torchrec.distributed.model_tracker.trackers.raw_id_tracker import RawIdTracker
+from torchrec.distributed.model_tracker.types import (
+    ModelTrackerConfig,
+    TrackerType,
+    UniqueRows,
+)
 
 from torchrec.distributed.planner import EmbeddingShardingPlanner, Topology
 from torchrec.distributed.sharding_plan import get_default_sharders
@@ -53,6 +61,7 @@
     none_throws,
     sharded_model_copy,
 )
+
 from torchrec.optim.fused import FusedOptimizerModule
 from torchrec.optim.keyed import CombinedOptimizer, KeyedOptimizer
 
@@ -294,12 +303,21 @@ def __init__(
         if init_data_parallel:
             self.init_data_parallel()
 
-        self.model_delta_tracker: Optional[ModelDeltaTrackerTrec] = (
+        self.model_delta_tracker: Optional[ModelDeltaTracker] = (
             self._init_delta_tracker(model_tracker_config, self._dmp_wrapped_module)
             if model_tracker_config is not None
+            and model_tracker_config.tracker_type == TrackerType.TREC
             else None
         )
 
+        if (
+            model_tracker_config is not None
+            and model_tracker_config.tracker_type == TrackerType.RAW_ID
+        ):
+            self.raw_id_tracker: Optional[RawIdTracker] = self.init_raw_id_tracker(
+                model_tracker_config, self._dmp_wrapped_module
+            )
+
     @property
     def module(self) -> nn.Module:
         """
@@ -370,7 +388,7 @@ def _init_dmp(self, module: nn.Module) -> nn.Module:
 
     def _init_delta_tracker(
         self, model_tracker_config: ModelTrackerConfig, module: nn.Module
-    ) -> ModelDeltaTrackerTrec:
+    ) -> ModelDeltaTracker:
         # Init delta tracker if config is provided
         return ModelDeltaTrackerTrec(
             model=module,
@@ -381,6 +399,15 @@ def _init_delta_tracker(
             fqns_to_skip=model_tracker_config.fqns_to_skip,
         )
 
+    def init_raw_id_tracker(
+        self, model_tracker_config: ModelTrackerConfig, module: nn.Module
+    ) -> Optional[RawIdTracker]:
+        return RawIdTracker(
+            model=module,
+            delete_on_read=model_tracker_config.delete_on_read,
+            fqns_to_skip=model_tracker_config.fqns_to_skip,
+        )
+
     def _init_optim(self, module: nn.Module) -> CombinedOptimizer:
         # pyre-ignore [6]
         return CombinedOptimizer(self._fused_optim_impl(module, []))
@@ -459,7 +486,7 @@ def init_parameters(module: nn.Module) -> None:
 
     def init_torchrec_delta_tracker(
         self, model_tracker_config: ModelTrackerConfig
-    ) -> ModelDeltaTrackerTrec:
+    ) -> ModelDeltaTracker:
         """
         Initializes the model delta tracker if it doesn't exists.
         """
@@ -470,7 +497,7 @@ def init_torchrec_delta_tracker(
 
         return none_throws(self.model_delta_tracker)
 
-    def get_model_tracker(self) -> ModelDeltaTrackerTrec:
+    def get_model_tracker(self) -> ModelDeltaTracker:
         """
         Returns the model tracker if it exists.
         """

torchrec/distributed/model_tracker/tests/test_model_delta_tracker.py

@@ -24,13 +24,18 @@
 from torchrec.distributed.embedding import EmbeddingCollectionSharder
 from torchrec.distributed.embedding_types import ModuleSharder, ShardingType
 from torchrec.distributed.embeddingbag import EmbeddingBagCollectionSharder
+from torchrec.distributed.model_tracker.model_delta_tracker import ModelDeltaTrackerTrec
 from torchrec.distributed.model_tracker.tests.utils import (
     EmbeddingTableProps,
     generate_planner_constraints,
     TestEBCModel,
     TestECModel,
 )
-from torchrec.distributed.model_tracker.types import ModelTrackerConfig, TrackingMode
+from torchrec.distributed.model_tracker.types import (
+    ModelTrackerConfig,
+    TrackerType,
+    TrackingMode,
+)
 
 from torchrec.distributed.planner import EmbeddingShardingPlanner, Topology
 from torchrec.distributed.test_utils.multi_process import (
@@ -228,7 +233,9 @@ def __init__(self, methodName="runTest") -> None:
                             sharding=ShardingType.ROW_WISE,
                         ),
                     ],
-                    model_tracker_config=ModelTrackerConfig(),
+                    model_tracker_config=ModelTrackerConfig(
+                        tracker_type=TrackerType.TREC
+                    ),
                 ),
                 FqnToFeatureNamesOutputTestParams(
                     expected_fqn_to_feature_names={
@@ -263,7 +270,9 @@ def __init__(self, methodName="runTest") -> None:
                             sharding=ShardingType.ROW_WISE,
                         ),
                     ],
-                    model_tracker_config=ModelTrackerConfig(),
+                    model_tracker_config=ModelTrackerConfig(
+                        tracker_type=TrackerType.TREC
+                    ),
                 ),
                 FqnToFeatureNamesOutputTestParams(
                     expected_fqn_to_feature_names={
@@ -296,7 +305,9 @@ def __init__(self, methodName="runTest") -> None:
                             sharding=ShardingType.ROW_WISE,
                         ),
                     ],
-                    model_tracker_config=ModelTrackerConfig(),
+                    model_tracker_config=ModelTrackerConfig(
+                        tracker_type=TrackerType.TREC,
+                    ),
                 ),
                 FqnToFeatureNamesOutputTestParams(
                     expected_fqn_to_feature_names={
@@ -332,7 +343,8 @@ def __init__(self, methodName="runTest") -> None:
                         ),
                     ],
                     model_tracker_config=ModelTrackerConfig(
-                        fqns_to_skip=["sparse_table_1"]
+                        fqns_to_skip=["sparse_table_1"],
+                        tracker_type=TrackerType.TREC,
                     ),
                 ),
                 FqnToFeatureNamesOutputTestParams(
@@ -368,7 +380,8 @@ def __init__(self, methodName="runTest") -> None:
                         ),
                     ],
                     model_tracker_config=ModelTrackerConfig(
-                        fqns_to_skip=["embedding_bags"]
+                        fqns_to_skip=["embedding_bags"],
+                        tracker_type=TrackerType.TREC,
                     ),
                 ),
                 FqnToFeatureNamesOutputTestParams(
@@ -399,7 +412,10 @@ def __init__(self, methodName="runTest") -> None:
                             sharding=ShardingType.ROW_WISE,
                         ),
                     ],
-                    model_tracker_config=ModelTrackerConfig(fqns_to_skip=["ec"]),
+                    model_tracker_config=ModelTrackerConfig(
+                        fqns_to_skip=["ec"],
+                        tracker_type=TrackerType.TREC,
+                    ),
                 ),
                 FqnToFeatureNamesOutputTestParams(
                     expected_fqn_to_feature_names={},
@@ -440,7 +456,9 @@ def test_fqn_to_feature_names(
                             sharding=ShardingType.ROW_WISE,
                         ),
                     ],
-                    model_tracker_config=ModelTrackerConfig(),
+                    model_tracker_config=ModelTrackerConfig(
+                        tracker_type=TrackerType.TREC,
+                    ),
                 ),
                 TrackerNotInitOutputTestParams(
                     dmp_tracker_atter="get_model_tracker",
@@ -461,7 +479,9 @@ def test_fqn_to_feature_names(
                             sharding=ShardingType.ROW_WISE,
                         ),
                     ],
-                    model_tracker_config=ModelTrackerConfig(),
+                    model_tracker_config=ModelTrackerConfig(
+                        tracker_type=TrackerType.TREC,
+                    ),
                 ),
                 TrackerNotInitOutputTestParams(
                     dmp_tracker_atter="get_unique",
@@ -514,6 +534,7 @@ def test_tracker_not_initialized(
                     model_tracker_config=ModelTrackerConfig(
                         tracking_mode=TrackingMode.ID_ONLY,
                         delete_on_read=True,
+                        tracker_type=TrackerType.TREC,
                     ),
                     model_inputs=[
                         ModelInput(
@@ -563,6 +584,7 @@ def test_tracker_not_initialized(
                     model_tracker_config=ModelTrackerConfig(
                         tracking_mode=TrackingMode.ID_ONLY,
                         delete_on_read=True,
+                        tracker_type=TrackerType.TREC,
                     ),
                     model_inputs=[
                         ModelInput(
@@ -610,6 +632,7 @@ def test_tracker_not_initialized(
                     model_tracker_config=ModelTrackerConfig(
                         tracking_mode=TrackingMode.ID_ONLY,
                         delete_on_read=True,
+                        tracker_type=TrackerType.TREC,
                     ),
                     model_inputs=[
                         ModelInput(
@@ -653,6 +676,7 @@ def test_tracker_not_initialized(
                     model_tracker_config=ModelTrackerConfig(
                         tracking_mode=TrackingMode.ID_ONLY,
                         delete_on_read=True,
+                        tracker_type=TrackerType.TREC,
                     ),
                     model_inputs=[
                         ModelInput(
@@ -715,6 +739,7 @@ def test_tracker_id_mode(
                     model_tracker_config=ModelTrackerConfig(
                         tracking_mode=TrackingMode.EMBEDDING,
                         delete_on_read=True,
+                        tracker_type=TrackerType.TREC,
                     ),
                     model_inputs=[
                         ModelInput(
@@ -763,6 +788,7 @@ def test_tracker_id_mode(
                     model_tracker_config=ModelTrackerConfig(
                         tracking_mode=TrackingMode.EMBEDDING,
                         delete_on_read=True,
+                        tracker_type=TrackerType.TREC,
                     ),
                     model_inputs=[
                         ModelInput(
@@ -808,6 +834,7 @@ def test_tracker_id_mode(
                     model_tracker_config=ModelTrackerConfig(
                         tracking_mode=TrackingMode.EMBEDDING,
                         delete_on_read=True,
+                        tracker_type=TrackerType.TREC,
                     ),
                     model_inputs=[
                         ModelInput(
@@ -881,6 +908,7 @@ def test_tracker_embedding_mode(
                     model_tracker_config=ModelTrackerConfig(
                         tracking_mode=TrackingMode.EMBEDDING,
                         delete_on_read=True,
+                        tracker_type=TrackerType.TREC,
                     ),
                     model_inputs=[
                         # First input: f1 has values 0,2,4,6 and f2 has values 8,10,12,14
@@ -975,6 +1003,7 @@ def test_tracker_embedding_mode(
                     model_tracker_config=ModelTrackerConfig(
                         tracking_mode=TrackingMode.ID_ONLY,
                         delete_on_read=True,
+                        tracker_type=TrackerType.TREC,
                     ),
                     model_inputs=[
                         # First input: f1 has values 0,2,4,6 and f2 has values 8,10,12,14
@@ -1088,6 +1117,7 @@ def test_multiple_get(
                         tracking_mode=TrackingMode.ID_ONLY,
                         delete_on_read=True,
                         consumers=["A", "B"],
+                        tracker_type=TrackerType.TREC,
                     ),
                     model_inputs=[
                         # First input: f1 has values 0,2,4,6 and f2 has values 8,10,12,14
@@ -1185,6 +1215,7 @@ def test_multiple_get(
                         tracking_mode=TrackingMode.ID_ONLY,
                         delete_on_read=False,
                         consumers=["A", "B"],
+                        tracker_type=TrackerType.TREC,
                     ),
                     model_inputs=[
                         # First input: f1 has values 0,2,4,6 and f2 has values 8,10,12,14
@@ -1282,6 +1313,7 @@ def test_multiple_get(
                         tracking_mode=TrackingMode.ID_ONLY,
                         delete_on_read=True,
                         consumers=["A", "B"],
+                        tracker_type=TrackerType.TREC,
                     ),
                     model_inputs=[
                         # First input: f1 has values 0,2,4,6 and f2 has values 8,10,12,14
@@ -1379,6 +1411,7 @@ def test_multiple_get(
                         tracking_mode=TrackingMode.ID_ONLY,
                         delete_on_read=False,
                         consumers=["A", "B"],
+                        tracker_type=TrackerType.TREC,
                     ),
                     model_inputs=[
                         # First input: f1 has values 0,2,4,6 and f2 has values 8,10,12,14
@@ -1486,6 +1519,7 @@ def test_multiple_consumers(
                     model_tracker_config=ModelTrackerConfig(
                         tracking_mode=TrackingMode.MOMENTUM_LAST,
                         delete_on_read=True,
+                        tracker_type=TrackerType.TREC,
                     ),
                     model_inputs=[
                         ModelInput(
@@ -1525,6 +1559,7 @@ def test_multiple_consumers(
                     model_tracker_config=ModelTrackerConfig(
                         tracking_mode=TrackingMode.MOMENTUM_LAST,
                         delete_on_read=True,
+                        tracker_type=TrackerType.TREC,
                     ),
                     model_inputs=[
                         ModelInput(
@@ -1581,6 +1616,7 @@ def test_duplication_with_momentum(
                     model_tracker_config=ModelTrackerConfig(
                         tracking_mode=TrackingMode.MOMENTUM_DIFF,
                         delete_on_read=True,
+                        tracker_type=TrackerType.TREC,
                     ),
                     model_inputs=[
                         ModelInput(
@@ -1620,6 +1656,7 @@ def test_duplication_with_momentum(
                     model_tracker_config=ModelTrackerConfig(
                         tracking_mode=TrackingMode.ROWWISE_ADAGRAD,
                         delete_on_read=True,
+                        tracker_type=TrackerType.TREC,
                     ),
                     model_inputs=[
                         ModelInput(
@@ -1679,6 +1716,7 @@ def _test_fqn_to_feature_names(
         )
 
         dt = dt_model.get_model_tracker()
+        assert isinstance(dt, ModelDeltaTrackerTrec)
         unittest.TestCase().assertEqual(
             dt.fqn_to_feature_names(), output_params.expected_fqn_to_feature_names
         )
@@ -1732,6 +1770,7 @@ def _test_id_mode(
         )
         features_list = model_input_generator(test_params.model_inputs, rank)
         dt = dt_model.get_model_tracker()
+        assert isinstance(dt, ModelDeltaTrackerTrec)
         for features in features_list:
             tracked_out = dt_model(features)
             baseline_out = baseline_model(features)
@@ -1826,6 +1865,7 @@ def _test_embedding_mode(
             # Only proceed with the rest of the test if models were created successfully
             features_list = model_input_generator(test_params.model_inputs, rank)
             dt = dt_model.get_model_tracker()
+            assert isinstance(dt, ModelDeltaTrackerTrec)
 
             orig_emb1 = (
                 # pyre-fixme[16]: Item `Tensor` of `Tensor | Module` has no attribute `ec`.
@@ -1939,6 +1979,7 @@ def _test_multiple_get(
         )
         features_list = model_input_generator(test_params.model_inputs, rank)
         dt = dt_model.get_model_tracker()
+        assert isinstance(dt, ModelDeltaTrackerTrec)
         table_fqns = dt.fqn_to_feature_names().keys()
         table_fqns_list = list(table_fqns)
         expected_emb1 = torch.tensor([])
@@ -2019,6 +2060,7 @@ def _test_multiple_consumer(
         )
         features_list = model_input_generator(test_params.model_inputs, rank)
         dt = dt_model.get_model_tracker()
+        assert isinstance(dt, ModelDeltaTrackerTrec)
         table_fqns = dt.fqn_to_feature_names().keys()
         table_fqns_list = list(table_fqns)
 
@@ -2082,6 +2124,7 @@ def _test_duplication_with_momentum(
         baseline_opt = torch.optim.Adam(baseline_model.parameters(), lr=0.1)
         features_list = model_input_generator(test_params.model_inputs, rank)
         dt = dt_model.get_model_tracker()
+        assert isinstance(dt, ModelDeltaTrackerTrec)
         table_fqns = dt.fqn_to_feature_names().keys()
         table_fqns_list = list(table_fqns)
         for features in features_list:
@@ -2147,6 +2190,7 @@ def _test_duplication_with_rowwise_adagrad(
         features_list = model_input_generator(test_params.model_inputs, rank)
 
         dt = dt_model.get_model_tracker()
+        assert isinstance(dt, ModelDeltaTrackerTrec)
         table_fqns = dt.fqn_to_feature_names().keys()
         table_fqns_list = list(table_fqns)