Apple Join the AI Race with MLX

Apple's machine learning research team released a new framework called MLX to build foundation models. This is very unlike Apple because usually Apple doesn't release open source software. The goal with MLX is to design a machine learning framework that developers can use to build models that can run efficiently on Apple Silicon. Here is a tweet from Auni Hanoon, who is part of the Apple Machine learning team. Just in time for the holidays, we are releasing some new software today from Apple Machine Learning Research. MLX is an efficient machine learning framework specifically designed for Apple Silicon. So you will be able to run this on your laptop. Here's an example of running llama 70 build model on M2 Ultra. Later in the video I'll show you how to do this on your own machine. So you can train a transformer LLM or fine tune Loras on it. You can do text generation with Mistral, image generation with stable diffusion and speech recognition with Whisper. So this is pretty amazing because now you have specialized software for Apple Silicon. They are also releasing MLX data that is supposed to be framework agnostic, efficient and flexible package for data loading. So it's going to support Pytorch, JAX as well as mlx. Apple has been releasing their M series chips. The latest One is the M3 and it has support for Neural Engine. So it has the capability to run AI applications. But now with the release of mlx, we will have specialized software to support the AI capabilities on Apple Silicon. And with this we might finally see a foundation model from Apple. But we'll have to wait for that a bit longer. Okay, so let's quickly look at the MLX framework. Later in the video I'll show you how to install this on your own machine. How to run a llama 2 model using MLX. And then at the end we're going to look at an example of how to fine tune a Llama 2 model using the same framework. So MLX is a numpy like array framework designed for efficient and flexible machine learning on Apple Silicon, brought to you by the Apple Machine Learning Research. Now, if you're not familiar with NumPy, NumPy is a python package that is specifically designed for scientific computation. One thing you will notice even in this is their focus is on machine learning, not on generative AI. The Python API closely follows NumPy with a few exceptions, and we're going to talk about those. MLX also has a fully featured C API which closely follows the Python API. So this is pretty great because not only you can run this in Python but you can run C code which is going to be a lot faster. Now how this is different from numpy? So the first thing is composable function transformation. So MLX has composable function transformation for automatic differentiation, automatic vectorization and computation graph optimization. So seems like MLX is taking a lot from Pytorch when it comes to vector computation. Then it supports lazy computation. So computation in MLX are lazy, which means that arrays are only materialized when needed. It has multi device support, so operation can run on any supported device. Since there is unified memory on these Apple silicons, so both CPU as well as GPU is using the same memory. That's why I think it's very important for it to have this multi device support. Now according to them, MLX is inspired from frameworks like Pytorch, JAX and Array Fire. Now a noticeable difference between these frameworks and MLX is the use of unified memory model. So arrays in MLX live in shared memory. Operation on MLX array can be performed on any of the supported device types without performing data copies. So usually if you have a GPU and cpu, some of the operations are performed on GPU and some are performed on cpu. So there is this back and forth between CPU and gpu. But since Apple Silicon uses this unified memory, so that's why MLX might be a lot more efficient on Apple Silicon. With this release, the Apple team released this MLX example repo on GitHub. Now they have pretty good examples in there which are well documented. So I'm going to show you how to run llama2 on your local machine. You can also do parameter efficient fine tuning with Lora. You can generate images with stable diffusion. This is going to be pretty amazing. And you can also do speech recognition with OpenAI's whisper. So let's look at a few examples of how to do some of these things. Okay, so first we need to install MLX on a local machine. And in order to do that we're going to be using pip. But before that I'm going to create a CONDA environment. So first we are going to create a new virtual environment. So I'm going to be using CONDA for that. So we need to use this command, conda create dash n. The virtual environment that I'm creating is called MLX and I'm going to be using Python 3.10. Now I have already a virtual environment by this name, so I'm not going to create it again, but I will simply activate that. And the way you do it is we're going to use conda activate MLX command. Now you can see that it has changed the virtual environment that we are currently using. Okay, next we are going to install the MLX package using pip. So we will use the pip Install MLX command. As you can see, it's already installed on my system. Now, in order to run some of these MLX examples, we need to clone the examples folder. We need to clone the examples folder. So come here, click on this green button, then simply copy the URL. So I'll go back to my terminal and here I'm going to use git clone and then provide the repo id. Now the repo has been cloned on my local system. Next I need to move to that. So we are going to use the change directory command to move to that directory. And now you can see that I am within the MLX example directory. Now first let me show you how to run a llama model using the MLX package. So we will need to move to this llama folder. So again we're going to be using the CD command. So CD llama and now we are within the llama folder. So in order to run the model, we will use this llama PY file. However, first we need to download a llama model. Now, in order to run llama 2 models with MLX, we need to convert them into this new NPC format. So there is already a Hugging face repo which has a llama 27 build model in this format. So we're going to be using that. There is also a script in the MLX package that will let you convert model weights to this new NPC format. But for this video we're going to be just using the one available on Hugging Face. Now here are the steps that we need to follow. We already installed the MLX package, so we don't need to do that. Now we need to install the Hugging face hub, the hugging face transfer packages. So I copied the commands here and let's just run those. And we have already copied or cloned the ML examples GitHub repo. So I'm not going to do that again. Next we need to set the Hugging face hub enable hugging face transfer environment variable. So here we're going to set that here. And this next command is going to both download the model and then convert it to the NPC format that is needed by the MLX package. So let's run that. Now we are all set to run the llama model. So for that we are going to provide the path of the model in the NPC format, then the tokenizer and we are going to run this using Python, llama, PY and your prompt. So the prompt in this case I'm just using the example prompt that is provided in the repo. But let's see how quick it's able to run the model. So, so this is basically real time processing from the model. So it loaded the model from the disk and now if I hit enter or press Enter it will start generation. So this is the real time and it's actually pretty amazing. It's able to do this pretty fast. So the full generation took five seconds. Okay, I want to run this command again, but now also want to look at the GPU usage on my M2 Max. So let's run this. Okay, it's loading the model from disk and now it should ask me to press Enter to start generation. So let's see if it's using the GPU or not. And you can see that the GPU usage went up. So this is pretty awesome. Now within these two runs, both the prompt processing time as well as the full generation time remain pretty consistent. So this is a really good news. Next let's look at an example of how to train transformer language model using mlx. So they have provided an example in this transformerlm folder. So all we need to do is just run the main PY with the GPU flag. Now by default it's using this PTB corpus. Now this is English pentry bank corpus and in particular the section of the corpus corresponding to the Articles of Wall Street Journal. Now if you want to provide your own data set, you can do that with the data set option. So let's look at the code. So here we have quite a few options that are available. If you want to run this on gpu, you will need to use the GPU flag. Again, you can provide the data set with the data set flag. You can define the context window or context size. By default it's 1024 Number of blocks of the transformer that you want to add, Then the number of heads. Now number of iterations. By default it's set to a very high value, but I think we'll just run like hundred iterations just to show an example of the GPU usage. You can also define the learning rate. Right. So all the stuff that you will need order to train a language model, you can just pass on those parameters in here. So I change the directory to the main example folder, then we're going to change the directory to the Transformer LM folder. And here we have the main PY file that we need to run in order to train our LLM. Okay, so keep in mind I'm not running a full training job in here. Just want to show you how it can be done. So we're going to use the Python made py, we're going to be using the GPU and I just kept the number of iterations to 100. Let's see what the GPU usage looks like. So this is around 150 million parameter model and you can see that it's actually doing pretty good in terms of the training speed. Right? So the loss is decreasing pretty nicely. The iterations per second are also pretty good for training. So I think it's, it's pretty nice and it's using the full capacity of the GPU on Apple M2 Max. I just restarted the training process and just want to run it for longer. In a subsequent video I'll show you how to test the mod the trained model. This is just a very small model, so probably it's not really useful, but we will see if we can fine tune something like llama7b with the MLX package if that is possible. I think that is going to add a lot of value to this package specifically for Apple Silicon users. I hope you found this video useful. Thanks for watching and as always, see you in the next one.