Merge VBE output [backend] reland (#5093)

Summary:
Pull Request resolved: #5093

X-link: https://github.com/facebookresearch/FBGEMM/pull/2100

----

# Context on the changes:

Currently, Torchrec merges the outputs of individual VBE TBE ops to be ordered by ranks using [_merge_variable_batch_embeddings](https://www.internalfb.com/code/fbsource/[3bd69d7fa3534144dcb0162ca59803a6c3ff6e70]/fbcode/torchrec/distributed/embedding_lookup.py?lines=593-604). This function seems to cause ~30% QPS regression compared to baseline (HBM+UVM) for Jupiter V1 model with VBE enabled.

To get rid of _merge_variable_batch_embeddings() function, we pre-allocate the `vbe_output` tensor which holds  outputs from all VBE ops and calculate `vbe_output_offsets` to allow each individual VBE ops to write to the correct location in the `vbe_output` tensor.

By default, `vbe_output` and `vbe_output_offsets` are `None`, which means VBE ops will return individual tensor the way it currently does. The feature is enabled when `vbe_output` and `vbe_output_offsets` are not `None`.

 ---
**NOTE**
1. This feature is currently supported for Sparse TBE.
2. The support is limited for CUDA.
3. For backward compatibility, we append the newly introduced `vbe_output` to the existing API. Hence, we need to make the `vbe_output` tensor as `optional` with default value as `None` (there's no default value for Tensor).
4. We *cannot* annotate  `vbe_output` because PyTorch registration does not support annotation of optional tensor.  Adding annotation will incur the following error below. This may cause some issues to support this on MTIA, if MTIA relies on tensor annotation.
```
E0903 09:50:32.966235 2850885 ExceptionTracer.cpp:227] exception stack complete
terminate called after throwing an instance of 'std::runtime_error'
  what():  expected ident but found '(' here:
split_embedding_codegen_lookup_adagrad_function_pt2(    Tensor placeholder_autograd_tensor,     Tensor[](a!) weights,     Tensor D_offsets,     SymInt total_D,     SymInt max_D,     Tensor hash_size_cumsum,     int total_hash_size_bits,     Tensor indices,     Tensor offsets,     int pooling_mode,     Tensor? indice_weights,     Tensor? feature_requires_grad,     int output_dtype,     Tensor?[](e!) aux_tensor,     int[] aux_int,     float[] aux_float,     bool[] aux_bool,     Tensor[](g!) momentum1, Tensor learning_rate_tensor, float[] optim_float,     SymInt max_B=-1,     SymInt max_B_feature_rank=-1,     SymInt vbe_output_size=-1,     Tensor?(t!) vbe_output=None ) -> Tensor
                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                            ~ <--- HERE
```

See https://docs.google.com/document/d/1h5YyeCjYmmN-CIFB98CrBf1uMksidPbNvM1rl8yZeds/edit?tab=t.0#heading=h.tdfkkc6ujdyl

----

This diff is a reland of D79704318 which all issues have been addressed.

## 1) pyper validation test

D79704318 was reverted as it broke pyper validation test (frontend/backend package compatibility issue), which blocks pyper releases. The issue is addressed in this diff.

Context: In pyper, changes in python would be included in frontend package (e.g., ads_dper3) and C++ in backend package (e.g., training_platform). If the diff contains both python and C++, there's a chance that some model will use mismatching packages. In other words, frontend package does not include the diff but backend does, and vice versa.

D83881544 is only enabling backend support (i.e., no one can actually use this feature, so TBE VBE will work as usual). Due to new Unified API changes, we need to pipeline optional tensor from frontend and requires python change.

Denote
- #0 as no D83881544 included
- #1 as D83881544 included

There are 4 scenarios:

(1) frontend #0 + old backend #0 - no issue
(2) frontend #1 + backend #1 - no issue
(3) frontend #0 + backend #1 - handled; TBE VBE will work normally.
(4) frontend #1 + backend #0 - no issue; the diff added warning that backend is old

There's another diff D79869613 in the stack that will enable frontend support (i.e., allow users to use this feature), which will go into __frontend package only__.
Now, the 1)-4) scenarios would remain the same, but new scenarios occur.

Denote
- #2 as D79869613 included

(5) frontend #2 + backend #1 - no issue, same as (2).
(6) frontend #2 (no feature enabled) + backend #0 - same as (4).
(7) frontend #2 (feature enabled) + backend #0 - **assertion error due to no backend support**, to prevent silent wrong behavior.

**To use the feature, this diff stack (D83881544 and D79869613) need to be included in both frontend and backend package.**

## 2) SEV

D79704318 caused SEV due to TBE v1 and v2 interfacing compatibility issue on lex_ig_o3_package. Unit tests to ensure v1 compatibility was added D83020965.

D83881544 passes the v1 compatibility test.

Detail on the root cause and fix:
https://docs.google.com/document/d/1XcYNfyiAn4aRFvjV0QG5aLiWKuuOWtJdLOMKNszZRpI/edit?tab=t.0#heading=h.psr4a2qn0mdk

------

Reviewed By: q10, renganxu

Differential Revision: D83881544

fbshipit-source-id: 5d63841bbf79a72219903e9d0f77ee3b998bc105

Author: spcyppt

fbgemm_gpu/codegen/genscript/generate_backward_split.py

@@ -422,6 +422,7 @@ def generate() -> None:
                 "lxu_cache_locations",  # 3
                 "uvm_cache_stats",  # 4
                 "prev_iter_dev",  # 5
+                "vbe_output_offsets",  # 6
             ],
             "aux_int": [
                 "iter",  # 0

fbgemm_gpu/codegen/genscript/optimizer_args.py

@@ -73,9 +73,7 @@ class OptimizerArgsSetItem:
     "row_counter_dev": "(q!)",
     "row_counter_uvm": "(r!)",
     "optim_tensor": "(s!)",
-    "delta_weights_host": "(t!)",
-    "delta_weights_dev": "(u!)",
-    "delta_weights_uvm": "(v!)",
+    "vbe_output": "(t!)",
 }
 
 ######################################################################

