ensemble_mode consolidation (fbgemm)

Summary: X-link: facebookresearch/FBGEMM#295 ensemble_mode consolidation (fbgemm) Differential Revision: D63634421
pytorch · Sep 30, 2024 · 3f65f9f · 3f65f9f
1 parent 93dcc07
commit 3f65f9f
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 14 deletions.
diff --git a/fbgemm_gpu/fbgemm_gpu/split_table_batched_embeddings_ops_training.py b/fbgemm_gpu/fbgemm_gpu/split_table_batched_embeddings_ops_training.py
@@ -150,6 +150,7 @@ class EnsembleModeDefinition:
     step_ema: float = 10000
     step_swap: float = 10000
     step_start: float = 0
+    step_ema_coef: float = 0.6
     step_mode: StepMode = StepMode.USE_ITER
 
 
@@ -924,10 +925,11 @@ def __init__(  # noqa C901
 
         if ensemble_mode is None:
             ensemble_mode = EnsembleModeDefinition()
-        self._ensemble_mode: Dict[str, int] = {
+        self._ensemble_mode: Dict[str, Union[int, float]] = {
             "step_ema": int(ensemble_mode.step_ema),
             "step_swap": int(ensemble_mode.step_swap),
             "step_start": int(ensemble_mode.step_start),
+            "step_ema_coef": float(ensemble_mode.step_ema_coef),
             "step_mode": int(ensemble_mode.step_mode.value),
         }
 
@@ -1002,6 +1004,7 @@ def __init__(  # noqa C901
                     if (
                         optimizer_state_dtypes is None
                         or "momentum1" not in optimizer_state_dtypes
+                        or optimizer == OptimType.ENSEMBLE_ROWWISE_ADAGRAD
                     )
                     else optimizer_state_dtypes["momentum1"].as_dtype()
                 )
@@ -1938,7 +1941,7 @@ def forward(  # noqa: C901
 
         raise ValueError(f"Invalid OptimType: {self.optimizer}")
 
-    def ensemble_and_swap(self, ensemble_mode: Dict[str, int]) -> None:
+    def ensemble_and_swap(self, ensemble_mode: Dict[str, Union[int, float]]) -> None:
         should_ema = self.iter.item() % int(ensemble_mode["step_ema"]) == 0
         should_swap = self.iter.item() % int(ensemble_mode["step_swap"]) == 0
         if should_ema or should_swap:
@@ -1947,7 +1950,7 @@ def ensemble_and_swap(self, ensemble_mode: Dict[str, int]) -> None:
             for i in range(len(self.embedding_specs)):
                 if should_ema:
                     coef_ema = (
-                        self.optimizer_args.momentum
+                        ensemble_mode["step_ema_coef"]
                         if self.iter.item() > int(ensemble_mode["step_start"])
                         else 0.0
                     )

diff --git a/fbgemm_gpu/test/tbe/training/backward_optimizers_test.py b/fbgemm_gpu/test/tbe/training/backward_optimizers_test.py
@@ -26,6 +26,7 @@
     CounterBasedRegularizationDefinition,
     CounterWeightDecayMode,
     CowClipDefinition,
+    EnsembleModeDefinition,
     GradSumDecay,
     LearningRateMode,
     SplitTableBatchedEmbeddingBagsCodegen,
@@ -34,13 +35,6 @@
     WeightDecayMode,
 )
 
-try:
-    from fbgemm_gpu.split_table_batched_embeddings_ops_training import (
-        EnsembleModeDefinition,
-    )
-except ImportError:
-    EnsembleModeDefinition = None
-
 from fbgemm_gpu.tbe.utils import (
     b_indices,
     get_table_batched_offsets_from_dense,
@@ -315,23 +309,22 @@ def execute_backward_optimizers_(  # noqa C901
             optimizer_kwargs["eta"] = eta
 
         if optimizer == OptimType.ENSEMBLE_ROWWISE_ADAGRAD:
-            (eps, step_ema, step_swap, step_start, step_mode, momentum) = (
+            (eps, step_ema, step_swap, step_start, step_mode) = (
                 1e-4,
                 1.0,
                 1.0,
                 0.0,
                 StepMode.USE_ITER,
-                0.8,
             )
             optimizer_kwargs["eps"] = eps
+            optimizer_kwargs["optimizer_state_dtypes"] = optimizer_state_dtypes
             optimizer_kwargs["ensemble_mode"] = EnsembleModeDefinition(
                 step_ema=step_ema,
                 step_swap=step_swap,
                 step_start=step_start,
+                step_ema_coef=momentum,
                 step_mode=step_mode,
             )
-            optimizer_kwargs["momentum"] = momentum
-            optimizer_kwargs["optimizer_state_dtypes"] = optimizer_state_dtypes
 
         cc = emb_op(
             embedding_specs=[