Skip to content

Fix bug in momentum type declaration in HIP TBE kernel #147

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

Open
aryaman-gupta wants to merge 5 commits into
base: aryaman/upstream
Choose a base branch
from
Open

Fix bug in momentum type declaration in HIP TBE kernel #147

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
4f22562
use acc_type for momentum instead of cache_t
aryaman-gupta Mar 19, 2026
0dfaed7
backward_adagrad_test.py: updates unit test to vary params
aryaman-gupta Mar 23, 2026
3a179ad
split env variable calls
aryaman-gupta Mar 23, 2026
77f752b
backward_adagrad_test.py: modifies test name and desc
aryaman-gupta Mar 23, 2026
ba53028
backward_adagrad_test: reset env variable after test
aryaman-gupta Mar 23, 2026
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
15 changes: 9 additions & 6 deletions ...mm_gpu/codegen/training/backward/rocm/embedding_backward_split_device_kernel_template.hip
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -32,6 +32,7 @@
#include <hip/hip_runtime.h>
#include <hip/hip_fp16.h>

#include <ATen/AccumulateType.h>
#include "fbgemm_gpu/rocm/split_embeddings_common.h"

namespace fbgemm_gpu::rocm {
Expand All @@ -48,8 +49,9 @@ struct rowwise_adagrad_optimizer_t
{
if constexpr(segment_split == 0)
{
cache_t * p_momentum = reinterpret_cast<cache_t*>(karg.p_momentum);
cache_t momentum = p_momentum[row_index]; // should be s_load
using momentum_t = at::acc_type<cache_t, true>;
momentum_t* p_momentum = reinterpret_cast<momentum_t*>(karg.p_momentum);
momentum_t momentum = p_momentum[row_index]; // should be s_load
// compute per row square sum
cache_t local_sum_squre = .0f;
if constexpr(weight_decay_mode == 1)
Expand All @@ -72,11 +74,11 @@ struct rowwise_adagrad_optimizer_t
}
}

cache_t avg_square =
wave_reduce<reduce_op_sum_t<cache_t>, cache_t, AMDGCN_WAVE_SIZE>(local_sum_squre) /
momentum_t avg_square =
static_cast<momentum_t>(wave_reduce<reduce_op_sum_t<cache_t>, cache_t, AMDGCN_WAVE_SIZE>(local_sum_squre)) /
embedding_dim;

cache_t momentum_new = momentum + avg_square;
momentum_t momentum_new = momentum + avg_square;

cache_t multiplier = karg.learning_rate / (sqrtf(momentum_new) + karg.eps);
cache_t correction;
Expand Down Expand Up @@ -164,7 +166,8 @@ __device__ void split_tbe_backward_hip_kernel_{{ kdesc }}(
const int64_t emb_idx = linear_index - hash_size;

p_emb_table += hash_size * emb_dim;
opt_karg.p_momentum = reinterpret_cast<void*>(reinterpret_cast<cache_t*>(opt_karg.p_momentum) + hash_size);
using momentum_t = at::acc_type<cache_t, true>;
opt_karg.p_momentum = reinterpret_cast<void*>(reinterpret_cast<momentum_t*>(opt_karg.p_momentum) + hash_size);

const int32_t segment_length = segment_end - segment_start;

Expand Down
2 changes: 2 additions & 0 deletions fbgemm_gpu/test/tbe/training/backward_adagrad_common.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -50,6 +50,7 @@
gpu_unavailable,
gradcheck,
optests,
skipIfNotRocm,
skipIfRocm,
TEST_WITH_ROCM,
use_cpu_strategy,
Expand All @@ -62,6 +63,7 @@
gpu_unavailable,
gradcheck,
optests,
skipIfNotRocm,
skipIfRocm,
TEST_WITH_ROCM,
use_cpu_strategy,
Expand Down
142 changes: 131 additions & 11 deletions fbgemm_gpu/test/tbe/training/backward_adagrad_test.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,7 @@
from fbgemm_gpu.split_table_batched_embeddings_ops_training import (
ComputeDevice,
SplitTableBatchedEmbeddingBagsCodegen,
WeightDecayMode,
)
from hypothesis import given, settings

Expand All @@ -38,6 +39,7 @@
gpu_unavailable,
optests,
PoolingMode,
skipIfNotRocm,
SparseType,
st,
)
Expand Down Expand Up @@ -234,28 +236,146 @@ def test_backward_adagrad_fp16_pmSUM_with_max_norm( # noqa C901
**kwargs,
)

@unittest.skipIf(*gpu_unavailable)
def test_backward_adagrad_fp16_pmSUM_D320(self) -> None:
def _test_backward_adagrad_rocm_kernel(
self,
T: int,
D: int,
B: int,
log_E: int,
L: int,
weights_precision: SparseType,
weighted: bool,
weight_decay_mode: WeightDecayMode,
) -> None:
"""Helper method for ROCm backward kernel tests."""
execute_backward_adagrad(
T=2,
# using D=80 since the test harness multiplies D by 4, so 80*4=320
D=80,
B=16,
log_E=4,
L=4,
T=T,
D=D,
B=B,
log_E=log_E,
L=L,
D_gradcheck=1,
weights_precision=SparseType.FP16,
weights_precision=weights_precision,
stochastic_rounding=False,
weighted=False,
weighted=weighted,
row_wise=True,
mixed=False,
mixed_B=False,
use_cache=False,
cache_algorithm=CacheAlgorithm.LRU,
pooling_mode=PoolingMode.SUM,
use_cpu=False,
output_dtype=SparseType.FP16,
output_dtype=weights_precision,
weight_decay_mode=weight_decay_mode,
)

@given(
T=st.integers(min_value=1, max_value=5),
D=st.sampled_from([16, 32, 40, 48, 64, 80]),
B=st.integers(min_value=1, max_value=128),
log_E=st.integers(min_value=3, max_value=5),
L=st.integers(min_value=2, max_value=20),
weights_precision=st.sampled_from([SparseType.FP16, SparseType.FP32]),
weighted=st.booleans(),
weight_decay_mode=st.sampled_from(
[
WeightDecayMode.NONE,
WeightDecayMode.L2,
WeightDecayMode.DECOUPLE,
]
),
)
@settings(**common_settings)
@unittest.skipIf(*gpu_unavailable)
@skipIfNotRocm("Test evaluates fallback kernel on ROCm")
def test_backward_adagrad_rocm_fallback_kernel(
self,
T: int,
D: int,
B: int,
log_E: int,
L: int,
weights_precision: SparseType,
weighted: bool,
weight_decay_mode: WeightDecayMode,
) -> None:
env_var = "FBGEMM_TBE_ROCM_HIP_BACKWARD_KERNEL"
original_value = os.environ.get(env_var)
os.environ[env_var] = "0"
logging.info(
f"Testing ROCm backward kernel with {env_var}=0 (stock)"
)
try:
self._test_backward_adagrad_rocm_kernel(
T=T,
D=D,
B=B,
log_E=log_E,
L=L,
weights_precision=weights_precision,
weighted=weighted,
weight_decay_mode=weight_decay_mode,
)
finally:
# Restore original value
if original_value is None:
os.environ.pop(env_var, None)
else:
os.environ[env_var] = original_value

@given(
T=st.integers(min_value=1, max_value=5),
D=st.sampled_from([16, 32, 40, 48, 64, 80]),
B=st.integers(min_value=1, max_value=128),
log_E=st.integers(min_value=3, max_value=5),
L=st.integers(min_value=2, max_value=20),
weights_precision=st.sampled_from([SparseType.FP16, SparseType.FP32]),
weighted=st.booleans(),
weight_decay_mode=st.sampled_from(
[
WeightDecayMode.NONE,
WeightDecayMode.L2,
WeightDecayMode.DECOUPLE,
]
),
)
@settings(**common_settings)
@unittest.skipIf(*gpu_unavailable)
@skipIfNotRocm("Test evaluates ROCm optimized backward kernel")
def test_backward_adagrad_rocm_optimized_kernel(
self,
T: int,
D: int,
B: int,
log_E: int,
L: int,
weights_precision: SparseType,
weighted: bool,
weight_decay_mode: WeightDecayMode,
) -> None:
env_var = "FBGEMM_TBE_ROCM_HIP_BACKWARD_KERNEL"
original_value = os.environ.get(env_var)
os.environ[env_var] = "1"
logging.info(
f"Testing ROCm backward kernel with {env_var}=1 (optimized)"
)
try:
self._test_backward_adagrad_rocm_kernel(
T=T,
D=D,
B=B,
log_E=log_E,
L=L,
weights_precision=weights_precision,
weighted=weighted,
weight_decay_mode=weight_decay_mode,
)
finally:
# Restore original value
if original_value is None:
os.environ.pop(env_var, None)
else:
os.environ[env_var] = original_value

@unittest.skipIf(*gpu_unavailable)
@unittest.skipIf(*gpu_memory_lt_gb(40))
Expand Down