v6e TPU での vLLM 推論

このチュートリアルでは、v6e TPU で vLLM 推論を実行する方法について説明します。また、Meta Llama-3.1-8B モデルのベンチマーク スクリプトを実行する方法も示します。

v6e TPU で vLLM を使用するには、vLLM クイックスタートをご覧ください。

GKE を使用している場合は、GKE チュートリアルもご覧ください。

始める前に

HuggingFace リポジトリの Llama3 ファミリー モデルを使用するには、同意契約に署名する必要があります。meta-llama/Llama-3.1-8B に移動し、同意契約に記入して、承認されるまで待ちます。

4 個のチップを搭載した TPU v6e をプロビジョニングする準備を行います。

1. Cloud TPU 環境を設定するガイドに沿って、 Google Cloud プロジェクトを設定して、Google Cloud CLI の構成を行い、Cloud TPU API を有効にして、Cloud TPU を使用するためのアクセス権があることを確認します。

2. Google Cloud で認証し、Google Cloud CLI のデフォルトのプロジェクトとゾーンを構成します。

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

容量を確保する

TPU 容量を確保する準備ができたら、Cloud TPU 割り当てで Cloud TPU 割り当ての詳細を確認してください。容量の確保についてご不明な点がございましたら、Cloud TPU のセールスチームまたはアカウント チームにお問い合わせください。

Cloud TPU 環境をプロビジョニングする

TPU VM は、GKE、GKE と XPK、またはキューに格納されたリソースとしてプロビジョニングできます。

前提条件

• プロジェクトに十分な TPUS_PER_TPU_FAMILY 割り当てがあることを確認します。これは、Google Cloud プロジェクト内でアクセスできるチップの最大数を指します。
• プロジェクトの次の TPU 割り当てが十分にあることを確認します。
  • TPU VM の割り当て
  • IP アドレスの割り当て
  • Hyperdisk Balanced の割り当て
• ユーザー プロジェクトの権限
  • GKE と XPK を使用している場合は、XPK の実行に必要な権限について、ユーザーまたはサービス アカウントに対する Google Cloud コンソールの権限をご覧ください。

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

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. Miniconda を PATH 変数に追加します。

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

6. ~/.bashrc を再読み込みして、PATH 変数に変更を適用します。

   source ~/.bashrc

7. Conda 環境を作成します。

   conda create -n vllm python=3.11 -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 -m pip install --editable .
    sudo apt-get install libopenblas-base libopenmpi-dev libomp-dev
    

モデルへのアクセス権を取得する

Hugging Face トークンをまだ生成していない場合は、新しいトークンを生成します。

1. [Your Profile] > [Settings] > [Access Tokens] の順に移動します。

2. [Create new token] を選択します。

3. 任意の名前と、少なくとも Read 権限を持つロールを指定します。

4. [Generate a token] を選択します。

5. 生成されたトークンをクリップボードにコピーし、環境変数として設定して、huggingface-cli で認証します。

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

vLLM サーバーを起動する

次のコマンドは、モデルの重みを Hugging Face Model Hub から TPU VM の /tmp ディレクトリにダウンロードし、一連の入力シェイプを事前コンパイルして、モデルのコンパイルを ~/.cache/vllm/xla_cache に書き込みます。

詳細については、vLLM のドキュメントをご覧ください。

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

vLLM ベンチマークを実行する

vLLM ベンチマーク スクリプトを実行します。

export MODEL="meta-llama/Llama-3.1-8B"
pip install pandas
pip install datasets
python benchmarks/benchmark_serving.py \
  --backend vllm \
  --model $MODEL  \
  --dataset-name random \
  --random-input-len 1820 \
  --random-output-len 128 \
  --random-prefix-len 0

クリーンアップ

TPU を削除します。

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