Meta Llama 3.1 with MLX LM and the MLX Python API as Context

mlx_api_prompt.py

import os
import mlx.core as mx
from mlx_lm import load, generate

filename = os.path.join(os.path.dirname(mx.__file__), "core/__init__.pyi")
with open(filename, 'r') as fid:
    prompt = fid.read()
prompt += "\nHow do you write a self-attention layer using the above API in MLX?"

model, tokenizer = load("mlx-community/meta-Llama-3.1-8B-Instruct-4bit")

messages = [{"role": "user", "content": prompt}]

prompt = tokenizer.apply_chat_template(
    messages, tokenize=False, add_generation_prompt=True
)

generate(
    model,
    tokenizer,
    prompt,
    512,
    verbose=True,
    temp=0.0,
    max_kv_size=4096,
)

Ah sorry about that. This relies on a not yet released MLX / MLX LM. We should have the release out which supports this by Thursday. In the meantime here are the instructions to build from source:

pip install git+https://github.com/ml-explore/mlx.git
pip install git+https://github.com/ml-explore/mlx-examples.git@use_fast_rope

Getting an error from this

Running command git checkout -b use_fast_rope --track origin/use_fast_rope
Switched to a new branch 'use_fast_rope'
branch 'use_fast_rope' set up to track 'origin/use_fast_rope'.
Resolved https://github.com/ml-explore/mlx-examples.git to commit 0a52a9d55a5c1bfa6b85ab63b259e4c86e98b62a
ERROR: git+https://github.com/ml-explore/mlx-examples.git@use_fast_rope does not appear to be a Python project: neither 'setup.py' nor 'pyproject.toml' found.

Any insights??

@seshakiran See the instructions I posted.

Alternatively, if you already have core/__init__.pyi, you can do pip install "git+https://github.com/ml-explore/mlx-examples.git@use_fast_rope#egg=mlx-lm&subdirectory=llms" --no-deps.

@awni Could see it being useful to kv cache MLX docs like this for porting.

@awni on that note above, have you or anyone else found if the method above (getting the methods for MLX via the library) act as the best docs for porting code to MLX?

Was about to start working on a LoRA trainer for FLUX.1 (starting with looking at mflux), and it's a beast of torch/diffusers/cuda code.

have you or anyone else found if the method above (getting the methods for MLX via the library) act as the best docs for porting code to MLX?

I haven't tried much there tbh. The API I use above includes the docstrings (from which a lot of the docs are autogenerated) so there would be substantial overlap between using that and using the actual docs.

have you or anyone else found if the method above (getting the methods for MLX via the library) act as the best docs for porting code to MLX?

I haven't tried much there tbh. The API I use above includes the docstrings (from which a lot of the docs are autogenerated) so there would be substantial overlap between using that and using the actual docs.

I've been using the MLX .md docs and formatting them with structure via https://github.com/simonw/files-to-prompt

The docstrings is a good idea.