fbgemm_gpu/codegen/training/backward/embedding_backward_split_host_template.cpp

@@ -109,7 +109,12 @@ enum SSDTensor {
       gwd_lower_bound,
       {%- endif %} {# /* if is_gwd */ #}
       {%- endif %} {# /* if not nobag */ #}
+      {%- if vbe and not dense %}
+      {{ "is_experimental" if has_experimental else "false" }},
+      std::nullopt /* vbe_output */
+      {%- else %}
       {{ "is_experimental" if has_experimental else "false" }}
+      {%- endif %}
     );
 
     if (is_annotate_trace_enabled) {
@@ -474,7 +479,12 @@ Tensor {{ fwd_mdesc }}_embedding{{ ndesc }}_codegen_forward{{ desc_suffix }}_cud
     const int64_t iter,
     const double gwd_lower_bound,
     {%- endif %}
+    {%- if vbe and not dense %}
+    const bool is_experimental,
+    std::optional<Tensor> vbe_output = std::nullopt
+    {%- else %}
     const bool is_experimental
+    {%- endif %}
 );
 
 Tensor
@@ -708,7 +718,7 @@ class {{ autograd_func }} :
     static auto generate_vbe_metadata_op =
         torch::Dispatcher::singleton()
             .findSchemaOrThrow("fbgemm::generate_vbe_metadata", "")
-            .typed<std::tuple<Tensor, Tensor>(const Tensor&, const Tensor&, const Tensor&, const Tensor&, const int64_t, const bool, const c10::SymInt, const int64_t, const c10::SymInt)>();
+            .typed<std::tuple<Tensor, Tensor>(const Tensor&, const Tensor&, const Tensor&, const Tensor&, const int64_t, const bool, const c10::SymInt, const int64_t, const c10::SymInt, const std::optional<Tensor>&)>();
 
     auto [
         vbe_row_output_offsets,
@@ -729,7 +739,8 @@ class {{ autograd_func }} :
         {%- endif %}
         max_B_feature_rank,
         info_B_num_bits,
-        /*total_B=*/offsets.sym_size(0) - 1
+        /*total_B=*/offsets.sym_size(0) - 1,
+        std::nullopt /* pre-allocated vbe_output is not supported in TBE interface V1 or Dense TBE */
         );
     {%- endif %}
 

fbgemm_gpu/codegen/training/forward/embedding_forward_split_meta_template.cpp

@@ -6,6 +6,7 @@
  * LICENSE file in the root directory of this source tree.
  */
 
+// clang-format off
 {#
 // @lint-ignore LINTIGNORE
 // @lint-ignore-every CLANGFORMAT
@@ -103,7 +104,12 @@ Tensor
     const int64_t iter,
     const double gwd_lower_bound,
     {%- endif %}
+    {%- if vbe and not dense %}
+    const bool is_experimental,
+    std::optional<Tensor> vbe_output
+    {%- else %}
     const bool is_experimental
+    {%- endif %}
 ) {
     // NB: omitted the device tests TENSORS_ON_SAME_CUDA_GPU_IF_NOT_OPTIONAL
     {%- if not nobag %}
@@ -210,4 +216,4 @@ TORCH_LIBRARY_FRAGMENT(fbgemm, m) {
 {%- endfor %} {#-/* for is_gwd */#}
 {%- endif %} {#/* if (not nobag or (not weighted and not vbe)) */#}
 {%- endfor %} {#-/* for nobag */#}
-    // clang-format on
+  // clang-format on

fbgemm_gpu/codegen/training/forward/embedding_forward_split_template.cu

@@ -6,10 +6,10 @@
  * LICENSE file in the root directory of this source tree.
  */
 
-{#
 // @lint-ignore LINTIGNORE
 // @lint-ignore-every CLANGFORMAT
 // clang-format off
+{#
 // Note: clang-format off doesn't work with this templaterized code,
 // so we need to keep lint-ignore-every.
 // See https://fburl.com/dw9ljh4h
@@ -391,7 +391,12 @@ batch_index_select_dim0_codegen_forward_cuda(
     const int64_t iter,
     const double gwd_lower_bound,
     {%- endif %}
+    {%- if vbe and not dense %}
+    const bool is_experimental,
+    std::optional<Tensor> vbe_output
+    {%- else %}
     const bool is_experimental
+    {%- endif %}
     {%- endif %} {#- /*if is_index_select*/ #}
 ) {
     {%- if not nobag or is_index_select %}
@@ -529,11 +534,24 @@ batch_index_select_dim0_codegen_forward_cuda(
                 o_dtype == SparseType::BF16 || o_dtype == SparseType::INT8);
 
     {%- if vbe %}
-    // Use a 2D tensor to make it compatible with 2D PackedTensorsAccessor of other output
+    {%- if dense %}
     output = at::empty(
         {1, vbe_output_size},
         dev_weights.options().dtype(getScalarType(o_dtype))
-    );
+      );
+    {%- else %}
+    // Use a 2D tensor to make it compatible with 2D PackedTensorsAccessor of other output
+    TENSORS_ON_SAME_CUDA_GPU_IF_NOT_OPTIONAL(vbe_row_output_offsets, vbe_output);
+    if (vbe_output.has_value()){
+      output = vbe_output.value().reshape({1, -1});
+    }
+    else {
+      output = at::empty(
+        {1, vbe_output_size},
+        dev_weights.options().dtype(getScalarType(o_dtype))
+      );
+    }
+    {%- endif %} {#-/* if dense */#}
     {%- else %}
     int64_t total_adjusted_D = total_D;
     if (o_dtype == SparseType::INT8) {
@@ -877,7 +895,12 @@ TORCH_LIBRARY_FRAGMENT(fbgemm, m) {
           "    int iter, "
           "    float gwd_lower_bound, "
           {%- endif %}
+          {%- if vbe and not dense %}
+          "    bool is_experimental,"
+          "    Tensor? vbe_output"
+          {%- else %}
           "    bool is_experimental"
+          {%- endif %}
           ") -> Tensor"
           {%- if not dense and not nobag and not vbe %}
           // only split_embedding_codegen_forward_[un]weighted_cuda