🍏🤖 mlx-playground

Run fast transformer decoders on your Macbooks' GPU! Working towards a fast reimplementation of GPT-2 and Llama-like models in mlx.

The aim is that the only dependencies are:

• mlx
• sentencepiece
• tqdm
• numpy

With an optional dev dependency of:

• transformers for downloading and converting weights

Accomplishments

• [x] ~~makemore llama reimplementation~~ (train your own w/ python train.py!)
• [x] BERT merged into mlx-examples
• [x] Phi-2 merged into mlx-examples
• [x] AdamW merged into mlx

Remaining Goals

This project will be considered complete once these goals are achieved.

• [ ] finetune BERT
• [ ] GPT-2 reimplementation and loading in MLX
• [ ] speculative decoding
• [ ] learning rate scheduling

Installation

poetry install --no-root

Phi-2

To download and convert the model:

python phi2/convert.py

That will fill in weights/phi-2.npz.

🚧 (Not yet done) To run the model:

python phi2/generate.py

Acknowledgements

Some great resources:

• Brian Kitano's LLaMa from Scratch
• PyTorch lab's gpt-fast