Running Inference on Local LLMs with vLLM
This guide explains how to run inference on local LLMs with vLLM using GPUs on DCC. Unlike Ollama, it does not require launching a separate inference server and can be run directly as a Python application.
The code for this example is available at: https://github.com/DukeRC/code/tree/main/Running-Inference-on-Local-LLMs-with-vLLM
As large language models (LLMs) continue to grow in popularity, efficient and scalable inference tools are becoming increasingly important in HPC environments. While Ollama is a strong option for smaller-scale inference, vLLM is better suited for larger-scale HPC workloads allowing for higher throughput, lower latency, and more concurrent requests. For a detailed performance comparison, see Ollama vs. vLLM: A deep dive into performance benchmarking.
This guide walks you through setting up and running inference on local LLMs with vLLM on DCC GPUs.
Initial setup
Assuming you have a Python environment available (e.g., Miniconda, Miniforge), first create a new environment and activate it.
Then install vLLM,
Additionally, install the dotenv package to manage environment variables,
Running Inference with vLLM
We will run inference on a single NVIDIA RTX 5000 Ada Generation GPU using the Qwen/Qwen2-7B-Instruct model from Hugging Face. We will first need to set up some environment variables to ensure that vLLM can access the model directory and connect to Hugging Face to download it if it is not already cached. The working directory, token file, and cache location below are set up once; the scripts are what you run each session.
First, create a directory in your work directory for this example,
Within this directory, create a .env file to store the environment variables,
echo 'HF_TOKEN=your_hugging_face_api_token' > /work/${USER}/vLLM/.env
chmod 600 /work/${USER}/vLLM/.env
The 600 permission ensures that only the owner can read and write the file, keeping your Hugging Face API token secure.
Next, add the following to your ~/.bashrc to keep vLLM's caches off your home directory,
export HF_HOME="/work/${USER}/.huggingface"
export FLASHINFER_WORKSPACE_BASE="/work/${USER}"
export VLLM_CACHE_ROOT="/work/${USER}/.cache/vllm"
HF_HOME covers downloaded model weights only. vLLM separately JIT-compiles FlashInfer kernels into $HOME/.cache/flashinfer and writes a torch.compile cache to $HOME/.cache/vllm, so all three need to point at /work to avoid filling your home quota; the symptom otherwise is OSError: [Errno 28] No space left on device. Note that FLASHINFER_WORKSPACE_BASE is a base directory, FlashInfer appends .cache/flashinfer to it itself.
Save the following Python script in the working directory as vllm_local.py,
#!/usr/bin/env python
#
# An example of using vLLM with a local model.
#
# Usage:
# 1. Set HF_HOME to control where Hugging Face cache is stored.
# export HF_HOME=/path/to/hf_cache
# 2. Set the HF_TOKEN environmental variable in a .env file in the root of the directory.
# 3. Run the script:
# ./vllm_local.py
import os
from dotenv import load_dotenv
os.environ["VLLM_CONFIGURE_LOGGING"] = "0"
from vllm import LLM, SamplingParams
# Model configuration
MODEL = "Qwen/Qwen2-7B-Instruct"
SAMPLING_PARAMS = SamplingParams(temperature=0.1, top_p=0.95, max_tokens=256)
# Array of prompts
PROMPTS = ["Tell me about Sri Lanka", "Why is the sky blue?", "Write a Python code that calculates the Fibonacci sequence up to 15."]
def main():
# Connect to Hugging Face with HF_TOKEN from .env
load_dotenv()
# Launch LLM
llm = LLM(model=MODEL, trust_remote_code=True)
outputs = llm.generate(PROMPTS, SAMPLING_PARAMS)
print("-" * 60)
print(f"Model: {MODEL}")
for output in outputs:
prompt = output.prompt
generated_text = output.outputs[0].text
print(f"Prompt: {prompt!r}")
print("Output:")
print(generated_text.strip())
print("-" * 60)
if __name__ == "__main__":
main()
Feel free to change the prompts in the PROMPTS array as you like. To run a single prompt, simply keep one entry in the array.
Running interactively
Request a slurm interactive session on an NVIDIA RTX 5000 Ada Generation GPU,
Once your allocation is ready, activate the conda environment we created earlier and run the script,
You should see the generated outputs for each prompt in the terminal.
Running as a batch job
To run the same script unattended, save the following as vllm_job.sh in the working directory,
#!/bin/bash
#SBATCH -J vllm-local # Job name
#SBATCH -p gpu-common # Partition name
#SBATCH -A rescomp # Account
#SBATCH --gres=gpu:5000_ada:1 # One RTX 5000 Ada GPU
#SBATCH --mem=100G # Memory
#SBATCH -t 1:00:00 # Time limit hrs:min:sec
#SBATCH -o vllm-%j.out # Standard output and error log
cd $SLURM_SUBMIT_DIR
# conda activate is undefined in a batch shell until conda.sh has been sourced
source "$(conda info --base)/etc/profile.d/conda.sh"
conda activate vllm-env
# Batch jobs do not read ~/.bashrc, so set the cache locations here
export HF_HOME="/work/${USER}/.huggingface"
export FLASHINFER_WORKSPACE_BASE="/work/${USER}"
export VLLM_CACHE_ROOT="/work/${USER}/.cache/vllm"
./vllm_local.py
and submit it with,
See the Slurm documentation for more on batch jobs.
Using multiple GPUs
To run a larger model across multiple GPUs on the same node, request more GPUs (e.g., gpu:5000_ada:2) and pass tensor_parallel_size to match,
Interactive chat session
If you would like a chat session instead of hard-coded prompts, use the following script,
#!/usr/bin/env python
#
# An example of using vLLM with a local model for interactive chat.
#
# Usage:
# 1. Set HF_HOME to control where Hugging Face cache is stored.
# export HF_HOME=/path/to/hf_cache
# 2. Set the HF_TOKEN environmental variable in a .env file in the root of the directory.
# 3. Run the script:
# ./vllm_local_chat.py
import os
from dotenv import load_dotenv
os.environ["VLLM_CONFIGURE_LOGGING"] = "0"
from vllm import LLM, SamplingParams
# Model configuration
MODEL = "Qwen/Qwen2-7B-Instruct"
SAMPLING_PARAMS = SamplingParams(temperature=0.1, top_p=0.95, max_tokens=256)
SYSTEM_PROMPT = "You are a helpful assistant."
def main() -> None:
# Connect to Hugging Face with HF_TOKEN from .env
load_dotenv()
llm = LLM(model=MODEL, trust_remote_code=True)
conversation = [{"role": "system", "content": SYSTEM_PROMPT}]
print(f"Model: {MODEL}")
print(f"Enter {sorted({"exit", "quit", "q"})} to exit.")
print("-" * 60)
while True:
try:
user_text = input("You: ").strip()
except (EOFError, KeyboardInterrupt):
print("\nExiting chat.")
break
if not user_text:
continue
if user_text.lower() in {"exit", "quit", "q"}:
print("Exiting chat.")
break
conversation.append({"role": "user", "content": user_text})
outputs = llm.chat(
conversation, sampling_params=SAMPLING_PARAMS, use_tqdm=False
)
assistant_text = outputs[0].outputs[0].text.strip()
print(f"Assistant: {assistant_text}\n")
conversation.append({"role": "assistant", "content": assistant_text})
if __name__ == "__main__":
main()
and run it with,