Getting Started with vLLM on TPUs

Key Takeaways

Running inference for your own large language model has gotten a lot easier [music] with the release of the new plugin for VLM that allows you to run on TPUs very efficiently and fast. [music] In this video, I'm going to walk you through the process of how to get VLM working on a TPU. If you're interested in learning more about how the backend works, there's this great blog post that goes through all the details of what's changed. [music] But the main thing to know is that VLM TPU is now powered by a unified backend. Not only is it faster than the older versions, but it offers a broader range of model support and features. So to gain access to a TPU, you first need to log into your Google account and go to the Google Cloud Console. Now, a great thing to note is that for a new account, you get $300 in free credit, which is just perfect for testing out and seeing how VLM on TPU works. Before you continue on, you will have to add some billing information. Now that we're in the Google Cloud Console, we can create a new project name to track this project of running VLM. Now that I'm in my new project, I can go and request for a TPU virtual machine. We can find that under the navigation menu, compute engine, TPUs. We'll need to enable this API before we can go forward. Now that I'm in my new Google Cloud Projects dashboard, we can see the TPU service is now enabled. So going back to the TPUs and now we can request a VM with TPUs. Let's give it a name. Then we need to choose the zone that our TPU is in. If you have a preference for zones, then you might want to choose a specific one. Otherwise, if you're interested in a specific TPU type, you can look in the Google Cloud Docs to see which TPUs are available in which zones. For this example, we want to use one of the V6 type TPUs. But it's important to note that availability may vary by zone. So, if you're not concerned about a specific zone, you can try other ones to see if the resources are available. I'm going to start with this Europe West 4. For our example, we're going to use a V6E chip. And I'm going to choose -1, which means we only want a single TPU chip. And I'm going to choose this V2 alpha TPU software version to be running on this chip. To help you in knowing what to choose for specific type of models, we do have a recipes repo which have recipes for exactly how to set up your system for different types of models. We're going to be serving the Llama 3.1 8 billion parameter model. And we can see here in the recipe it says one chip is sufficient for this model type. Okay. So let's go ahead and create this TPU. Alternatively to the web interface you can use the G-Cloud CLI. So here I have in my terminal G-Cloud is installed and I'm going to link this to my account by doing G-Cloud off login and selecting this account that I'm using to run this demo. Now, we want to set up the G-Cloud CLI to work with our correct project. And we can find the full project ID here on the Google Cloud project page. I'm going to copy this. If we prefer to request that TPU using the Google Cloud console, we could use this full command that we have here in our recipe. We just need to be sure that we set our TPU name, the zone, and the name of our full project. But we can see that our initial request here did go through. And we do have this TPU that now appears at with a status of being ready. We can also check the status of our TPUs by running the G-Cloud command. Here we can see that our VLM TPU is ready to go. So next, what we want to do is SSH into this TPU machine from our console. First, I'm going to set variables for our TPU name, our zone, our project name, and now this command will let us SSH into the machine. Now, we're into the TPU virtual machine. Before we get started, I'm going to pip install TPU info, which will let us display some information about the TPU so we can see it's there. Now, running TPU info, we can see that this machine indeed does have one V6E chip. I'm going to show you a few different ways to run VLM TPU on this machine. The first one is by far the simplest, and that's just using the pre-built Docker container for running VLM TPU. Now, we're going to use the commands from this recipe to start up the Docker image that has VLM TPU installed. First, setting this Docker URI and then running the Docker container. Now, we're going to change this size to 17 GB because we're running a smaller LLM. And also note here that we're port forwarding port 8000, which is what the VLM server will be running on. This might take a little bit of time as Docker downloads these different parts of the container. Now, we can see that we're inside this Docker image that has the latest version of VLM TPU and we can serve the model. But before we do so, we need to set up a few things. a hugging face home directory where our model will be downloaded and also our hugging face token. Because we're running a model that requires us to agree to terms and conditions, we need to make sure that we have agreed to the llama 3.1 terms and conditions in hugging face and then put our token here that we can find under our hugging face account. This is just a temporary token that I created for this demo and I'll delete afterwards. Now we can run this VLM serve command which will start up a VLM server with this model that we've set. We've made sure to tell it only has one TPU and we've set some other parameters here. Now immediately here we can see it's detected the TPU. The next thing it's going to do is download all those model weights off of hugging face. Then once the model weights are downloaded, it's going to go through some compile steps. And after a little while, we can now see that our server is up and running. It says that application startup is complete. There is a server running on that port 8000 that now we can access and query this model. Now I'm going to leave this running and I'm going to start a new terminal session where we'll connect to this VM. So in this new terminal window, I'm going to reconnect to the same VM. And now that we're in this VM, we can query against this model that we're hosted. And we can quickly test this by just running a curl command against that port 8000. We're calling the correct model and we're giving it a prompt. You can see here we got the full response from the model. If we want to run a more comprehensive benchmarking, we can use the VLM benchmarking tool. So let's just go back into this VLM Docker instance. Now from inside this Docker container, we can run this benchmark serve command which will run a bunch of data through this endpoint and then provide us some statistics in the end. As this is running, we can switch back over to our server that's running. You could see that server has received a bunch of different requests that are coming through. And now that this benchmark is complete, we have a bunch of different statistics about the run results. But we can also run VLM on TPU by installing it as a pit package. This part isn't necessary, but I'm going to install UV, which is helpful for maintaining different Python environments. I'm going to make a folder for our working directory. CD into that. Now, we're going to use UV to create a virtual environment with Python 3.12. It's important to run that version of Python because older versions may not work with VLM TPU. Then it's as simple as installing using pip install vlmtpu. And now we can run the same command as before just using uv run now. And we'll run this instead of in the docker container. This time we'll be running it directly on the TPU VM. And after a little while we see that the application startup is complete just like it was within the docker image. Once we're done using our TPU we want to make sure that we turn it off. And we can do that in the console or using the command line tool. I'm just going to stop it here in the console. If you're interested in learning more about TPU inference or contributing, check out our GitHub repo. We also have all of our documentation which includes instructions for installation and links to these recipes that I mentioned before. Thanks so much for watching this tutorial and we hope you have a good time running VLM on TPUs.