Repository navigation
Conversation
The 65,536-bin LUT can map a value to a neighbouring code near zero, so 0.0 was stored as +2.125e-5 * absmax. Use the LUT entry as a start and walk to the nearest code with the default backend's midpoint rule.
This branch has not been deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
On CPU,
quantize_blockwiselooks codes up in a 65,536-bin table. Near zero a bin can be wider than the gap between neighbouring codes, so an element can get a neighbouring code instead of the nearest one. In the default signed map, 0.0 is stored as +2.125e-5 * absmax. WithAdam8biton CPU, a weight whose gradient is always 0 then gets a non-zero first moment and drifts, whiletorch.optim.Adamleaves it unchanged. CUDA and the default backend pick the nearest code.The fix keeps the table but uses its entry only as a starting index. It then steps to the exact bucket with the same midpoint rule as the default backend (ties go to the lower index), usually 0 or 1 steps. The x86 loop and the NEON loop both use it. The NEON change is the same edit, but I could not compile it here.
I added
test_blockwise_quantization_nearest_code: every code value must map to its own index and 0.0 must round-trip, on both the signed and unsigned dynamic map.Testing, CPU build of main (
a616886cacf86701514f0cb1d6be760f967481c0) in Docker:Adam8bitagainsttorch.optim.Adamon a zero-gradient coordinate): exit 1 on main (index 131 instead of 127, drift of -3.6e-6), exit 0 with the fix.pytest tests/test_functional.py -k "cpu and blockwise_quantization": the new test fails on main (2 failed) and all 4 selected tests pass with the fix.tests/test_optim.pyon CPU and the aarch64 build.Fixes #2118