Add configuration knob for ENSEMBLE_ROWWISE_ADAGRAD, frontend (#2955)

Summary:
Pull Request resolved: #2955

X-link: facebookresearch/FBGEMM#55

- Add configuration knob for ENSEMBLE_ROWWISE_ADAGRAD, frontend

Reviewed By: spcyppt

Differential Revision: D60986449

fbgemm_gpu/fbgemm_gpu/split_embedding_configs.py

@@ -33,6 +33,7 @@ class EmbOptimType(enum.Enum):
     SHAMPOO_V2 = "shampoo_v2"  # not currently supported for sparse embedding tables
     MADGRAD = "madgrad"
     EXACT_ROWWISE_WEIGHTED_ADAGRAD = "exact_row_wise_weighted_adagrad"  # deprecated
+    ENSEMBLE_ROWWISE_ADAGRAD = "ensemble_row_wise_adagrad"
     NONE = "none"
 
     def __str__(self) -> str:

fbgemm_gpu/fbgemm_gpu/split_table_batched_embeddings_ops_training.py

@@ -25,7 +25,7 @@
 
 import fbgemm_gpu.split_embedding_codegen_lookup_invokers as invokers
 
-# from fbgemm_gpu.config import FeatureGateName
+from fbgemm_gpu.config import FeatureGate, FeatureGateName
 from fbgemm_gpu.runtime_monitor import (
     AsyncSeriesTimer,
     TBEStatsReporter,
@@ -1331,6 +1331,12 @@ def _generate_vbe_metadata(
             self.current_device,
         )
 
+    @torch.jit.ignore
+    def _feature_is_enabled(self, feature: FeatureGateName) -> bool:
+        # Define proxy method so that it can be marked with @torch.jit.ignore
+        # This allows models using this class to compile correctly
+        return FeatureGate.is_enabled(feature)
+
     def forward(  # noqa: C901
         self,
         indices: Tensor,
@@ -1549,6 +1555,17 @@ def forward(  # noqa: C901
             offsets=self.row_counter_offsets,
             placements=self.row_counter_placements,
         )
+
+        if self.optimizer == OptimType.ENSEMBLE_ROWWISE_ADAGRAD:
+            if self._feature_is_enabled(FeatureGateName.TBE_ENSEMBLE_ROWWISE_ADAGRAD):
+                raise AssertionError(
+                    "ENSEMBLE_ROWWISE_ADAGRAD feature has not landed yet (see D60189486 stack)"
+                )
+            else:
+                raise AssertionError(
+                    "ENSEMBLE_ROWWISE_ADAGRAD is an inactive or deprecated feature!"
+                )
+
         if self._used_rowwise_adagrad_with_counter:
             if (
                 self._max_counter_update_freq > 0