Run ML inference by using vLLM on GPUs

Run in Google Colab View source on GitHub

vLLM is a fast and user-friendly library for LLM inference and serving. vLLM optimizes LLM inference with mechanisms like PagedAttention for memory management and continuous batching for increasing throughput. For popular models, vLLM has been shown to increase throughput by a multiple of 2 to 4. With Apache Beam, you can serve models with vLLM and scale that serving with just a few lines of code.

This notebook demonstrates how to run machine learning inference by using vLLM and GPUs in three ways:

  • locally without Apache Beam
  • locally with the Apache Beam local runner
  • remotely with the Dataflow runner

It also shows how to swap in a different model without modifying your pipeline structure by changing the configuration.

Requirements

This notebook assumes that a GPU is enabled in Colab. If this setting isn't enabled, the locally executed sections of this notebook might not work. To enable a GPU, in the Colab menu, click Runtime > Change runtime type. For Hardware accelerator, choose a GPU accelerator. If you can't access a GPU in Colab, you can run the Dataflow section of this notebook.

To run the Dataflow section, you need access to the following resources:

Install dependencies

Before creating your pipeline, download and install the dependencies required to develop with Apache Beam and vLLM. vLLM is supported in Apache Beam versions 2.60.0 and later.

pip install openai>=1.52.2
pip install vllm>=0.6.3
pip install triton>=3.1.0
pip install apache-beam[gcp]==2.61.0
pip install nest_asyncio # only needed in colab
pip