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,