Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

mlx: Add initial mlx implementation #58

Open
wants to merge 1 commit into
base: main
Choose a base branch
from
Open

mlx: Add initial mlx implementation #58

wants to merge 1 commit into from

Conversation

tmc
Copy link

@tmc tmc commented Oct 9, 2024

This adds an initial adaptation of the jax and pytorch implementations to mlx.

@artus-LYTiQ
Copy link

Lol, you beat me to it. Will check and compare

@artus-LYTiQ
Copy link

gemini_kvcache.py.txt
gemini_main.py.txt
gemini_model.py.txt
gemini_sampler.py.txt
gemini_stats.py.txt
gemini_weights.py.txt

Hope you can make something out of the attached files. I tried to one-shot the Jax code to MLX as it is really quite similar. But the subtle differences of course let that fail. Then I fiddled around with Sonnet 3.5 until it ran but it is obviously buggy and my brain is fried for today. @tmc

@samefarrar
Copy link

If you're interested, and xjdr would rather keep them seperate, I have a MLX fork that runs on my M2 Macbook out the box. Have implemented server, SSE, and a frontend for changing system prompts etc. still WIP though but caught up to the main branch (frog branch incoming).

@artus-LYTiQ
Copy link

Stuck at work for the next five or six hours but would love to integrate the three efforts (or push the most promising further). The core question will imo be how xjdr will want to treat imperfect/deviating implementations and how fast we can come up with tooling that measures how big those deviations actually are.

@d-r-w
Copy link

d-r-w commented Oct 10, 2024

If you're interested, and xjdr would rather keep them seperate, I have a MLX fork that runs on my M2 Macbook out the box. Have implemented server, SSE, and a frontend for changing system prompts etc. still WIP though but caught up to the main branch (frog branch incoming).

Your fork is really cool but as you mention it's very different
Love the idea of a out-of-box adapter assuming there aren't random bugs introduced by claude

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

4 participants