Serve Qwen2-7B-Instruct with vLLM on TPUs

This tutorial serves the Qwen/Qwen2-7B-Instruct model using the vLLM TPU serving framework on a v6e TPU VM.

Objectives

1. Set up your environment.
2. Run vLLM with Qwen2-7B-Instruct.
3. Send an inference request.
4. Run a benchmark workload.
5. Clean up.

Costs

This tutorial uses billable components of Google Cloud, including:

• Cloud TPU
• Compute Engine

To generate a cost estimate based on your projected usage, use the pricing calculator.

Before you begin

Before going through this tutorial, follow the instructions in the Set up the Cloud TPU environment page. The instructions guide you through the steps needed to create a Google Cloud project and configure it to use Cloud TPU. You may also use an existing Google Cloud project. If you choose to do so, you can skip the create a Google Cloud project step and start with Set up your environment to use Cloud TPU.

You need a Hugging Face access token to use this tutorial. You can sign up for a free account at Hugging Face. Once you have an account, generate an access token:

1. On the Welcome to Hugging Face page, click your account avatar and select Access tokens.
2. On the Access Tokens page, click Create new token.
3. Select the Read token type and enter a name for your token.
4. Your access token is displayed. Save the token in a safe place.

Set up your environment

1. Create a Cloud TPU v6e VM using the queued resources API. For qwen2-7b-instruct, we recommend using a v6e-1 TPU.

   export PROJECT_ID=YOUR_PROJECT_ID
   export TPU_NAME=qwen2-7b-instruct-tutorial
   export ZONE=us-east5-a
   export QR_ID=qwen2-7b-instruct-qr
   
   gcloud alpha compute tpus queued-resources create $QR_ID \
    --node-id $TPU_NAME \
    --project $PROJECT_ID \
    --zone $ZONE \
    --accelerator-type v6e-1 \
    --runtime-version v2-alpha-tpuv6e
   

2. Check to make sure your TPU VM is ready.

   gcloud compute tpus queued-resources describe $QR_ID \
     --project $PROJECT_ID \
     --zone $ZONE
   

   Tip: Initially, the status of the queued resource request will be WAITING_FOR_RESOURCES. When your TPU VM has been created, the status will change to ACTIVE. This can take some time depending on resource availability. If you experience a very long wait time, you might consider trying a different zone. Wait until the status is ACTIVE before proceeding.

   For example, when the status is ACTIVE:

   name: projects/your-project-id/locations/your-zone/queuedResources/your-queued-resource-id
     state:
     state: ACTIVE
     tpu:
     nodeSpec:
     - node:
         acceleratorType: v6e-1
         bootDisk: {}
         networkConfig:
             enableExternalIps: true
         queuedResource: projects/your-project-number/locations/your-zone/queuedResources/your-queued-resource-id
         runtimeVersion: v2-alpha-tpuv6e
         schedulingConfig: {}
         serviceAccount: {}
         shieldedInstanceConfig: {}
         useTpuVm: true
         nodeId: your-node-id
         parent: projects/your-project-number/locations/your-zone
   

3. Connect to the TPU VM.

     gcloud compute tpus tpu-vm ssh $TPU_NAME \
       --project $PROJECT_ID \
       --zone $ZONE
   

Run vLLM with Qwen2-7B-instruct

1. Set your Hugging Face token.

     export HF_TOKEN="YOUR_HF_TOKEN"
   

2. Inside the TPU VM, run the vLLM Docker container in detached mode and start the vLLM server. This command uses a shared memory size of 10 GB.

   export DOCKER_URI="vllm/vllm-tpu:v0.18.0"
   export CONTAINER_NAME="${USER}-vllm"
   export MAX_MODEL_LEN=4096
   export TP=1 # number of chips
   
   sudo docker run -d --name "${CONTAINER_NAME}" \
       --privileged --net=host \
       -v /dev/shm:/dev/shm \
       --shm-size 10gb \
       -e "HF_HOME=/dev/shm" \
       -e "HF_TOKEN=${HF_TOKEN}" \
       -p 8000:8000 "${DOCKER_URI}" \
           vllm serve Qwen/Qwen2-7B-Instruct \
               --seed 42 \
               --gpu-memory-utilization 0.98 \
               --max-num-batched-tokens 1024 \
               --max-num-seqs 128 \
               --tensor-parallel-size $TP \
               --max-model-len $MAX_MODEL_LEN
   

3. Check the server logs to confirm it's running.

   sudo docker logs -f "${CONTAINER_NAME}"
   

   When the vLLM server is running you see an output that resembles the following. After the output displays, press CTRL+C to return to the terminal.

   (APIServer pid=7) INFO:     Started server process [7]
   (APIServer pid=7) INFO:     Waiting for application startup.
   (APIServer pid=7) INFO:     Application startup complete.
   

Send an inference request

Once the vLLM server is running, you can send requests to the API. For more information, see the vLLM API reference documentation.

1. Send a test request to the server using curl.

     sudo docker exec -ti "${CONTAINER_NAME}" \
       curl http://localhost:8000/v1/completions \
         -H "Content-Type: application/json" \
         -d '{
             "model": "Qwen/Qwen2-7B-Instruct",
             "prompt": "The future of AI is",
             "max_tokens": 200,
             "temperature": 0
           }'
   

The response is returned in JSON format.

Run a benchmark workload

You can run benchmarks against the running server from your second terminal.

1. Inside the container, install the datasets library.

   sudo docker exec -it "${CONTAINER_NAME}" pip install datasets
   

2. Inside the container, run the vllm bench serve command.

   sudo docker exec -it "${CONTAINER_NAME}" \
       vllm bench serve \
           --backend vllm \
           --model "Qwen/Qwen2-7B-Instruct"  \
           --dataset-name random \
           --num-prompts 1000 \
           --seed 100
   

The benchmark results look like the following:

============ Serving Benchmark Result ============
Successful requests:                     1000
Benchmark duration (s):                  45.35
Total input tokens:                      1024000
Total generated tokens:                  126848
Request throughput (req/s):              22.05
Output token throughput (tok/s):         2797.15
Peak output token throughput (tok/s):    4258.00
Peak concurrent requests:                1000.00
Total Token throughput (tok/s):          25377.57
---------------Time to First Token----------------
Mean TTFT (ms):                          21332.46
Median TTFT (ms):                        21330.37
P99 TTFT (ms):                           42436.47
-----Time per Output Token (excl. 1st token)------
Mean TPOT (ms):                          37.36
Median TPOT (ms):                        38.56
P99 TPOT (ms):                           38.69
---------------Inter-token Latency----------------
Mean ITL (ms):                           37.35
Median ITL (ms):                         38.55
P99 ITL (ms):                            39.43
==================================================

Clean up

To avoid incurring charges to your Google Cloud account for the resources used in this tutorial, either delete the project that contains the resources, or keep the project and delete the individual resources.

1. In your terminal, type exit to disconnect from the TPU VM.

Delete your resources

You can delete the project which will delete all resources or you can keep the project and delete the resources.

Delete your project

To delete your Google Cloud project and all associated resources run:

    gcloud projects delete $PROJECT_ID

Delete TPU resources

Delete your Cloud TPU resources. The following command deletes both the queued resource request and the TPU VM using the --force parameter.

  gcloud alpha compute tpus queued-resources delete $QR_ID \
    --project=$PROJECT_ID \
    --zone=$ZONE \
    --force

What's next

• Learn more about vLLM on Cloud TPU.
• Learn more about Cloud TPU.