Update the rowwise adagrad optimizer to leverage optimizer state offloading, v4, frontend (#4249)

Summary:
X-link: facebookresearch/FBGEMM#1328


- Follow up to D75329024by plumbing the flag for optimizer state offloading to the TBE frontend

Differential Revision: D75336208

Author: q10

fbgemm_gpu/codegen/genscript/generate_backward_split.py

@@ -10,6 +10,7 @@
 
 import itertools
 import sys
+from copy import deepcopy
 from typing import List
 
 try:
@@ -164,6 +165,10 @@ def generate_backward_split_gpu(**kwargs: Any) -> None:
         if not kwargs.get("dense"):
             # Generate CUDA autograd
 
+            # Extract the aux_args and ssd_aux_args for later use
+            aux_args = kwargs["aux_args"]
+            ssd_aux_args = kwargs["ssd_aux_args"]
+
             for ssd in [True, False] if kwargs.get("has_ssd_support") else [False]:
                 template_filepath = (
                     "training/backward/embedding_backward_split_host_template.cpp"
@@ -195,6 +200,10 @@ def generate_backward_split_gpu(**kwargs: Any) -> None:
                     )
 
                 if kwargs.get("has_cpu_support") or kwargs.get("has_gpu_support"):
+                    # Since the template file only uses aux_args, reset the key
+                    # based on whether we are generated for SSD variant or not
+                    kwargs["aux_args"] = ssd_aux_args if ssd else aux_args
+
                     # Generates Python invoker for CUDA + CPU, and PT2
                     template = CodeTemplate.load(
                         "training/python/split_embedding_codegen_lookup_invoker.template"
@@ -433,28 +442,44 @@ def generate() -> None:
                 "mixed_D",  # 6
             ],
         }
-        # ssd-specific argument
+
+        # SSD-specific arguments
         ssd_aux_bool = [
+            # When set to true, the per-row optimizer state will offloaded to
+            # the end of each row in the SSD cache.
             "enable_optimizer_offloading",  # 7
         ]
+
         assert (
             list(aux_args.keys()) == aux_names
         ), f"{aux_names} must match {aux_args.keys()}"
 
+        ssd_aux_args = deepcopy(aux_args)
+        ssd_aux_args["aux_bool"].extend(ssd_aux_bool)
+
         all_optimizers = []
         ssd_optimizers = []
 
         for optimizer in optimizers:
             optim = optimizer["optimizer"]
+
             if (
                 optimizer["has_cpu_support"] or optimizer["has_gpu_support"]
             ) and optim != "dense":
                 all_optimizers.append(optim)
                 if optimizer["has_ssd_support"]:
                     ssd_optimizers.append(optim)
+
             BackwardSplitGenerator.generate_backward_split(
-                ssd_tensors=ssd_tensors, aux_args=aux_args, **optimizer
+                ssd_tensors=ssd_tensors,
+                # Both aux_args and ssd_aux_args will be passed in, since
+                # generate_backward_split will generate both SSD and non-SSD
+                # variants
+                aux_args=aux_args,
+                ssd_aux_args=ssd_aux_args,
+                **optimizer,
             )
+
         BackwardSplitGenerator.generate_rocm_backward_split()
 
         # Generate common device kernels for backwards
@@ -465,11 +490,10 @@ def generate() -> None:
         BackwardSplitGenerator.generate_backward_indices()
 
         # Generate headers for backwards
-        BackwardSplitGenerator.generate_backward_header(aux_args, aux_names)
-        aux_args["aux_bool"].extend(ssd_aux_bool)
-        BackwardSplitGenerator.generate_backward_header(
-            aux_args, aux_names, is_ssd=True
-        )
+        for is_ssd in [True, False]:
+            BackwardSplitGenerator.generate_backward_header(
+                (ssd_aux_args if is_ssd else aux_args), aux_names, is_ssd=is_ssd
+            )
 
         BackwardSplitGenerator.generate_python_sources(all_optimizers, ssd_optimizers)
 

fbgemm_gpu/codegen/training/python/lookup_args.template

@@ -48,6 +48,7 @@ class CommonArgs(NamedTuple):
     use_homogeneous_placements: bool
     {%- if ssd %}
     ssd_tensors: Dict[str, torch.Tensor]
+    enable_optimizer_offloading: bool
     {%- endif %}
     learning_rate_tensor: torch.Tensor
     info_B_num_bits: int

fbgemm_gpu/codegen/training/python/split_embedding_codegen_lookup_invoker.template

