v6e TPU에서 vLLM 추론

이 튜토리얼에서는 v6e TPU에서 vLLM 추론을 실행하는 방법을 보여줍니다. 또한 Meta Llama-3.1 8B 모델의 벤치마크 스크립트를 실행하는 방법도 보여줍니다.

v6e TPU에서 vLLM을 시작하려면 vLLM 빠른 시작을 참고하세요.

GKE를 사용하는 경우 GKE 튜토리얼도 참고하세요.

시작하기 전에

HuggingFace 저장소에서 Llama3 모델 제품군을 사용하려면 동의 계약에 서명해야 합니다. https://huggingface.co/meta-llama/Llama-3.1-8B로 이동하여 동의 계약을 작성하고 승인될 때까지 기다립니다.

칩 4개가 있는 TPU v6e를 프로비저닝할 준비를 합니다.

  1. Google 계정에 로그인합니다. 아직 계정이 없다면 새 계정을 만듭니다.
  2. Google Cloud 콘솔의 프로젝트 선택기 페이지에서 Google Cloud 프로젝트를 선택하거나 만듭니다.
  3. Google Cloud 프로젝트에 결제를 사용 설정합니다. Google Cloud를 사용하려면 결제가 필요합니다.
  4. gcloud alpha 구성요소를 설치합니다.

  5. 다음 명령어를 실행하여 최신 버전의 gcloud 구성요소를 설치합니다.

    gcloud components update

  6. Cloud Shell을 사용하여 다음 gcloud 명령어를 통해 TPU API를 사용 설정합니다. Google Cloud 콘솔에서도 사용 설정할 수 있습니다.

    gcloud services enable tpu.googleapis.com

  7. TPU VM의 서비스 ID를 만듭니다.

    gcloud alpha compute tpus tpu-vm service-identity create --zone=ZONE

  8. TPU 서비스 계정을 만들고 Google Cloud 서비스에 대한 액세스 권한을 부여합니다.

    서비스 계정을 사용하면 Google Cloud TPU 서비스가 다른 Google Cloud 서비스에 액세스할 수 있습니다. 사용자 관리형 서비스 계정이 권장됩니다. 다음 가이드를 따라 역할을 만들고 부여합니다. 다음 역할이 필요합니다.

    • TPU 관리자: TPU를 만드는 데 필요합니다.
    • 스토리지 관리자: Cloud Storage에 액세스하는 데 필요합니다.
    • 로그 작성자: Logging API를 사용하여 로그를 작성하는 데 필요
    • 모니터링 측정항목 작성자: Cloud Monitoring에 측정항목을 쓰는 데 필요

  9. Google Cloud로 인증하고 Google Cloud CLI의 기본 프로젝트와 영역을 구성합니다.

    gcloud auth login
gcloud config set project PROJECT_ID
gcloud config set compute/zone ZONE

용량 확보

Cloud TPU 영업팀 또는 계정팀에 문의하여 TPU 할당량을 요청하고 용량에 관해 문의하세요.

Cloud TPU 환경 프로비저닝

v6e TPU는 GKE, GKE 및 XPK 또는 대기열에 추가된 리소스로 프로비저닝할 수 있습니다.

기본 요건

  • Google Cloud 프로젝트 내에서 액세스할 수 있는 최대 칩 수를 지정하는 TPUS_PER_TPU_FAMILY 할당량이 프로젝트에 충분한지 확인합니다.
  • 이 튜토리얼은 다음 구성으로 테스트되었습니다.
    • Python 3.10 or later
    • Nightly 소프트웨어 버전:
      • nightly JAX 0.4.32.dev20240912
      • nightly LibTPU 0.1.dev20240912+nightly
    • 안정화 소프트웨어 버전:
      • JAX + v0.4.35의 JAX 라이브러리
  • 프로젝트에 다음에 충분한 TPU 할당량이 있는지 확인합니다.
    • TPU VM 할당량
    • IP 주소 할당량
    • 하이퍼디스크 균형 할당량
  • 사용자 프로젝트 권한

TPU v6e 프로비저닝

   gcloud alpha compute tpus queued-resources create QUEUED_RESOURCE_ID \
      --node-id TPU_NAME \
       --project PROJECT_ID \
       --zone ZONE \
       --accelerator-type v6e-4 \
       --runtime-version v2-alpha-tpuv6e \
       --service-account SERVICE_ACCOUNT

명령어 플래그 설명

