Adding support for tracking optimizers states in Model Delta Tracker.

Summary:
### Overview

This diff adds support for tracking optimizer states  in the Model Delta Tracker system. It introduces a new tracking mode called `MOMENTUM_LAST` that enables tracking of momentum values from optimizers to support approximate top-k delta-row selection.

### Key Changes

#### 1. Optimizer State Tracking Support

*   To support tracking of optimizer states I have added `optim_state_tracker_fn` attribute to `GroupedEmbeddingsLookup` and `GroupedPooledEmbeddingsLookup` classes responsible for traversing over the BatchedFused modules.
*   Implemented `register_optim_state_tracker_fn()` method in both classes to register the trackable callable
*   Tracking calls are invoked after each lookup operation.

#### 2. Model Delta Tracker Changes

*   Added `record_momentum()` method to track momentum values from optimizer states and its support in record_lookup function.
*   Added validation and optim tracker function logic to support the new `MOMENTUM_LAST` mode

#### 3. New Tracking Mode

*   Added `TrackingMode.MOMENTUM_LAST` to [`**types.py**`](command:code-compose.open?%5B%22%2Ffbcode%2Ftorchrec%2Fdistributed%2Fmodel_tracker%2Ftypes.py%22%2Cnull%5D "/fbcode/torchrec/distributed/model_tracker/types.py")
*   Maps to `EmbdUpdateMode.LAST` to capture the most recent momentum values

Differential Revision: D76868111

Author: aliafzal

torchrec/distributed/embedding.py

@@ -1515,7 +1515,7 @@ def compute_and_output_dist(
             ):
                 embs = lookup(features)
                 if self.post_lookup_tracker_fn is not None:
