Step by Step Guide: vLLM TPU Pod Slice Inference

사전설정

gcloud


tssh () {
  gcloud alpha compute tpus tpu-vm ssh $1 --zone us-central2-b
}

tsshm () {
  gcloud compute tpus tpu-vm ssh "$1" --zone us-central2-b --command "$2"
}

tssha () {
  gcloud compute tpus tpu-vm ssh "$1" --zone us-central2-b --worker=all --command "$2"
}

위와 같은 커맨드를 설정하면

tssh v4-64 : node-id가 v4-64인 클러스터의 첫 노드에 SSH 접속

tsshm v4-64 command : node-id가 v4-64인 클러스터의 Master(1st) 노드에서 커맨드 실행

tssha v4-64 command : node-id가 v4-64인 클러스터의 모든 노드에서 다함께 커맨드 실행 (pdsh같이)

vLLM 최신 레포 clone


tssha v4-64 "git clone https://github.com/vllm-project/vllm"
# tssha v4-64 "git clone https://github.com/Beomi/vllm-tpu-test"

vLLM Docker(TPU) build


tssha v4-64 "cd vllm && mv Dockerfile.tpu Dockerfile && sudo docker build -t vllm ."
# tssha v4-64 "cd vllm-tpu-test && mv Dockerfile.tpu Dockerfile && sudo docker build -t vllm ."

vLLM Docker Ray Cluster Start

Master(1st) Node에서 ray 실행 — IP를 반드시 체크해 둘 것! (VLLM_HOST_IP 에 세팅해야 함)


tsshm v4-64 'sudo docker run -t -d -e HF_TOKEN=INPUT_YOUR_TOKEN -e VLLM_HOST_IP=마스터노드IP --privileged --net host --shm-size=16G -it vllm ray start --head --block'
# 예시: tsshm v4-64 'sudo docker run -t -d -e HF_TOKEN=INPUT_YOUR_TOKEN -e VLLM_HOST_IP=10.130.0.53 --privileged --net host --shm-size=16G -it vllm ray start --head --block'

Master Node의 IP 체크:

위에서 실행한 Master Node에 띄워진 docker image의 ID를 사용

Sub Node들에서 아래 명령어 실행 (VLLM_HOST_IP 넣어야 함)


tssha v4-64 'sudo docker ps --filter ancestor=vllm --format "{{.ID}}" | grep -q . || sudo docker run -t -d -e HF_TOKEN=YOUR_TOKEN --privileged --net host --shm-size=16G -it vllm ray start --address <FIRST_NODE_IP>:6379 --block'

# 예시:
# tssha v4-64 'sudo docker ps --filter ancestor=vllm --format "{{.ID}}" | grep -q . || sudo docker run -t -d -e VLLM_HOST_IP=10.130.0.53 -e HF_TOKEN=YOUR_TOKEN --privileged --net host --shm-size=16G -it vllm ray start --address 10.130.0.53:6379 --block'

vLLM Docker Ray Cluster 상태 체크

아래 명령어로 master node에 접속해서, 도커 잘 떠있는지 확인하기.


tssh v4-64

→ 현재 모델이 agitated_goldstine 이름으로 떠 있으니, 해당 도커에 접속


sudo docker exec -it agitated_goldstine bash

→ 그 상태에서 ray status 명령어를 통해서 모든 TPU가 정상적으로 Ray Cluster에 연결되었는지 체크 할 수 있다.

현재 띄운 예시가 TPUv4-64로 32cores 클러스터이기 때문에 현재 사용이 0.0/32.0 TPU로 나온다.

참고: `ray` 커맨드로 시작한 모든 도커 끄기


tssha v4-64 'sudo docker ps | grep ray | awk "{print \$NF}" | xargs sudo docker stop'

Llama-3.1-instruct 모델 띄우기 w/ TP=4

위에서와 같이 Master node의 도커 쉘 내로 들어간 상태에서 (ray status로 보이는 상태에서) 아래 명령어로 vllm을 띄운다.


vllm serve NousResearch/Meta-Llama-3-8B-Instruct --dtype auto --api-key token-abc123 --tensor-parallel-size 4 --pipeline-parallel-size 8 --device tpu --distributed_executor_backend ray

문제: 자원을 Ray Cluster에서 받아오지 않음.

문제: Zero division error 발생 (Local Rank = 0)


group: <torch.distributed.distributed_c10d.ProcessGroup object at 0x7f49fc498ab0>
global_rank: 0
global_world_size: 4
num_nodes: 8

혹시 spawn 방식으로 하면 될까? → 🙅‍♂️


    "VLLM_WORKER_MULTIPROC_METHOD":
    lambda: os.getenv("VLLM_WORKER_MULTIPROC_METHOD", "fork"),
--> VLLM_WORKER_MULTIPROC_METHOD 를 `spawn` 으로 바꿔보기?


tsshm v4-64 'sudo docker run -t -d -e HF_TOKEN=INPUT_YOUR_TOKEN -e VLLM_HOST_IP=10.130.0.53 -e VLLM_WORKER_MULTIPROC_METHOD=spawn --privileged --net host --shm-size=16G -it vllm ray start --head --block'


tssha v4-64 'sudo docker ps --filter ancestor=vllm --format "{{.ID}}" | grep -q . || sudo docker run -t -d -e VLLM_HOST_IP=10.130.0.53 -e HF_TOKEN=YOUR_TOKEN -e VLLM_WORKER_MULTIPROC_METHOD=spawn --privileged --net host --shm-size=16G -it vllm ray start --address 10.130.0.53:6379 --block'

Reference

https://docs.vllm.ai/en/latest/serving/distributed_serving.html#multi-node-inference-and-serving

https://docs.vllm.ai/en/latest/getting_started/tpu-installation.html

Step by Step Guide: vLLM TPU Pod Slice Inference