mrq/bitsandbytes-rocm
1
0
Fork 1
Code Issues 1 Pull Requests Packages Projects Releases Wiki Activity

Fixed noisy tests for 8-bit Lion.

Browse Source
This commit is contained in:
Tim Dettmers 2023-04-11 08:42:41 -07:00
parent 0b2ebcdab9
commit 792af5c883
1 changed files with 13 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

17
tests/test_optim.py
View File

@ -18,6 +18,13 @@ import bitsandbytes.functional as F
k = 20
def assert_most_approx_close(a, b, rtol=1e-3, atol=1e-3, max_error_count=0):
idx = torch.isclose(a, b, rtol, atol)
error_count = (idx == 0).sum().item()
if error_count > max_error_count:
print(f"Too many values not close: assert {sumval} < {count}")
torch.testing.assert_allclose(a, b, rtol, atol)
def get_temp_dir():
path = f"/tmp/autoswap/{str(uuid.uuid4())}"
@ -306,7 +313,9 @@ def test_optimizer8bit(dim1, dim2, gtype, optim_name):
bnb_optimizer.step()
torch_optimizer.step()
torch.testing.assert_allclose(p1, p2.float(), atol=patol, rtol=prtol)
# since Lion can have pretty noisy updates where things lie at the boundary
# allow up to 5 errors for Lion
assert_most_approx_close(p1, p2.float(), patol, prtol, max_error_count=5)
dequant_states = []
for name1, name2, qmap, max_val in str2statenames[optim_name]:
@ -388,9 +397,9 @@ def test_optimizer8bit(dim1, dim2, gtype, optim_name):
== 0
)
assert num_not_close.sum().item() < 20
torch.testing.assert_allclose(
p1, p2.float(), atol=patol, rtol=prtol
)
# since Lion can have pretty noisy updates where things lie at the boundary
# allow up to 5 errors for Lion
assert_most_approx_close(p1, p2.float(), patol, prtol, max_error_count=5)
# the parameters diverge quickly. Here we keep them close
# together so we can test against the Adam error