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 アカウントにログインします。Google アカウントをまだお持ちでない場合は、新しいアカウントを登録します。
  2. Google Cloud コンソールで、プロジェクト セレクタ ページから Google Cloud プロジェクトを選択するか作成します。
  3. Google Cloud プロジェクトに対する課金を有効にします。Google Cloud の使用にはすべて課金が必要です。
  4. gcloud alpha components をインストールします。

  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

容量を確保する

TPU の割り当てをリクエストし、容量について質問するには、Cloud TPU の営業チームまたはアカウント チームにお問い合わせください。

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

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

前提条件

  • プロジェクトに十分な TPUS_PER_TPU_FAMILY 割り当てがあることを確認します。これは、Google Cloud プロジェクト内でアクセスできるチップの最大数を指定します。
  • このチュートリアルは、次の構成でテストされました。
    • Python 3.10 or later
    • ナイトリー ソフトウェアのバージョン:
      • 夜間 JAX 0.4.32.dev20240912
      • ナイトリー LibTPU 0.1.dev20240912+nightly
    • 安定版ソフトウェアのバージョン:
      • v0.4.35 の JAX + JAX Lib
  • プロジェクトに次の 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. 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. [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