변수 설명
NODE_ID 큐에 추가된 리소스 요청이 할당될 때 생성되는 TPU의 사용자 할당 ID입니다.
PROJECT_ID Google Cloud 프로젝트 이름. 기존 프로젝트를 사용하거나 새 프로젝트를 만듭니다.>
ZONE 지원되는 영역에 대해서는 TPU 리전 및 영역 문서를 참고하세요.
ACCELERATOR_TYPE 지원되는 가속기 유형은 가속기 유형 문서를 참고하세요.
RUNTIME_VERSION v2-alpha-tpuv6e
SERVICE_ACCOUNT Google Cloud 콘솔 -> IAM -> 서비스 계정에서 찾을 수 있는 서비스 계정의 이메일 주소입니다.

예: tpu-service-account@<your_project_ID>.iam.gserviceaccount.com.com

list 또는 describe 명령어를 사용하여 큐에 추가된 리소스의 상태를 쿼리합니다.

   gcloud alpha compute tpus queued-resources describe ${QUEUED_RESOURCE_ID}  \
      --project ${PROJECT_ID} --zone ${ZONE}

대기열에 추가된 리소스 요청 상태의 전체 목록은 대기열에 추가된 리소스 문서를 참고하세요.

SSH를 사용하여 TPU에 연결

  gcloud compute tpus tpu-vm ssh TPU_NAME

종속 항목 설치

  1. Miniconda의 디렉터리를 만듭니다.

    mkdir -p ~/miniconda3

  2. Miniconda 설치 프로그램 스크립트를 다운로드합니다.

    wget https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh -O ~/miniconda3/miniconda.sh

  3. Miniconda를 설치합니다.

    bash ~/miniconda3/miniconda.sh -b -u -p ~/miniconda3

  4. Miniconda 설치 프로그램 스크립트를 삭제합니다.

    rm -rf ~/miniconda3/miniconda.sh

  5. PATH 변수에 Miniconda를 추가합니다.

    export PATH="$HOME/miniconda3/bin:$PATH"

  6. ~/.bashrc를 새로고침하여 PATH 변수에 변경사항을 적용합니다.

    source ~/.bashrc

  7. Conda 환경을 만듭니다.

    conda create -n vllm python=3.10 -y
conda activate vllm

  8. vLLM 저장소를 클론하고 vLLM 디렉터리로 이동합니다.

    git clone https://github.com/vllm-project/vllm.git && cd vllm

  9. 기존 torch 및 torch-xla 패키지를 삭제합니다.

    pip uninstall torch torch-xla -y

  10. 다른 빌드 종속 항목을 설치합니다.

    pip install -r requirements-tpu.txt
VLLM_TARGET_DEVICE="tpu" python setup.py develop
sudo apt-get install libopenblas-base libopenmpi-dev libomp-dev

모델 액세스 권한 얻기

토큰을 아직 만들지 않았다면 새 Hugging Face 토큰을 생성합니다.

  1. 내 프로필 > 설정 > 액세스 토큰을 클릭합니다.
  2. 새 토큰을 선택합니다.
  3. 원하는 이름과 Read 이상의 권한이 있는 역할을 지정합니다.
  4. 토큰 생성을 선택합니다.

  5. 생성된 토큰을 클립보드에 복사하고 환경 변수로 설정한 다음 huggingface-cli로 인증합니다.

    export TOKEN=YOUR_TOKEN
git config --global credential.helper store
huggingface-cli login --token $TOKEN

벤치마킹 데이터 다운로드

  1. /data 디렉터리를 만들고 Hugging Face에서 ShareGPT 데이터 세트를 다운로드합니다.

    mkdir ~/data && cd ~/data
wget https://huggingface.co/datasets/anon8231489123/ShareGPT_Vicuna_unfiltered/resolve/main/ShareGPT_V3_unfiltered_cleaned_split.json

vLLM 서버 시작

다음 명령어는 Hugging Face Model Hub에서 TPU VM의 /tmp 디렉터리로 모델 가중치를 다운로드하고, 다양한 입력 도형을 사전 컴파일하고, 모델 컴파일을 ~/.cache/vllm/xla_cache에 씁니다.

자세한 내용은 vLLM 문서를 참고하세요.

   cd ~/vllm
   vllm serve "meta-llama/Meta-Llama-3.1-8B" --download_dir /tmp --num-scheduler-steps 4 --swap-space 16 --disable-log-requests --tensor_parallel_size=4 --max-model-len=2048 &> serve.log &

vLLM 벤치마크 실행

vLLM 벤치마킹 스크립트를 실행합니다.

   python benchmarks/benchmark_serving.py \
       --backend vllm \
       --model "meta-llama/Meta-Llama-3.1-8B"  \
       --dataset-name sharegpt \
       --dataset-path ~/data/ShareGPT_V3_unfiltered_cleaned_split.json  \
       --num-prompts 1000

삭제

TPU를 삭제합니다.

gcloud compute tpus queued-resources delete QUEUED_RESOURCE_ID \
    --project PROJECT_ID \
    --zone ZONE \
    --force \
    --async