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.

conda create -n vllm-env python=3.13 -y
conda activate vllm-env

Then install vLLM,

pip install vllm

Additionally, install the dotenv package to manage environment variables,

pip install python-dotenv

Running Inference with vLLM

We will run inference on a single NVIDIA RTX 5000 Ada Generation GPU using the Qwen/Qwen2-7B 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.

First, create a directory in your work directory for this example,

mkdir -p /work/${USER}/vLLM

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 set the Hugging Face cache directory,

export HF_HOME="/work/${USER}/.huggingface"

Save the following Python script in the working directory as vllm_local.py,

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"
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.

Let's request a slurm interactive session on an NVIDIA RTX 5000 Ada Generation GPU. You may also submit the script as a slurm batch job if you prefer.

srun -p gpu-common -A rescomp --gres=gpu:5000_ada:1 -t 1:00:00 --mem=100G --pty bash -i

Once your allocation is ready, activate the conda environment we created earlier and run the script,

conda activate vllm-env
./vllm_local.py

You should see the generated outputs for each prompt in the terminal.

If you would like a chat session instead of hard-coded prompts, use the following script,

vllm_local_chat.py

#!/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"
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()