Partial revert of D66986498 (Optimized backward pass for ROCm devices…

…, pt 1), 2nd attempt Summary: - Full revert of D66986498 for debugging https://fb.workplace.com/groups/fbgemmusers/permalink/9438488366231860/ (D68724497 didn't fix the regressions) Differential Revision: D68861705
pytorch · Jan 29, 2025 · beb18d6 · beb18d6
1 parent f258326
commit beb18d6
Show file tree

Hide file tree

Showing 17 changed files with 11 additions and 1,688 deletions.
diff --git a/fbgemm_gpu/cmake/TbeTraining.cmake b/fbgemm_gpu/cmake/TbeTraining.cmake
@@ -42,14 +42,6 @@ handle_genfiles(gen_py_files_training)
 handle_genfiles(gen_py_files_defused_optim)
 
 
-################################################################################
-# FBGEMM_GPU Generated HIP-Specific Sources
-################################################################################
-
-get_tbe_sources_list(gen_hip_files_training)
-handle_genfiles_rocm(gen_hip_files_training)
-
-
 ################################################################################
 # TBE C++ Training Targets
 ################################################################################
@@ -160,8 +152,6 @@ gpu_cpp_library(
     ${gen_cpu_files_training}
   GPU_SRCS
     ${gen_gpu_files_training}
-  HIP_SPECIFIC_SRCS
-    ${gen_hip_files_training}
   GPU_FLAGS
     ${TORCH_CUDA_OPTIONS}
   DEPS

diff --git a/fbgemm_gpu/cmake/tbe_sources.py b/fbgemm_gpu/cmake/tbe_sources.py
@@ -473,25 +473,6 @@
     ]
 )
 
-gen_hip_files_training = [
-    "gen_embedding_backward_split_{}{}_device_kernel_hip.hip".format(
-        "weighted" if weighted else "unweighted",
-        "_nobag" if nobag else "",
-    )
-    for nobag in [
-        True,
-        False,
-    ]
-    for weighted in (
-        [
-            True,
-            False,
-        ]
-        if not nobag
-        else [False]
-    )
-]
-
 ################################################################################
 # Python Training Code
 ################################################################################

diff --git a/fbgemm_gpu/codegen/genscript/generate_backward_split.py b/fbgemm_gpu/codegen/genscript/generate_backward_split.py
@@ -310,27 +310,6 @@ def generate_backward_indices() -> None:
                 ssd=ssd,
             )
 
