Skip to content

Inference with vLLM on Polaris

vLLM is an open-source library designed to optimize the inference and serving. Originally developed at UC Berkeley's Sky Computing Lab, it has evolved into a community-driven project. The library is built around the innovative PagedAttention algorithm, which significantly improves memory management by reducing waste in Key-Value (KV) cache memory.

Provided Installation

vLLM (version 0.11.0) is available as part of the conda module. Please use the following commands:

module use /soft/modulefiles
module load conda
conda activate

Then, you can import vllm as follows 

>>> import vllm
>>> print(vllm.__version__)
'0.11.0rc2.dev147+g47b933954.d20251006'

Access Model Weights

Model weights for commonly used open-weight models are downloaded and available in the following directory on Polaris.

1
/eagle/datasets/model-weights/hub

To ensure your workflows utilize the preloaded model weights and datasets, update the following environment variables in your session. Some models hosted on Hugging Face may be gated, requiring additional authentication. To access these gated models, you will need a Hugging Face authentication token.

1
2
3
4
5
6
export HF_HOME="/eagle/datasets/model-weights"
export HF_DATASETS_CACHE="/eagle/datasets/model-weights"
export HF_MODULES_CACHE="/eagle/datasets/model-weights"
export HF_TOKEN="YOUR_HF_TOKEN"
export RAY_TMPDIR="/tmp"
export TMPDIR="/tmp"

Serving Small Models on a Single GPU

For small models that fit within the memory of a single A100 GPU (40 GB), no additional configuration is required to serve the model. Simply use the default tensor parallelism size (TP) of 1 when serving the model. This ensures the model is run on a single GPU without the need for distributed setup. Models with fewer than 15 billion parameters typically fit within a single GPU when using half precision (e.g., bfloat16).

For example, the following command serves meta-llama/Llama-3.1-8B-Instruct on a single GPU of a single node.

1
vllm serve meta-llama/Llama-3.1-8B-Instruct --port 8000 --dtype bfloat16

After the server output says Application startup complete., it is ready to accept prompt requests, for example with the following simple Python code (can background the server process)

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
from openai import OpenAI

openai_api_base = f"http://localhost:8000/v1"
client = OpenAI(
    api_key="EMPTY",
    base_url=openai_api_base,
)

prompt = "Hi, can you introduce yourself?"
response = client.chat.completions.create(
    model="meta-llama/Llama-3.1-8B-Instruct",
    messages=[{"role": "user", "content": prompt}],
    temperature=0.0,
    max_tokens=1024,
    stream=False
)
print(f"\n{response.choices[0].message.content}\n")

Serving Medium Models on Multiple GPUs (Single Node)

To serve larger models which require multiple GPUs (TP>1) but still only a single node, a more advanced setup is necessary. This involves setting the VLLM_HOST_IP and the TP size. Models with about 70 billion parameters can usually fit within a single node utilizing half precition.

The following commands demonstrate how to serve the meta-llama/Llama-3.1-70B-Instruct on 4 GPUs on a single node. 

1
2
3
4
export VLLM_HOST_IP=$(getent hosts $(hostname).hsn.cm.polaris.alcf.anl.gov | awk '{ print $1 }' | tr ' ' '\n' | sort | head -n 1)
export no_proxy="localhost,127.0.0.1"

vllm serve meta-llama/Llama-3.1-70B-Instruct --port 8000 --tensor-parallel-size 4 --dtype bfloat16 --trust-remote-code --max-model-len 8192

Serving Large Models on Multiple GPUs and Nodes

To serve models across multiple nodes, we suggest users take advantage of the provided setup_ray_cluster.sh script to setup a Ray cluster across nodes before running vllm serve.

Setup script
setup_ray_cluster.sh
  1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
# This script is intended to be sourced, not executed directly.
# On the head node, run:
#   source ./setup_ray_cluster.sh
#   main

########################################################################
# FUNCTIONS
########################################################################

# Setup environment and variables needed to setup ray and vllm
setup_environment() {
    echo "[$(hostname)] Setting up the environment..."

    # Set proxy configurations
    export HTTP_PROXY="http://proxy.alcf.anl.gov:3128"
    export HTTPS_PROXY="http://proxy.alcf.anl.gov:3128"
    export http_proxy="http://proxy.alcf.anl.gov:3128"
    export https_proxy="http://proxy.alcf.anl.gov:3128"
    export ftp_proxy="http://proxy.alcf.anl.gov:3128"

    # Define the common setup script path (make sure this file is accessible on all nodes)
    SCRIPT_DIR=$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd)
    SCRIPT_PATH="${SCRIPT_DIR}/$(basename "${BASH_SOURCE[0]}")"
    export COMMON_SETUP_SCRIPT=$SCRIPT_PATH

    # Load modules and activate your conda environment
    module use /soft/modulefiles
    module load conda
    conda activate

    export CUDA_VISIBLE_DEVICES=0,1,2,3
    export RAY_EXPERIMENTAL_NOSET_CUDA_VISIBLE_DEVICES=1
    export OMP_NUM_THREADS=8
    export TORCH_LLM_ALLREDUCE=1
    export HF_HOME="/eagle/datasets/model-weights"
    export HF_DATASETS_CACHE="/flare/datasets/model-weights"
    export TMPDIR="/tmp"
    export RAY_TMPDIR="/tmp"
    export VLLM_IMAGE_FETCH_TIMEOUT=60

    ulimit -c unlimited

    # Derive the node's HSN IP address (modify the getent command as needed)
    export HSN_IP_ADDRESS=$(getent hosts "$(hostname).hsn.cm.polaris.alcf.anl.gov" | awk '{ print $1 }' | sort | head -n 1)
    export VLLM_HOST_IP="$HSN_IP_ADDRESS"

    echo "[$(hostname)] Environment setup complete. HSN_IP_ADDRESS is $HSN_IP_ADDRESS"
}

# Stop any running Ray processes
stop_ray() {
    echo "[$(hostname)] Stopping Ray (if running)..."
    ray stop -f
}

# Start Ray head node
start_ray_head() {
    echo "[$(hostname)] Starting Ray head..."
    ray start --num-gpus=4 --num-cpus=32 --head --node-ip-address="$HSN_IP_ADDRESS" --temp-dir=/tmp

    # Wait until Ray reports that the head node is up
    echo "[$(hostname)] Waiting for Ray head to be up..."
    until ray status &>/dev/null; do
        sleep 5
        echo "[$(hostname)] Waiting for Ray head..."
    done
    echo "[$(hostname)] ray status: $(ray status)"
    echo "[$(hostname)] Ray head node is up."
}

# Start Ray worker node
start_ray_worker() {
    echo "[$(hostname)] Starting Ray worker, connecting to head at $RAY_HEAD_IP..."
    echo "HSN IP Address : $HSN_IP_ADDRESS"
    ray start --num-gpus=4 --num-cpus=32 --address="$RAY_HEAD_IP:6379" --node-ip-address="$HSN_IP_ADDRESS" --temp-dir=/tmp

    echo "[$(hostname)] Waiting for Ray worker to be up..."
    until ray status &>/dev/null; do
        sleep 5
        echo "[$(hostname)] Waiting for Ray worker..."
    done
    echo "[$(hostname)] ray status: $(ray status)"
    echo "[$(hostname)] Ray worker node is up."
}

########################################################################
# MAIN SCRIPT LOGIC
########################################################################

