Renaming embeddings to states to account for optimizers states. (#3116)

Summary:
Pull Request resolved: #3116

We have named the field 'embeddings' in the IndexedLookup and DeltaRows data classes for preserving embeddings and model tracker output. Now that we want to expand this functionality to preserve optimizer state as well, renaming the 'embeddings' field to 'states'.

Reviewed By: TroyGarden

Differential Revision: D76867584

fbshipit-source-id: 2659c267714940a33da1ba8371e37345c601faa9

Author: aliafzal

torchrec/distributed/model_tracker/delta_store.py

@@ -21,29 +21,27 @@
 
 def _compute_unique_rows(
     ids: List[torch.Tensor],
-    embeddings: Optional[List[torch.Tensor]],
+    states: Optional[List[torch.Tensor]],
     mode: EmbdUpdateMode,
 ) -> DeltaRows:
     r"""
     To calculate unique ids and embeddings
     """
     if mode == EmbdUpdateMode.NONE:
-        assert (
-            embeddings is None
-        ), f"{mode=} == EmbdUpdateMode.NONE but received embeddings"
+        assert states is None, f"{mode=} == EmbdUpdateMode.NONE but received embeddings"
         unique_ids = torch.cat(ids).unique(return_inverse=False)
-        return DeltaRows(ids=unique_ids, embeddings=None)
+        return DeltaRows(ids=unique_ids, states=None)
     else:
         assert (
-            embeddings is not None
+            states is not None
         ), f"{mode=} != EmbdUpdateMode.NONE but received no embeddings"
 
         cat_ids = torch.cat(ids)
-        cat_embeddings = torch.cat(embeddings)
+        cat_states = torch.cat(states)
 
         if mode == EmbdUpdateMode.LAST:
             cat_ids = cat_ids.flip(dims=[0])
-            cat_embeddings = cat_embeddings.flip(dims=[0])
+            cat_states = cat_states.flip(dims=[0])
 
         # Get unique ids and inverse mapping (each element's index in unique_ids).
         unique_ids, inverse = cat_ids.unique(sorted=False, return_inverse=True)
@@ -65,8 +63,8 @@ def _compute_unique_rows(
         )
 
         # Use first occurrence indices to select corresponding embedding row.
-        unique_embedings = cat_embeddings[first_occurrence]
-        return DeltaRows(ids=unique_ids, embeddings=unique_embedings)
+        unique_states = cat_states[first_occurrence]
+        return DeltaRows(ids=unique_ids, states=unique_states)
 
 
 class DeltaStore:
@@ -90,11 +88,11 @@ def append(
         batch_idx: int,
         table_fqn: str,
         ids: torch.Tensor,
-        embeddings: Optional[torch.Tensor],
+        states: Optional[torch.Tensor],
     ) -> None:
         table_fqn_lookup = self.per_fqn_lookups.get(table_fqn, [])
         table_fqn_lookup.append(
-            IndexedLookup(batch_idx=batch_idx, ids=ids, embeddings=embeddings)
+            IndexedLookup(batch_idx=batch_idx, ids=ids, states=states)
         )
         self.per_fqn_lookups[table_fqn] = table_fqn_lookup
 
@@ -132,21 +130,21 @@ def compact(self, start_idx: int, end_idx: int) -> None:
                 new_per_fqn_lookups[table_fqn] = lookups
                 continue
             ids = [lookup.ids for lookup in lookups_to_compact]
-            embeddings = (
-                [none_throws(lookup.embeddings) for lookup in lookups_to_compact]
+            states = (
+                [none_throws(lookup.states) for lookup in lookups_to_compact]
                 if self.embdUpdateMode != EmbdUpdateMode.NONE
                 else None
             )
             delta_rows = _compute_unique_rows(
-                ids=ids, embeddings=embeddings, mode=self.embdUpdateMode
+                ids=ids, states=states, mode=self.embdUpdateMode
             )
             new_per_fqn_lookups[table_fqn] = (
                 lookups[:index_l]
                 + [
                     IndexedLookup(
                         batch_idx=start_idx,
                         ids=delta_rows.ids,
-                        embeddings=delta_rows.embeddings,
+                        states=delta_rows.states,
                     )
                 ]
                 + lookups[index_r:]
@@ -163,9 +161,9 @@ def get_delta(self, from_idx: int = 0) -> Dict[str, DeltaRows]:
             compact_ids = [
                 lookup.ids for lookup in lookups if lookup.batch_idx >= from_idx
             ]
-            compact_embeddings = (
+            compact_states = (
                 [
-                    none_throws(lookup.embeddings)
+                    none_throws(lookup.states)
                     for lookup in lookups
                     if lookup.batch_idx >= from_idx
                 ]
@@ -174,6 +172,6 @@ def get_delta(self, from_idx: int = 0) -> Dict[str, DeltaRows]:
             )
 
             delta_per_table_fqn[table_fqn] = _compute_unique_rows(
-                ids=compact_ids, embeddings=compact_embeddings, mode=self.embdUpdateMode
+                ids=compact_ids, states=compact_states, mode=self.embdUpdateMode
             )
         return delta_per_table_fqn

torchrec/distributed/model_tracker/model_delta_tracker.py

@@ -183,7 +183,7 @@ def record_ids(self, kjt: KeyedJaggedTensor) -> None:
                 batch_idx=self.curr_batch_idx,
                 table_fqn=table_fqn,
                 ids=torch.cat(ids_list),
-                embeddings=None,
+                states=None,
             )
 
     def record_embeddings(
@@ -223,7 +223,7 @@ def record_embeddings(
                 batch_idx=self.curr_batch_idx,
                 table_fqn=table_fqn,
                 ids=torch.cat(ids_list),
-                embeddings=torch.cat(per_table_emb[table_fqn]),
+                states=torch.cat(per_table_emb[table_fqn]),
             )
 
     def get_delta_ids(self, consumer: Optional[str] = None) -> Dict[str, torch.Tensor]: