From 6a1be54fdc22ce829c9a1ddf6d9837ec29c36f35 Mon Sep 17 00:00:00 2001 From: "jiufei.liu" Date: Wed, 23 Sep 2026 15:04:24 +0800 Subject: [PATCH] test: align post-reload param check threshold with pre-reload check Commit 36f5c4f4 raised the max_error_count of the parameter comparison after each optimizer step from 10 to 15 to tolerate Lion's noisy boundary updates, but did not update the identical comparison that runs right after the state_dict save/load roundtrip, which still allows only 10. Both assertions compare the same tensors: between them only state_dict(), optimizer recreation and load_state_dict() run, and load_state_dict restores optimizer state without touching parameters. The mismatch count after the roundtrip is therefore identical to the one measured before it. With the thresholds out of sync, a step producing 11-15 mismatched elements passes the first check but fails the second, misleadingly pointing at the state_dict roundtrip instead of the tolerated Lion noise. Set the post-reload threshold to 15 to match. --- tests/test_optim.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/test_optim.py b/tests/test_optim.py index 29736311d..a5ab82e06 100644 --- a/tests/test_optim.py +++ b/tests/test_optim.py @@ -218,8 +218,8 @@ def test_optimizer32bit(dim1, dim2, gtype, optim_name, device): buf.seek(0) bnb_optimizer.load_state_dict(torch.load(buf)) # 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) + # allow up to 15 errors for Lion + assert_most_approx_close(p1, p2.float(), atol=atol, rtol=rtol, max_error_count=15) 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