main() {
    # Require HF_TOKEN env variable
    if [[ -z "${HF_TOKEN}" ]]; then
        echo "Error: HF_TOKEN is not set!" >&2
        return 1
    fi

    # Ensure that the script is being run within a PBS job
    if [ -z "$PBS_NODEFILE" ]; then
        echo "Error: PBS_NODEFILE not set. This script must be run within a PBS job allocation."
        return 1
    fi

    # Read all nodes from the PBS_NODEFILE into an array.
    mapfile -t nodes_full < "$PBS_NODEFILE"
    num_nodes=${#nodes_full[@]}

    echo "Allocated nodes ($num_nodes):"
    printf " - %s\n" "${nodes_full[@]}"

    # Require at least 2 nodes (one head + one worker)
    if [ "$num_nodes" -lt 2 ]; then
        echo "Error: Need at least 2 nodes to launch the Ray cluster."
        return 1
    fi

    # The first node will be our Ray head.
    head_node_full="${nodes_full[0]}"

    # All remaining nodes will be the workers.
    worker_nodes_full=("${nodes_full[@]:1}")

    # It is a good idea to run this master script on the designated head node.
    current_node=$(hostname -f)
    echo "[$(hostname)] Running on head node."

    # --- Setup and start the head node ---
    setup_environment
    stop_ray
    start_ray_head

    # Export the head node's IP so that workers can join.
    export RAY_HEAD_IP="$HSN_IP_ADDRESS"
    export RAY_ADDRESS="$RAY_HEAD_IP:6379"
    echo "[$(hostname)] RAY_HEAD_IP exported as $RAY_HEAD_IP"
    echo "[$(hostname)] RAY_ADDRESS exported as $RAY_ADDRESS"

    # --- Launch Ray workers on each of the other nodes via SSH ---
    for worker in "${worker_nodes_full[@]}"; do
        echo "[$(hostname)] Launching Ray worker on $worker..."
        ssh "$worker" "bash -l -c 'set -x; export RAY_HEAD_IP=${RAY_HEAD_IP}; export COMMON_SETUP_SCRIPT=${COMMON_SETUP_SCRIPT} ;source \$COMMON_SETUP_SCRIPT; setup_environment; stop_ray; start_ray_worker'" &
    done

    # Wait for all background SSH jobs to finish.
    wait

    echo "[$(hostname)] Ray cluster is up and running with $num_nodes nodes."
}

The following example serves meta-llama/Llama-3.1-70B-Instruct model using 2 nodes. It sets tensor parallelism TP=4 for intra-node communications and pipeline parallelism PP=2 for inter-node communication, for a total of 16 shards. 

1
2
3
4
source /path/to/setup_ray_cluster.sh
main
ray status # should show 2 active nodes and 8 GPUs total
vllm serve meta-llama/Llama-3.1-70B-Instruct --port 8000 --tensor-parallel-size 4 --pipeline-parallel-size 2 --dtype bfloat16 --trust-remote-code --max-model-len 8192

General guidelines to keep in mind when serving models across nodes:

Guidelines for vLLM Model Serving

  • Setting --max-model-len can be important in order to fit the model on the GPUs.
  • Tensor parallelism size must evenly divide the number of attention heads of the model. For example, the Llama-3.1-70B-Instruct model has 64 attention heads, so valid TP values are 1, 2, 4, 8. However, setting TP size equal to the number of GPUs on the node (4 for Polaris) is preferred to utilize intra-node communications for this parallism dimension.
  • Pipeline parallelism size must evenly divide the number of hidden layers in the model. For example, the Llama-3.1-70B-Instruct model has 80 layers, so PP values of 1, 2, 4, 5, etc. are valid. In this case PP is set to the number of nodes used.
  • The product TP x PP indicates the total number of GPUs used to serve the model, which is usually defined by the number of parameters in the model and the memory of the individual GPUs. As a back of the envelope calculation, when using half precision such as bfloat16, (num. billion paramemers x 2) / GPU GB mem gives the number of GPUs needed.

Scaling vLLM Workflows

To scale vLLM workflows on ALCF system there are a few recommended approaches depending on the user's needs and setup. These approaches are described in detail in the GettingStarted repository along with example scripts for each.