Adding support for MOMENTUM_DIFF and ROWWISE_ADAGRAD optimizer states (#3144)

Summary:
Pull Request resolved: #3144

This diff extends the Model Delta Tracker to support two new tracking modes: `MOMENTUM_DIFF` and `ROWWISE_ADAGRAD`, which enable tracking of rowwise optimizer states for more sophisticated gradient analysis.

Differential Revision: D76918891

Author: aliafzal

torchrec/distributed/model_tracker/model_delta_tracker.py

@@ -8,14 +8,20 @@
 # pyre-strict
 import logging as logger
 from collections import Counter, OrderedDict
-from typing import Dict, Iterable, List, Optional
+from typing import Dict, Iterable, List, Optional, Tuple
 
 import torch
+from fbgemm_gpu.split_embedding_configs import EmbOptimType as OptimType
+from fbgemm_gpu.split_table_batched_embeddings_ops import (
+    SplitTableBatchedEmbeddingBagsCodegen,
+)
 
 from torch import nn
+from torchrec.distributed.batched_embedding_kernel import BatchedFusedEmbedding
 
 from torchrec.distributed.embedding import ShardedEmbeddingCollection
 from torchrec.distributed.embedding_lookup import (
+    BatchedFusedEmbeddingBag,
     GroupedEmbeddingsLookup,
     GroupedPooledEmbeddingsLookup,
 )
@@ -26,6 +32,8 @@
     EmbdUpdateMode,
     TrackingMode,
 )
+from torchrec.distributed.utils import none_throws
+
 from torchrec.sparse.jagged_tensor import KeyedJaggedTensor
 
 UPDATE_MODE_MAP: Dict[TrackingMode, EmbdUpdateMode] = {
@@ -42,6 +50,14 @@
     # This mode supports approximate top-k delta-row selection, can be
     # obtained by running momentum.norm().topk().
     TrackingMode.MOMENTUM_LAST: EmbdUpdateMode.LAST,
+    # MOMENTUM_DIFF keeps a running sum of the square of the gradients per row.
+    # Within each publishing interval, we track the starting value of this running
+    # sum on all used rows and then do a lookup when ``get_delta`` is called to query
+    # the latest sum. Then we can compute the delta of the two values and return them
+    # together with the row ids.
+    TrackingMode.MOMENTUM_DIFF: EmbdUpdateMode.FIRST,
+    # The same as MOMENTUM_DIFF. Adding for backward compatibility.
+    TrackingMode.ROWWISE_ADAGRAD: EmbdUpdateMode.FIRST,
 }
 
 # Tracking is current only supported for ShardedEmbeddingCollection and ShardedEmbeddingBagCollection.
