Make AdEMAMix32bit inherit AdEMAMix so it allocates the double state buffer - #2076
Open
caiotheodoro wants to merge 1 commit into
Open
Make AdEMAMix32bit inherit AdEMAMix so it allocates the double state buffer#2076caiotheodoro wants to merge 1 commit into
caiotheodoro wants to merge 1 commit into
Conversation
…buffer AdEMAMix32bit and PagedAdEMAMix32bit subclassed Optimizer2State directly, so they never used AdEMAMix.init_state and allocated state1 with the shape of the parameter instead of (2, *p.shape). The ademamix kernels read m2 from the second half of state1, so on the CPU and default backends the first step raised RuntimeError, and the t_alpha/t_beta3 schedulers were silently ignored. Subclass AdEMAMix with optim_bits=32, mirroring AdEMAMix8bit, and add a test comparing both classes against AdEMAMix(optim_bits=32).
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
1 participant
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.
Fixes #2075.
AdEMAMix32bitandPagedAdEMAMix32bitsubclassedOptimizer2Statedirectly, sostate1was allocated asp.shapeinstead of(2, *p.shape)and thet_alpha/t_beta3scheduler inAdEMAMix.update_stepnever ran. On the CPU and default backends the firststep()raisedRuntimeError: output with shape [] doesn't match the broadcast shape [4096]. This re-parentsAdEMAMix32bitontoAdEMAMixwithoptim_bits=32, mirroringAdEMAMix8bit.PagedAdEMAMix32bitinherits the change. The constructor signature and defaults are unchanged.New test
test_ademamix32bit_matches_ademamixsteps both classes next toAdEMAMix(optim_bits=32)for 5 steps, scheduled and unscheduled, and assertsstate1.shape == (2, 4096)plus equalstate1,state2and parameters. Onmainit fails 6 of 6 with the RuntimeError above. On this branch 6 pass and 2 skip (paged on CPU, following the existing convention intest_optimizer32bit).Checks run on macOS arm64 (CPU build plus MPS), torch 2.14.0, Python 3.13.12:
pytest tests/test_optim.py -k ademamix -n 4: 32 failed, 90 passed, 26 skipped on this branch; 32 failed, 84 passed, 24 skipped onmain. The 32 failures are the same MPS-only cases on both trees (optimizer_update_8bit_blockwiseand paged buffers are not implemented on MPS)._ReferenceAdEMAMixonmainand on this branch. The 72AdEMAMix32bitandPagedAdEMAMix32bitcases go from RuntimeError to parameters bit-identical withAdEMAMix(optim_bits=32). The other 108 cases produce identical bytes on both trees.pre-commit run --all-files: all hooks pass.I could not run the CUDA kernel.
kOptimizer32bit2Stateloads m2 fromstate1[n + i], so on CUDA the old single-size buffer was read past its end.Drafted with Claude Opus / Fable 5.1. Reviewed by Muse Spark 1.3 and GLM 5.3 as judges before submission.