Skip to content

Fix flaky Lion optimizer test on HIP/AMD GPUs #95

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
kudomcho wants to merge 1 commit into
base: rocm_enabled_multi_backend
Choose a base branch
from
Open

Fix flaky Lion optimizer test on HIP/AMD GPUs #95

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
16 changes: 7 additions & 9 deletions tests/test_optim.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -160,9 +160,11 @@ def test_optimizer32bit(dim1, dim2, gtype, optim_name):
rtol=rtol,
)

# since Lion can have pretty noisy updates where things lie at the boundary
# allow up to 10 errors for Lion
assert_most_approx_close(p1, p2.float(), atol=atol, rtol=rtol, max_error_count=10)
# Lion uses sgn() which amplifies FMA rounding differences at the
# sign boundary; HIP/clang contracts FMA differently from nvcc,
# producing more boundary flips on AMD GPUs.
lion_max_err = max(10, p1.numel() // 200000)
assert_most_approx_close(p1, p2.float(), atol=atol, rtol=rtol, max_error_count=lion_max_err)

if i % (k // 5) == 0 and i > 0:
path = get_temp_dir()
Expand All @@ -172,18 +174,14 @@ def test_optimizer32bit(dim1, dim2, gtype, optim_name):
bnb_optimizer = str2optimizers[optim_name][1]([p2])
bnb_optimizer.load_state_dict(torch.load(join(path, "opt.pt")))
rm_path(path)
# since Lion can have pretty noisy updates where things lie at the boundary
# allow up to 10 errors for Lion
assert_most_approx_close(p1, p2.float(), atol=atol, rtol=rtol, max_error_count=10)
assert_most_approx_close(p1, p2.float(), atol=atol, rtol=rtol, max_error_count=lion_max_err)
for name1, name2 in str2statenames[optim_name]:
# since Lion can have pretty noisy updates where things lie at the boundary
# allow up to 10 errors for Lion
assert_most_approx_close(
torch_optimizer.state[p1][name1],
bnb_optimizer.state[p2][name2],
atol=atol,
rtol=rtol,
max_error_count=10,
max_error_count=lion_max_err,
)

if gtype != torch.float32:
Expand Down