-    @staticmethod
-    def generate_rocm_backward_split(**kwargs: Any) -> None:
-        # Generate backward device kernels based on weighted (True/False), VBE
-        # (True/False), no bag (True/False)
-        template_filepath = (
-            "training/backward/rocm/embedding_backward_split_device_kernel_template.hip"
-        )
-
-        BackwardSplitGenerator.render_backward_templates(
-            template_filepath,
-            "",
-            "{}gen_embedding_backward_{}_device_kernel_hip.hip",
-            {
-                "has_gpu_support": True,
-                "has_vbe_support": False,
-                "has_ssd_support": False,
-                "dense": False,
-                "gen_once": False,
-            },
-        )
-
     @staticmethod
     def generate_python_sources(
         all_optimizers: List[str], ssd_optimizers: List[str]
@@ -390,7 +369,6 @@ def generate() -> None:
             BackwardSplitGenerator.generate_backward_split(
                 ssd_tensors=ssd_tensors, **optimizer
             )
-        BackwardSplitGenerator.generate_rocm_backward_split()
 
         # Generate common device kernels for backwards
         BackwardSplitGenerator.generate_backward_device()

diff --git a/fbgemm_gpu/codegen/training/backward/embedding_backward_dense_host_cpu.cpp b/fbgemm_gpu/codegen/training/backward/embedding_backward_dense_host_cpu.cpp
@@ -171,8 +171,7 @@ Tensor split_embedding_codegen_lookup_dense_function(
         Tensor>& /* vbe_B_offsets_rank_per_feature = std::nullopt */,
     c10::SymInt /* max_B = -1 */,
     c10::SymInt /* max_B_feature_rank = -1 */,
-    c10::SymInt /* vbe_output_size = -1 */,
-    bool /* mixed_D = true */) {
+    c10::SymInt /* vbe_output_size = -1 */) {
   return SplitLookupFunction_Dense_Op::apply(
       host_weights,
       weights_offsets,
@@ -191,15 +190,15 @@ Tensor split_embedding_codegen_lookup_dense_function(
 // Deprecated for fb namespace! Please use fbgemm namespace instead!
 TORCH_LIBRARY_FRAGMENT(fb, m) {
   m.def(
-      "dense_embedding_codegen_lookup_function(Tensor dev_weights, Tensor weights_offsets, 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=0, Tensor? B_offsets=None, Tensor? vbe_output_offsets_feature_rank=None, Tensor? vbe_B_offsets_rank_per_feature=None, SymInt max_B=-1, SymInt max_B_feature_rank=-1, SymInt vbe_output_size=-1, bool mixed_D=True) -> Tensor");
+      "dense_embedding_codegen_lookup_function(Tensor dev_weights, Tensor weights_offsets, 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=0, Tensor? B_offsets=None, Tensor? vbe_output_offsets_feature_rank=None, Tensor? vbe_B_offsets_rank_per_feature=None, SymInt max_B=-1, SymInt max_B_feature_rank=-1, SymInt vbe_output_size=-1) -> Tensor");
   DISPATCH_TO_CPU(
       "dense_embedding_codegen_lookup_function",
       split_embedding_codegen_lookup_dense_function);
 }
 
 TORCH_LIBRARY_FRAGMENT(fbgemm, m) {
   m.def(
-      "dense_embedding_codegen_lookup_function(Tensor dev_weights, Tensor weights_offsets, 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=0, Tensor? B_offsets=None, Tensor? vbe_output_offsets_feature_rank=None, Tensor? vbe_B_offsets_rank_per_feature=None, SymInt max_B=-1, SymInt max_B_feature_rank=-1, SymInt vbe_output_size=-1, bool mixed_D=True) -> Tensor");
+      "dense_embedding_codegen_lookup_function(Tensor dev_weights, Tensor weights_offsets, 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=0, Tensor? B_offsets=None, Tensor? vbe_output_offsets_feature_rank=None, Tensor? vbe_B_offsets_rank_per_feature=None, SymInt max_B=-1, SymInt max_B_feature_rank=-1, SymInt vbe_output_size=-1) -> Tensor");
   DISPATCH_TO_CPU(
       "dense_embedding_codegen_lookup_function",
       split_embedding_codegen_lookup_dense_function);

diff --git a/fbgemm_gpu/codegen/training/backward/embedding_backward_split_host_template.cpp b/fbgemm_gpu/codegen/training/backward/embedding_backward_split_host_template.cpp
@@ -152,7 +152,6 @@ enum SSDTensor {
           {%- else %}
           D_offsets,
           max_D,
-          mixed_D,
           {%- endif %} {# /* if nobag */ #}
           hash_size_cumsum,
           total_hash_size_bits,
@@ -225,7 +224,6 @@ enum SSDTensor {
         Variable(), // D_offsets
         Variable(), // total_D
         Variable(), // max_D
-        Variable(), // mixed_D
         {%- endif %}
         Variable(), // hash_size_cumsum
         Variable(), //total_hash_size_bits
@@ -306,7 +304,6 @@ enum SSDTensor {
           D_offsets,
           total_D,
           max_D,
-          mixed_D,
           {%- endif %}
           hash_size_cumsum,
           total_hash_size_bits,
@@ -487,7 +484,6 @@ Tensor
     {%- else %}
     const Tensor& D_offsets,
     const c10::SymInt max_D,
-    const bool mixed_D,
     {%- endif %}
     const Tensor& hash_size_cumsum,
     const int64_t total_hash_size_bits,
@@ -570,7 +566,6 @@ class {{ autograd_func }} :
     const Tensor& D_offsets,
     const c10::SymInt total_D,
     const c10::SymInt max_D,
-    const bool mixed_D,
     {%- else %}
     const c10::SymInt D,
     {%- endif %}
@@ -767,7 +762,6 @@ class {{ autograd_func }} :
 
     {%- if not nobag %}
     ctx->saved_data["max_D"] = max_D;
-    ctx->saved_data["mixed_D"] = mixed_D;
     ctx->saved_data["pooling_mode"] = pooling_mode;
     {%- else %}
     ctx->saved_data["D"] = D;
@@ -883,7 +877,6 @@ class {{ autograd_func }} :
 
     {%- if not nobag %}
     auto max_D = ctx->saved_data["max_D"].toSymInt();
-    const auto mixed_D = ctx->saved_data["mixed_D"].toBool();
     auto pooling_mode = ctx->saved_data["pooling_mode"].toInt();
     {%- else %}
     auto D = ctx->saved_data["D"].toSymInt();
@@ -1079,11 +1072,10 @@ Tensor {{ bwd_mdesc }}_embedding_codegen_lookup_{{ optimizer }}_function(
     {%- if ssd %}
     const std::optional<at::TensorList>& ssd_tensors = std::nullopt,
     {%- endif %}
-    const double gwd_lower_bound = 0,
+    const double gwd_lower_bound = 0
     {%- else %}
-    const c10::SymInt vbe_output_size = -1,
+    const c10::SymInt vbe_output_size = -1
     {%- endif %}
-    const bool mixed_D = true
 ) {
   // TODO: refactor into macro
   {%- if has_gpu_support %}
@@ -1199,8 +1191,7 @@ TORCH_LIBRARY_FRAGMENT({{ lib_name }}, m) {
           {%- if ssd %}
           "    Tensor[]? ssd_tensors=None,"
           {%- endif %}
-          "   float gwd_lower_bound=0, "
-          "   bool mixed_D=True"
+          "   float gwd_lower_bound=0 "
           ") -> Tensor",
           {PT2_COMPLIANT_TAG});