Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Partial revert of D66986498 (Optimized backward pass for ROCm devices, pt 1), 2nd attempt #3637

Open
wants to merge 1 commit into
base: main
Choose a base branch
from
Open

Partial revert of D66986498 (Optimized backward pass for ROCm devices, pt 1), 2nd attempt #3637

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 0 additions & 10 deletions fbgemm_gpu/cmake/TbeTraining.cmake
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
################################################################################
Expand Down Expand Up @@ -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
Expand Down
19 changes: 0 additions & 19 deletions fbgemm_gpu/cmake/tbe_sources.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
################################################################################
Expand Down
22 changes: 0 additions & 22 deletions fbgemm_gpu/codegen/genscript/generate_backward_split.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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]
Expand Down Expand Up @@ -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()
Expand Down
7 changes: 3 additions & 4 deletions fbgemm_gpu/codegen/training/backward/embedding_backward_dense_host_cpu.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -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);
Expand Down
15 changes: 3 additions & 12 deletions fbgemm_gpu/codegen/training/backward/embedding_backward_split_host_template.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -152,7 +152,6 @@ enum SSDTensor {
{%- else %}
D_offsets,
max_D,
mixed_D,
{%- endif %} {# /* if nobag */ #}
hash_size_cumsum,
total_hash_size_bits,
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -306,7 +304,6 @@ enum SSDTensor {
D_offsets,
total_D,
max_D,
mixed_D,
{%- endif %}
hash_size_cumsum,
total_hash_size_bits,
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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 %}
Expand Down Expand Up @@ -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;
Expand Down Expand Up @@ -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();
Expand Down Expand Up @@ -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 %}
Expand Down Expand Up @@ -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});

Expand Down
Loading