@@ -99,6 +115,7 @@ def __init__(
 
         # from module FQN to ShardedEmbeddingCollection/ShardedEmbeddingBagCollection
         self.tracked_modules: Dict[str, nn.Module] = {}
+        self.table_to_fqn: Dict[str, str] = {}
         self.feature_to_fqn: Dict[str, str] = {}
         # Generate the mapping from FQN to feature names.
         self.fqn_to_feature_names()
@@ -180,6 +197,11 @@ def record_lookup(
         # 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)
+        elif (
+            self._mode == TrackingMode.MOMENTUM_DIFF
+            or self._mode == TrackingMode.ROWWISE_ADAGRAD
+        ):
+            self.record_rowwise_optim_state(emb_module, kjt)
         else:
             raise NotImplementedError(f"Tracking mode {self._mode} is not supported")
 
@@ -278,6 +300,60 @@ def record_momentum(
                 states=per_key_states,
             )
 
+    def record_rowwise_optim_state(
+        self,
+        emb_module: nn.Module,
+        kjt: KeyedJaggedTensor,
+    ) -> None:
+        opt_states: List[List[torch.Tensor]] = (
+            # pyre-fixme[16]: Item `Tensor` of `Tensor | Module` has no attribute
+            #  `split_optimizer_states`.
+            emb_module._emb_module.split_optimizer_states()
+        )
+        proxy: torch.Tensor = torch.cat([state[0] for state in opt_states])
+        states = proxy[kjt.values()]
+        assert (
+            kjt.values().numel() == states.numel()
+        ), f"number of ids and states mismatch, expect {kjt.values()=}, {kjt.values().numel()}, but got {states.numel()} "
+        offsets: torch.Tensor = torch.ops.fbgemm.asynchronous_complete_cumsum(
+            torch.tensor(kjt.length_per_key(), dtype=torch.int64)
+        )
+        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_latest(self) -> Dict[str, torch.Tensor]:
+        ret: Dict[str, torch.Tensor] = {}
+        for module in self.tracked_modules.values():
+            # pyre-fixme[29]:
+            for lookup in module._lookups:
+                for embs_module in lookup._emb_modules:
+                    assert isinstance(
+                        embs_module, (BatchedFusedEmbeddingBag, BatchedFusedEmbedding)
+                    ), f"expect BatchedFusedEmbeddingBag or BatchedFusedEmbedding, but {type(embs_module)} found"
+                    tbe = embs_module._emb_module
+
+                    assert isinstance(tbe, SplitTableBatchedEmbeddingBagsCodegen)
+                    table_names = [t.name for t in embs_module._config.embedding_tables]
+                    opt_states = tbe.split_optimizer_states()
+                    assert len(table_names) == len(opt_states)
+
+                    for i, table_name in enumerate(table_names):
+                        emb_fqn = self.table_to_fqn[table_name]
+                        table_state = opt_states[i][0]
+                        assert (
+                            emb_fqn not in ret
+                        ), f"a table with {emb_fqn} already exists"
+                        ret[emb_fqn] = table_state
+
+        return ret
+
     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
@@ -289,7 +365,13 @@ def get_delta_ids(self, consumer: Optional[str] = None) -> Dict[str, torch.Tenso
         per_table_delta_rows = self.get_delta(consumer)
         return {fqn: delta_rows.ids for fqn, delta_rows in per_table_delta_rows.items()}
 
-    def get_delta(self, consumer: Optional[str] = None) -> Dict[str, DeltaRows]:
+    def get_delta(
+        self,
+        consumer: Optional[str] = None,
+        top_percentage: Optional[float] = 1.0,
+        per_table_percentage: Optional[Dict[str, Tuple[float, str]]] = None,
+        sorted_by_indices: Optional[bool] = True,
+    ) -> Dict[str, DeltaRows]:
         """
         Return a dictionary of hit local IDs and parameter states / embeddings for each sparse feature. The Values are first keyed by submodule FQN.
 
@@ -314,6 +396,65 @@ def get_delta(self, consumer: Optional[str] = None) -> Dict[str, DeltaRows]:
         self.per_consumer_batch_idx[consumer] = index_end
         if self._delete_on_read:
             self.store.delete(up_to_idx=min(self.per_consumer_batch_idx.values()))
+
+        if self._mode in (TrackingMode.MOMENTUM_DIFF, TrackingMode.ROWWISE_ADAGRAD):
+            square_sum_map = self.get_latest()
+            for fqn, rows in tracker_rows.items():
+                assert (
+                    fqn in square_sum_map
+                ), f"{fqn} not found in {square_sum_map.keys()}"
+                # compute delta sum_t(g^2) for t in [t1, t2] through
+                # sum_t2(g^2) - sum_t1(g^2)
+                # pyre-fixme[58]: `-` is not supported for operand types `Tensor`
+                #  and `Optional[Tensor]`.
+                rows.states = square_sum_map[fqn][rows.ids] - rows.states
+
+                if rows.states is not None:
+                    default_k = rows.states.size(-1)
+                    top_k = (
+                        int(top_percentage * default_k)
+                        if top_percentage is not None
+                        else default_k
+                    )
+
+                    if (
+                        per_table_percentage is not None
+                        and per_table_percentage.get(fqn) is not None
+                    ):
+                        per_table_k = int(per_table_percentage[fqn][0] * default_k)
+                        policy = per_table_percentage[fqn][1]
+
+                        if policy == "MIN":
+                            top_k = min(top_k, per_table_k)
+                        elif policy == "MAX":
+                            top_k = max(top_k, per_table_k)
+                        elif policy == "OVERRIDE":
+                            top_k = per_table_k
+                        else:
+                            logger.warning(
+                                f"Unknown policy {policy}, will keep using original top_k {top_k}"
+                            )
+
+                    logger.info(f"get_unique {fqn=} {top_k=} {default_k=}")
+
+                    if top_k >= default_k:
+                        continue
+
+                    if sorted_by_indices:
+                        sorted_indices, _ = torch.sort(
+                            torch.topk(
+                                none_throws(rows.states), top_k, sorted=False
+                            ).indices,
+                            stable=False,
+                        )
+                        rows.ids = rows.ids[sorted_indices]
+                        rows.states = none_throws(rows.states)[sorted_indices]
+                    else:
+                        rows.states, indices = torch.topk(
+                            none_throws(rows.states), top_k, sorted=False
+                        )
+                        rows.ids = rows.ids[indices]
+
         return tracker_rows
 
     def get_tracked_modules(self) -> Dict[str, nn.Module]:
@@ -330,7 +471,6 @@ def fqn_to_feature_names(self) -> Dict[str, List[str]]:
             return self._fqn_to_feature_map
 
         table_to_feature_names: Dict[str, List[str]] = OrderedDict()
-        table_to_fqn: Dict[str, str] = OrderedDict()
         for fqn, named_module in self._model.named_modules():
             split_fqn = fqn.split(".")
             # Skipping partial FQNs present in fqns_to_skip
@@ -356,13 +496,13 @@ def fqn_to_feature_names(self) -> Dict[str, List[str]]:
                 # will incorrectly match fqn with all the table names that have the same prefix
                 if table_name in split_fqn:
                     embedding_fqn = self._clean_fqn_fn(fqn)
-                    if table_name in table_to_fqn:
+                    if table_name in self.table_to_fqn:
                         # Sanity check for validating that we don't have more then one table mapping to same fqn.
                         logger.warning(
-                            f"Override {table_to_fqn[table_name]} with {embedding_fqn} for entry {table_name}"
+                            f"Override {self.table_to_fqn[table_name]} with {embedding_fqn} for entry {table_name}"
                         )
-                    table_to_fqn[table_name] = embedding_fqn
-            logger.info(f"Table to fqn: {table_to_fqn}")
+                    self.table_to_fqn[table_name] = embedding_fqn
+            logger.info(f"Table to fqn: {self.table_to_fqn}")
         flatten_names = [
             name for names in table_to_feature_names.values() for name in names
         ]
@@ -375,15 +515,15 @@ def fqn_to_feature_names(self) -> Dict[str, List[str]]:
 
         fqn_to_feature_names: Dict[str, List[str]] = OrderedDict()
         for table_name in table_to_feature_names:
-            if table_name not in table_to_fqn:
+            if table_name not in self.table_to_fqn:
                 # This is likely unexpected, where we can't locate the FQN associated with this table.
                 logger.warning(
-                    f"Table {table_name} not found in {table_to_fqn}, skipping"
+                    f"Table {table_name} not found in {self.table_to_fqn}, skipping"
                 )
                 continue
-            fqn_to_feature_names[table_to_fqn[table_name]] = table_to_feature_names[
-                table_name
-            ]
+            fqn_to_feature_names[self.table_to_fqn[table_name]] = (
+                table_to_feature_names[table_name]
+            )
         self._fqn_to_feature_map = fqn_to_feature_names
         return fqn_to_feature_names
 
@@ -451,6 +591,24 @@ def _validate_and_init_tracker_fns(self) -> None:
                         (GroupedEmbeddingsLookup, GroupedPooledEmbeddingsLookup),
                     )
                     lookup.register_optim_state_tracker_fn(self.record_lookup)
+            elif (
+                self._mode == TrackingMode.ROWWISE_ADAGRAD
+                or self._mode == TrackingMode.MOMENTUM_DIFF
+            ):
+                # pyre-ignore[29]:
+                for lookup in module._lookups:
+                    assert isinstance(
+                        lookup,
+                        (GroupedEmbeddingsLookup, GroupedPooledEmbeddingsLookup),
+                    ) and all(
+                        # TorchRec maps ROWWISE_ADAGRAD to EXACT_ROWWISE_ADAGRAD
+                        # pyre-ignore[16]:
+                        emb._emb_module.optimizer == OptimType.EXACT_ROWWISE_ADAGRAD
+                        # pyre-ignore[16]:
+                        or emb._emb_module.optimizer == OptimType.PARTIAL_ROWWISE_ADAM
+                        for emb in lookup._emb_modules
+                    )
+                    lookup.register_optim_state_tracker_fn(self.record_lookup)
             else:
                 raise NotImplementedError(
                     f"Tracking mode {self._mode} is not supported"