-                    self.post_lookup_tracker_fn(features, embs)
+                    self.post_lookup_tracker_fn(self, features, embs)
 
             with maybe_annotate_embedding_event(
                 EmbeddingEvent.OUTPUT_DIST, self._module_fqn, sharding_type

torchrec/distributed/embedding_lookup.py

@@ -10,7 +10,7 @@
 import logging
 from abc import ABC
 from collections import OrderedDict
-from typing import Any, cast, Dict, Iterator, List, Optional, Tuple, Union
+from typing import Any, Callable, cast, Dict, Iterator, List, Optional, Tuple, Union
 
 import torch
 import torch.distributed as dist
@@ -206,6 +206,10 @@ def __init__(
         )
 
         self.grouped_configs = grouped_configs
+        # Model tracker function to tracker optimizer state
+        self.optim_state_tracker_fn: Optional[
+            Callable[[nn.Module, KeyedJaggedTensor, torch.Tensor], None]
+        ] = None
 
     def _create_embedding_kernel(
         self,
@@ -305,7 +309,13 @@ def forward(
             self._feature_splits,
         )
         for emb_op, features in zip(self._emb_modules, features_by_group):
-            embeddings.append(emb_op(features).view(-1))
+            lookup = emb_op(features).view(-1)
+            embeddings.append(lookup)
+
+            # Model tracker optimizer state function, will only be set called
+            # when model tracker is configured to track optimizer state
+            if self.optim_state_tracker_fn is not None:
+                self.optim_state_tracker_fn(emb_op, features, lookup)
 
         return embeddings_cat_empty_rank_handle(embeddings, self._dummy_embs_tensor)
 
@@ -409,6 +419,19 @@ def purge(self) -> None:
             # pyre-fixme[29]: `Union[Module, Tensor]` is not a function.
             emb_module.purge()
 
+    def register_optim_state_tracker_fn(
+        self,
+        record_fn: Callable[[nn.Module, KeyedJaggedTensor, torch.Tensor], None],
+    ) -> None:
+        """
+        Model tracker function to tracker optimizer state
+
+         Args:
+             record_fn (Callable[[nn.Module, KeyedJaggedTensor, torch.Tensor], None]): A custom record function to be called after lookup is done.
+
+        """
+        self.optim_state_tracker_fn = record_fn
+
 
 class CommOpGradientScaling(torch.autograd.Function):
     @staticmethod
@@ -481,6 +504,10 @@ def __init__(
             if scale_weight_gradients and get_gradient_division()
             else 1
         )
+        # Model tracker function to tracker optimizer state
+        self.optim_state_tracker_fn: Optional[
+            Callable[[nn.Module, KeyedJaggedTensor, torch.Tensor], None]
+        ] = None
 
     def _create_embedding_kernel(
         self,
@@ -608,7 +635,12 @@ def forward(
                         features._weights, self._scale_gradient_factor
                     )
 
-                embeddings.append(emb_op(features))
+                lookup = emb_op(features)
+                embeddings.append(lookup)
+                # Model tracker optimizer state function, will only be set called
+                # when model tracker is configured to track optimizer state
+                if self.optim_state_tracker_fn is not None:
+                    self.optim_state_tracker_fn(emb_op, features, lookup)
 
                 if features.variable_stride_per_key() and len(self._emb_modules) > 1:
                     stride_per_rank_per_key = list(
@@ -738,6 +770,19 @@ def purge(self) -> None:
             # pyre-fixme[29]: `Union[Module, Tensor]` is not a function.
             emb_module.purge()
 
+    def register_optim_state_tracker_fn(
+        self,
+        record_fn: Callable[[nn.Module, KeyedJaggedTensor, torch.Tensor], None],
+    ) -> None:
+        """
+        Model tracker function to tracker optimizer state
+
+         Args:
+             record_fn (Callable[[nn.Module, KeyedJaggedTensor, torch.Tensor], None]): A custom record function to be called after lookup is done.
+
+        """
+        self.optim_state_tracker_fn = record_fn
+
 
 class MetaInferGroupedEmbeddingsLookup(
     BaseEmbeddingLookup[KeyedJaggedTensor, torch.Tensor], TBEToRegisterMixIn

torchrec/distributed/embedding_types.py

@@ -373,7 +373,7 @@ def __init__(
         self._lookups: List[nn.Module] = []
         self._output_dists: List[nn.Module] = []
         self.post_lookup_tracker_fn: Optional[
-            Callable[[KeyedJaggedTensor, torch.Tensor], None]
+            Callable[[nn.Module, KeyedJaggedTensor, torch.Tensor], None]
         ] = None
         self.post_odist_tracker_fn: Optional[Callable[..., None]] = None
 
@@ -426,14 +426,14 @@ def train(self, mode: bool = True):  # pyre-ignore[3]
 
     def register_post_lookup_tracker_fn(
         self,
-        record_fn: Callable[[KeyedJaggedTensor, torch.Tensor], None],
+        record_fn: Callable[[nn.Module, KeyedJaggedTensor, torch.Tensor], None],
     ) -> None:
         """
         Register a function to be called after lookup is done. This is used for
         tracking the lookup results and optimizer states.
 
         Args:
-            record_fn (Callable[[KeyedJaggedTensor, torch.Tensor], None]): A custom record function to be called after lookup is done.
+            record_fn (Callable[[nn.Module, KeyedJaggedTensor, torch.Tensor], None]): A custom record function to be called after lookup is done.
 
         """
         if self.post_lookup_tracker_fn is not None:

torchrec/distributed/embeddingbag.py

@@ -1459,7 +1459,7 @@ def compute_and_output_dist(
             ):
                 embs = lookup(features)
                 if self.post_lookup_tracker_fn is not None:
-                    self.post_lookup_tracker_fn(features, embs)
+                    self.post_lookup_tracker_fn(self, features, embs)
 
             with maybe_annotate_embedding_event(
                 EmbeddingEvent.OUTPUT_DIST,

torchrec/distributed/model_tracker/model_delta_tracker.py

@@ -13,7 +13,12 @@
 import torch
 
 from torch import nn
+
 from torchrec.distributed.embedding import ShardedEmbeddingCollection
+from torchrec.distributed.embedding_lookup import (
+    GroupedEmbeddingsLookup,
+    GroupedPooledEmbeddingsLookup,
+)
 from torchrec.distributed.embeddingbag import ShardedEmbeddingBagCollection
 from torchrec.distributed.model_tracker.delta_store import DeltaStore
 from torchrec.distributed.model_tracker.types import (
@@ -27,9 +32,16 @@
     # Only IDs are tracked, no additional state is stored.
     TrackingMode.ID_ONLY: EmbdUpdateMode.NONE,
     # TrackingMode.EMBEDDING utilizes EmbdUpdateMode.FIRST to ensure that
-    # the earliest embedding values are stored since the last checkpoint or snapshot.
-    # This mode is used for computing topk delta rows, which is currently achieved by running (new_emb - old_emb).norm().topk().
+    # the earliest embedding values are stored since the last checkpoint
+    # or snapshot. This mode is used for computing topk delta rows, which
+    # is currently achieved by running (new_emb - old_emb).norm().topk().
     TrackingMode.EMBEDDING: EmbdUpdateMode.FIRST,
+    # TrackingMode.MOMENTUM utilizes EmbdUpdateMode.LAST to ensure that
+    # the most recent momentum values—capturing the accumulated gradient
+    # direction and magnitude—are stored since the last batch.
+    # This mode supports approximate top-k delta-row selection, can be
+    # obtained by running momentum.norm().topk().
+    TrackingMode.MOMENTUM_LAST: EmbdUpdateMode.LAST,
 }
 
 # Tracking is current only supported for ShardedEmbeddingCollection and ShardedEmbeddingBagCollection.
@@ -141,7 +153,9 @@ def trigger_compaction(self) -> None:
             # Update the current compact index to the end index to avoid duplicate compaction.
             self.curr_compact_index = end_idx
 
-    def record_lookup(self, kjt: KeyedJaggedTensor, states: torch.Tensor) -> None:
+    def record_lookup(
+        self, emb_module: nn.Module, kjt: KeyedJaggedTensor, states: torch.Tensor
+    ) -> None:
         """
         Records the IDs from a given KeyedJaggedTensor and their corresponding embeddings/parameter states.
 
@@ -152,6 +166,7 @@ def record_lookup(self, kjt: KeyedJaggedTensor, states: torch.Tensor) -> None:
         (in ID_ONLY mode) or both IDs and their corresponding embeddings (in EMBEDDING mode).
 
         Args:
+            emb_module (nn.Module): The embedding module in which the lookup was performed.
             kjt (KeyedJaggedTensor): The KeyedJaggedTensor containing IDs to record.
             states (torch.Tensor): The embeddings or states corresponding to the IDs in the kjt.
         """
@@ -162,7 +177,9 @@ def record_lookup(self, kjt: KeyedJaggedTensor, states: torch.Tensor) -> None:
         # In EMBEDDING mode, we track per feature IDs and corresponding embeddings received in the current batch.
         elif self._mode == TrackingMode.EMBEDDING:
             self.record_embeddings(kjt, states)
-
+        # In MOMENTUM_LAST mode, we track per feature IDs and corresponding momentum values received in the current batch.
+        elif self._mode == TrackingMode.MOMENTUM_LAST:
+            self.record_momentum(emb_module, kjt)
         else:
             raise NotImplementedError(f"Tracking mode {self._mode} is not supported")
 
@@ -228,6 +245,39 @@ def record_embeddings(
                 states=torch.cat(per_table_emb[table_fqn]),
             )
 
+    def record_momentum(
+        self,
+        emb_module: nn.Module,
+        kjt: KeyedJaggedTensor,
+    ) -> None:
+        # FIXME: this is the momentum from last iteration, use momentum from current iter
+        # for correctness.
+        # pyre-ignore Undefined attribute [16]:
+        momentum = emb_module._emb_module.momentum1_dev
+        # FIXME: support multiple tables per group, information can be extracted from
+        # module._config (i.e., GroupedEmbeddingConfig)
+        # pyre-ignore Undefined attribute [16]:
+        states = momentum.view(-1, emb_module._config.embedding_dims()[0])[
+            kjt.values()
+        ].norm(dim=1)
+
+        offsets: torch.Tensor = torch.ops.fbgemm.asynchronous_complete_cumsum(
+            torch.tensor(kjt.length_per_key(), dtype=torch.int64)
+        )
+        assert (
+            kjt.values().numel() == states.numel()
+        ), f"number of ids and states mismatch, expect {kjt.values()=}, {kjt.values().numel()}, but got {states.numel()} "
+
+        for i, key in enumerate(kjt.keys()):
+            fqn = self.feature_to_fqn[key]
+            per_key_states = states[offsets[i] : offsets[i + 1]]
+            self.store.append(
+                batch_idx=self.curr_batch_idx,
+                table_fqn=fqn,
+                ids=kjt[key].values(),
+                states=per_key_states,
+            )
+
     def get_delta_ids(self, consumer: Optional[str] = None) -> Dict[str, torch.Tensor]:
         """
         Return a dictionary of hit local IDs for each sparse feature. Ids are
@@ -380,13 +430,31 @@ def _clean_fqn_fn(self, fqn: str) -> str:
     def _validate_and_init_tracker_fns(self) -> None:
         "To validate the mode is supported for the given module"
         for module in self.tracked_modules.values():
+            # EMBEDDING mode is only supported for ShardedEmbeddingCollection
             assert not (
                 isinstance(module, ShardedEmbeddingBagCollection)
                 and self._mode == TrackingMode.EMBEDDING
             ), "EBC's lookup returns pooled embeddings and currently, we do not support tracking raw embeddings."
-            # register post lookup function
-            # pyre-ignore[29]
-            module.register_post_lookup_tracker_fn(self.record_lookup)
+
+            if (
+                self._mode == TrackingMode.ID_ONLY
+                or self._mode == TrackingMode.EMBEDDING
+            ):
+                # register post lookup function
+                # pyre-ignore[29]
+                module.register_post_lookup_tracker_fn(self.record_lookup)
+            elif self._mode == TrackingMode.MOMENTUM_LAST:
+                # pyre-ignore[29]:
+                for lookup in module._lookups:
+                    assert isinstance(
+                        lookup,
+                        (GroupedEmbeddingsLookup, GroupedPooledEmbeddingsLookup),
+                    )
+                    lookup.register_optim_state_tracker_fn(self.record_lookup)
+            else:
+                raise NotImplementedError(
+                    f"Tracking mode {self._mode} is not supported"
+                )
             # register auto compaction function at odist
             if self._auto_compact:
                 # pyre-ignore[29]