Enable Partial Rowwise Adam in SSD TBE (#4525)

Summary:
Pull Request resolved: #4525

X-link: facebookresearch/FBGEMM#1572

This diff enables Partial Rowwise Adam in SSD TBE, though just for the non KV-ZCH use case.

Reviewed By: emlin

Differential Revision: D77880362

Author: q10

fbgemm_gpu/fbgemm_gpu/tbe/ssd/training.py

@@ -180,6 +180,7 @@ def __init__(
         # Set the optimizer
         assert optimizer in (
             OptimType.EXACT_ROWWISE_ADAGRAD,
+            OptimType.PARTIAL_ROWWISE_ADAM,
         ), f"Optimizer {optimizer} is not supported by SSDTableBatchedEmbeddingBags"
         self.optimizer = optimizer
 
@@ -2205,7 +2206,7 @@ def forward(
         self.step += 1
 
         # Increment the iteration (value is used for certain optimizers)
-        self._increment_iteration()
+        iter_int = self._increment_iteration()
 
         if self.optimizer == OptimType.EXACT_SGD:
             raise AssertionError(
@@ -2226,6 +2227,24 @@ def forward(
                 common_args, self.optimizer_args, momentum1
             )
 
+        elif self.optimizer == OptimType.PARTIAL_ROWWISE_ADAM:
+            momentum2 = invokers.lookup_args_ssd.Momentum(
+                # pyre-ignore[6]
+                dev=self.momentum2_dev,
+                # pyre-ignore[6]
+                host=self.momentum2_host,
+                # pyre-ignore[6]
+                uvm=self.momentum2_uvm,
+                # pyre-ignore[6]
+                offsets=self.momentum2_offsets,
+                # pyre-ignore[6]
+                placements=self.momentum2_placements,
+            )
+
+            return invokers.lookup_partial_rowwise_adam_ssd.invoke(
+                common_args, self.optimizer_args, momentum1, momentum2, iter_int
+            )
+
     @torch.jit.ignore
     def _split_optimizer_states_non_kv_zch(
         self,

fbgemm_gpu/test/tbe/ssd/ssd_split_tbe_training_test.py

@@ -473,6 +473,9 @@ def generate_ssd_tbes(
         weighted: bool,
         lr: float = 0.01,  # from SSDTableBatchedEmbeddingBags
         eps: float = 1.0e-8,  # from SSDTableBatchedEmbeddingBags
+        weight_decay: float = 0.0,  # used by LARS-SGD, LAMB, ADAM, and Rowwise Adagrad
+        beta1: float = 0.9,  # used by Partial Rowwise Adam
+        beta2: float = 0.999,  # used by Partial Rowwise Adam
         ssd_shards: int = 1,  # from SSDTableBatchedEmbeddingBags
         optimizer: OptimType = OptimType.EXACT_ROWWISE_ADAGRAD,
         cache_set_scale: float = 1.0,
@@ -487,6 +490,7 @@ def generate_ssd_tbes(
         lazy_bulk_init_enabled: bool = False,
         backend_type: BackendType = BackendType.SSD,
         enable_raw_embedding_streaming: bool = False,
+        optimizer_state_dtypes: Dict[str, SparseType] = {},  # noqa: B006
     ) -> Tuple[SSDTableBatchedEmbeddingBags, List[torch.nn.EmbeddingBag]]:
         """
         Generate embedding modules (i,e., SSDTableBatchedEmbeddingBags and
@@ -561,6 +565,9 @@ def generate_ssd_tbes(
                 ssd_uniform_init_upper=0.1,
                 learning_rate=lr,
                 eps=eps,
+                weight_decay=weight_decay,
+                beta1=beta1,
+                beta2=beta2,
                 ssd_rocksdb_shards=ssd_shards,
                 optimizer=optimizer,
                 pooling_mode=pooling_mode,
@@ -796,6 +803,16 @@ def split_optimizer_states_(
             bucket_asc_ids_list, no_snapshot=False, should_flush=True
         )
 
+    def assert_close_(self, test: torch.Tensor, ref: torch.Tensor) -> None:
+        tolerance = 1.0e-4 if test.dtype == torch.float else 1.0e-2
+
+        torch.testing.assert_close(
+            test.float().cpu(),
+            ref.float().cpu(),
+            atol=tolerance,
+            rtol=tolerance,
+        )
+
     @given(
         **default_st, backend_type=st.sampled_from([BackendType.SSD, BackendType.DRAM])
     )
@@ -1014,6 +1031,165 @@ def test_ssd_backward_adagrad(
                 rtol=tolerance,
             )
 
+    @given(
+        **default_st,
+        backend_type=st.sampled_from([BackendType.SSD, BackendType.DRAM]),
+    )
+    @settings(verbosity=Verbosity.verbose, max_examples=MAX_EXAMPLES, deadline=None)
+    def test_ssd_backward_partial_rowwise_adam(
+        self,
+        T: int,
+        D: int,
+        B: int,
+        log_E: int,
+        L: int,
+        weighted: bool,
+        cache_set_scale: float,
+        pooling_mode: PoolingMode,
+        weights_precision: SparseType,
+        output_dtype: SparseType,
+        share_table: bool,
+        trigger_bounds_check: bool,
+        mixed_B: bool,
+        backend_type: BackendType,
+    ) -> None:
+        assume(not weighted or pooling_mode == PoolingMode.SUM)
+        # VBE is currently not supported for PARTIAL_ROWWISE_ADAM optimizer
+        assume(not mixed_B)
+
+        # Constants
+        lr = 0.5
+        eps = 0.2
+        ssd_shards = 2
+        beta1 = 0.9
+        beta2 = 0.99
+        weight_decay = 0.01
+
+        # Generate embedding modules and inputs
+        (
+            emb,
+            emb_ref,
+        ) = self.generate_ssd_tbes(
+            T,
+            D,
+            B,
+            log_E,
+            L,
+            weighted,
+            lr=lr,
+            eps=eps,
+            weight_decay=weight_decay,
+            beta1=beta1,
+            beta2=beta2,
+            ssd_shards=ssd_shards,
+            optimizer=OptimType.PARTIAL_ROWWISE_ADAM,
+            cache_set_scale=cache_set_scale,
+            pooling_mode=pooling_mode,
+            weights_precision=weights_precision,
+            output_dtype=output_dtype,
+            share_table=share_table,
+            backend_type=backend_type,
+        )
+
+        Es = [emb.embedding_specs[t][0] for t in range(T)]
+        (
+            indices_list,
+            per_sample_weights_list,
+            indices,
+            offsets,
+            per_sample_weights,
+            batch_size_per_feature_per_rank,
+        ) = self.generate_inputs_(
+            B,
+            L,
+            Es,
+            emb.feature_table_map,
+            weights_precision=weights_precision,
+            trigger_bounds_check=trigger_bounds_check,
+            # mixed_B=mixed_B,
+        )
+
+        # Execute forward
+        output_ref_list, output = self.execute_ssd_forward_(
+            emb,
+            emb_ref,
+            indices_list,
+            per_sample_weights_list,
+            indices,
+            offsets,
+            per_sample_weights,
+            B,
+            L,
+            weighted,
+            batch_size_per_feature_per_rank=batch_size_per_feature_per_rank,
+        )
+
+        # Execute backward
+        self.execute_ssd_backward_(
+            output_ref_list,
+            output,
+            B,
+            D,
+            pooling_mode,
+            batch_size_per_feature_per_rank,
+        )
+
+        emb.flush()
+
+        tolerance = 1.0e-2
+
+        emb_test_weights = emb.debug_split_embedding_weights()
+        split_optimizer_states = self.split_optimizer_states_(emb)
+
+        for _, t in self.get_physical_table_arg_indices_(emb.feature_table_map):
+            (m1, m2) = split_optimizer_states[t]
+            print(f"{t=} {m1.shape=}, {m2.shape=}")
+
+        for f, t in self.get_physical_table_arg_indices_(emb.feature_table_map):
+            (m1, m2) = split_optimizer_states[t]
+            # Some optimizers have non-float momentum values
+            # pyre-ignore[16]
+            ref_grad = emb_ref[f].weight.grad.cpu().to_dense()
+            ref_weights = emb_ref[f].weight.cpu()
+
+            # Compare momentum2 values: (1 - beta2) * dL^2
+            m2_ref = (ref_grad.pow(2).mean(dim=1)) * (1.0 - beta2)
+            self.assert_close_(m2, m2_ref)
+
+            # Compare momentum1 values: (1 - beta1) * dL
+            m1_ref = ref_grad * (1.0 - beta1)
+            print(f"{m1_ref.shape=}, {m1.shape=}")
+            self.assert_close_(m1, m1_ref)
+
+            # Bias corrections
+            iter_ = emb.iter.item()
+            v_hat_t = m2_ref / (1 - beta2**iter_)
+            v_hat_t = v_hat_t.view(v_hat_t.numel(), 1)
+            m_hat_t = m1_ref / (1 - beta1**iter_)
+
+            # Weight update
+            ref_weights_updated = (
+                torch.addcdiv(
+                    ref_weights,
+                    value=-lr,
+                    tensor1=m_hat_t,
+                    tensor2=v_hat_t.sqrt_().add_(eps),
+                )
+                - lr * weight_decay * ref_weights
+            )
+
+            if weights_precision == SparseType.FP16:
+                # Round the reference weight the same way that TBE does
+                ref_weights_updated = ref_weights_updated.half().float()
+
+            # Compare weights
+            torch.testing.assert_close(
+                emb_test_weights[t].float().cuda(),
+                ref_weights_updated.cuda(),
+                atol=tolerance,
+                rtol=tolerance,
+            )
+
     @given(
         bulk_init_chunk_size=st.sampled_from([0, 204800]),
         lazy_bulk_init_enabled=st.booleans(),