Allow torchrec delta tracker to be initialized post init (#3472)

Summary:
Pull Request resolved: #3472

internal

General Context: We are in the process of transition to a unified DeltaTracker and this is 6/n diffs representing  changes towards the transition.

Context: MRS DeltaTracker is initialized right before training, to allow for OSS DeltaTracker to have similar behavior adding a post DMP init function for initializing ModelDeltaTracker if not initialized.

Differential Revision: D80615308

fbshipit-source-id: 10c73274793580e2c1d0e7e6efd9377cab66fd99

Author: aliafzal

torchrec/distributed/model_parallel.py

@@ -50,6 +50,7 @@
     append_prefix,
     copy_to_device,
     filter_state_dict,
+    none_throws,
     sharded_model_copy,
 )
 from torchrec.optim.fused import FusedOptimizerModule
@@ -456,6 +457,19 @@ def init_parameters(module: nn.Module) -> None:
 
         module.apply(init_parameters)
 
+    def init_torchrec_delta_tracker(
+        self, model_tracker_config: ModelTrackerConfig
+    ) -> ModelDeltaTrackerTrec:
+        """
+        Initializes the model delta tracker if it doesn't exists.
+        """
+        if self.model_delta_tracker is None:
+            self.model_delta_tracker = self._init_delta_tracker(
+                model_tracker_config, self._dmp_wrapped_module
+            )
+
+        return none_throws(self.model_delta_tracker)
+
     def get_model_tracker(self) -> ModelDeltaTrackerTrec:
         """
         Returns the model tracker if it exists.