(Closes #35, updates #52) Changes for Jax-Metal #57
Conversation
This file contains 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
The new implementation of top_k is enabled with a flag. It should check the current device and use jax.lax.top_k if not running on jax-metal. This is not a general-purpose implementation of top_k but hopefully sufficient for sampling.
…ers code back for gpu/cpu; added jax-metal library to poetry"
|
Added back the scaled_rope that was accidentally deleted by me before. Tested with jax-metal and with jax-cpu. @nix sure, we have to wait for them. My hunch is that he will want to run the original code on "his" devices (read: TPU and jax-gpu) and we should try to come as close to that as possible but not force him to run watered down code. But now @Arrabonae has the choice. I will try to integrate the Frog branch next. |
Merging frog into main as his ideas have proven to be valid
|
Integrated the Frog branch and ran main.py on jax-metal as well as jax-cpu. Our lass is happy. Can someone validate on his rig, preferably on a jax-gpu, too? Note that I did not test the eval_main.py yet due to the README update by xjdr. |
|
Frog entropix/main.py successfully tested under
|
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.
Based on #52
Added conditionals based on jax.extended.backend.get_backend(). Added jax-metal to toml. Verified that it is running on both Metal and CPU.