@@ -92,6 +92,7 @@ torch.ops.load_library("//deeplearning/fbgemm/fbgemm_gpu:sparse_ops_cpu")
         "Please check the frontend and backend version. "
     )
     {{ arg_type }}.append(dict_{{ arg_type }}["{{ var }}"])
+    
     {%- endfor %}
 {%- endmacro %}
 
@@ -203,12 +204,9 @@ def invoke(
         "use_uniq_cache_locations_bwd": common_args.use_uniq_cache_locations_bwd,
         "use_homogeneous_placements": common_args.use_homogeneous_placements,
         "apply_global_weight_decay": apply_global_weight_decay,
-        {%- if not ssd %}
-        "mixed_D": mixed_D
-        {%- else %}
         "mixed_D": mixed_D,
-        # TODO: Update this when frontend is ready to land
-        "enable_optimizer_offloading": False
+        {%- if ssd %}
+        "enable_optimizer_offloading": common_args.enable_optimizer_offloading,
         {%- endif %}
     }
     dict_optim_int: Dict[str, int] = {}

fbgemm_gpu/fbgemm_gpu/tbe/ssd/training.py

@@ -1828,6 +1828,7 @@ def forward(
                 "post_bwd_evicted_indices": post_bwd_evicted_indices_cpu,
                 "actions_count": actions_count_cpu,
             },
+            enable_optimizer_offloading=self.enable_optimizer_offloading,
             # pyre-fixme[6]: Expected `lookup_args_ssd.VBEMetadata` but got `lookup_args.VBEMetadata`
             vbe_metadata=vbe_metadata,
             learning_rate_tensor=self.learning_rate_tensor,

fbgemm_gpu/test/tbe/ssd/ssd_split_tbe_training_test.py

@@ -1595,9 +1595,7 @@ def test_kv_db_forward(
     @given(
         **default_st,
         num_buckets=st.integers(min_value=10, max_value=15),
-        opt_offloading=st.just(
-            False
-        ),  # make it st.booleans when Benson's opt offloading diff is landed
+        enable_optimizer_offloading=st.booleans(),
         backend_type=st.sampled_from([BackendType.SSD, BackendType.DRAM]),
     )
     @settings(verbosity=Verbosity.verbose, max_examples=MAX_EXAMPLES, deadline=None)
@@ -1617,7 +1615,7 @@ def test_kv_emb_state_dict(
         trigger_bounds_check: bool,
         mixed_B: bool,
         num_buckets: int,
-        opt_offloading: bool,
+        enable_optimizer_offloading: bool,
         backend_type: BackendType,
     ) -> None:
         # Constants
@@ -1653,7 +1651,7 @@ def test_kv_emb_state_dict(
             output_dtype=output_dtype,
             share_table=share_table,
             num_buckets=num_buckets,
-            enable_optimizer_offloading=opt_offloading,
+            enable_optimizer_offloading=enable_optimizer_offloading,
             backend_type=backend_type,
         )
 
@@ -1791,8 +1789,6 @@ def test_kv_emb_state_dict(
             self.assertLess(table_index, len(emb_state_dict_list))
             assert len(split_optimizer_states[table_index]) == num_ids
             opt = split_optimizer_states[table_index]
-            if opt_offloading:
-                opt = opt[bucket_asc_ids_list[table_index].view(-1)]
             new_ref_weight = torch.addcdiv(
                 emb_r_w.float(),
                 value=-lr,
@@ -1822,6 +1818,7 @@ def test_kv_emb_state_dict(
     @given(
         **default_st,
         num_buckets=st.integers(min_value=10, max_value=15),
+        enable_optimizer_offloading=st.booleans(),
     )
     @settings(verbosity=Verbosity.verbose, max_examples=MAX_EXAMPLES, deadline=None)
     def test_kv_opt_state_w_offloading(
@@ -1840,6 +1837,7 @@ def test_kv_opt_state_w_offloading(
         trigger_bounds_check: bool,
         mixed_B: bool,
         num_buckets: int,
+        enable_optimizer_offloading: bool,
     ) -> None:
         # Constants
         lr = 0.5
@@ -1875,7 +1873,7 @@ def test_kv_opt_state_w_offloading(
             output_dtype=output_dtype,
             share_table=share_table,
             num_buckets=num_buckets,
-            enable_optimizer_offloading=False,
+            enable_optimizer_offloading=enable_optimizer_offloading,
         )
 
         # Generate inputs