Enabling DMP to init raw_id_tracker if configured (meta-pytorch#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

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/types.py

@@ -75,6 +75,20 @@ class UpdateMode(Enum):
     LAST = "last"
 
 
+class TrackerType(Enum):
+    r"""
+    To identify the type of tracker.
+
+    Enums:
+    TREC: Used for Trec EC/EBC.
+    RAW_ID: Used for MPZCH.
+    """
+
+    NONE = "none"
+    TREC = "trec"
+    RAW_ID = "raw_id"
+
+
 @dataclass
 class ModelTrackerConfig:
     r"""
@@ -92,3 +106,4 @@ class ModelTrackerConfig:
     delete_on_read: bool = True
     auto_compact: bool = False
     fqns_to_skip: List[str] = field(default_factory=list)
+    tracker_type: TrackerType = TrackerType.NONE