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. Cloud TPU 環境を設定するガイドに沿って、Cloud TPU を使用するための適切なアクセス権があることを確認します。

  2. TPU VM のサービス ID を作成します。

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

  3. TPU サービス アカウントを作成し、 Google Cloud サービスへのアクセス権を付与します。

    サービス アカウントにより、 Google Cloud TPU サービスが他の Google Cloudサービスにアクセスできるようになります。ユーザー管理のサービス アカウントの使用をおすすめします。サービス アカウントは、Google Cloud コンソールまたは gcloud コマンドを使用して作成できます。

    gcloud コマンドライン ツールを使用してサービス アカウントを作成します。

    gcloud iam service-accounts create your-service-account-name \
--description="your-sa-description" \
--display-name="your-sa-display-name"
export SERVICE_ACCOUNT_NAME=your-service-account-name

    Google Cloud コンソールからサービス アカウントを作成します。

    1. Google Cloud コンソールで [サービス アカウント] ページに移動します。
    2. [サービス アカウントを作成] をクリックします。
    3. サービス アカウント名を入力します。
    4. （省略可）サービス アカウントの説明を入力します。
    5. [作成] をクリックして続行します。
    6. サービス アカウントに付与するロールを選択します。
    7. [続行] をクリックします。
    8. （省略可）サービス アカウントを管理できるユーザーまたはグループを指定します。
    9. [完了] をクリックして、サービス アカウントの作成を完了します。

    サービス アカウントを作成したら、次の手順でサービス アカウントのロールを付与します。

    次のロールが必要です。

    • TPU 管理者: TPU を作成するために必要です
    • ストレージ管理者: Cloud Storage にアクセスするために必要です。
    • ログ書き込み
    • モニタリング指標の書き込み: Cloud Monitoring に指標を書き込むために必要

    ユーザーに IAM ロールを割り当てるには、管理者から roles/resourcemanager.projectIamAdmin が付与されている必要があります。プロジェクト IAM 管理者 roles/resourcemanager.projectIamAdmin ロールを持つユーザーも、このロールを付与できます。

    次の gcloud コマンドを使用して、サービス アカウントのロールを追加します。

    gcloud projects add-iam-policy-binding ${PROJECT_ID} \
   --member serviceAccount:${SERVICE_ACCOUNT_NAME}@${PROJECT_ID}.iam.gserviceaccount.com \
   --role roles/tpu.admin
gcloud projects add-iam-policy-binding ${PROJECT_ID} \
   --member serviceAccount:${SERVICE_ACCOUNT_NAME}@${PROJECT_ID}.iam.gserviceaccount.com \
   --role roles/storage.admin
gcloud projects add-iam-policy-binding ${PROJECT_ID} \
   --member serviceAccount:${SERVICE_ACCOUNT_NAME}@${PROJECT_ID}.iam.gserviceaccount.com \
   --role roles/logging.logWriter
gcloud projects add-iam-policy-binding ${PROJECT_ID} \
   --member serviceAccount:${SERVICE_ACCOUNT_NAME}@${PROJECT_ID}.iam.gserviceaccount.com \
   --role roles/monitoring.metricWriter

    Google Cloud コンソールを使用してロールを割り当てることもできます。

    Google Cloud コンソールで、次のロールを選択します。

    1. サービス アカウントを選択し、[プリンシパルを追加] をクリックします。
    2. [新しいプリンシパル] フィールドに、サービス アカウントのメールアドレスを入力します。
    3. [ロールを選択] プルダウンで、ロール（ストレージ管理者など）を検索して選択します。
    4. [保存] をクリックします。

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

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

容量を確保する

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

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

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

前提条件

  • このチュートリアルは Python 3.10 以降でテストされています。
  • プロジェクトに十分な TPUS_PER_TPU_FAMILY 割り当てがあることを確認します。これは、Google Cloud プロジェクト内でアクセスできるチップの最大数を指定します。
  • プロジェクトに次の TPU 割り当てがあることを確認します。
    • TPU VM の割り当て
    • IP アドレスの割り当て
    • Hyperdisk Balanced の割り当て
  • ユーザー プロジェクトの権限

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. 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 setup.py develop
sudo apt-get install libopenblas-base libopenmpi-dev libomp-dev

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

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

  1. [Your Profile] > [Settings] > [Access Tokens] の順にクリックします。
  2. [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

ベンチマーク データをダウンロードする

  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