複数の GPU ノードを使用して LLM（Llama3.1 405B）をサービングする

概要

このチュートリアルでは、vLLM サービング フレームワークと LeaderWorkerSet（LWS）API を利用し、Google Kubernetes Engine（GKE）の複数のノードで画像処理装置（GPU）を使用して Llama 3.1 405b をサービングする方法について説明します。

このドキュメントは、AI / ML ワークロードをデプロイして提供する際に、マネージド Kubernetes での詳細な制御、スケーラビリティ、復元力、ポータビリティ、費用対効果が求められる場合の出発点として適しています。

LeaderWorkerSet（LWS）

LWS は、AI / ML マルチノード推論ワークロードの一般的なデプロイ パターンに対応する Kubernetes deployment API です。LWS を使用すると、複数の Pod をグループとして扱うことができます。

vLLM を使用したマルチホスト サービング

1 つの GPU ノードに収まらない非常に大きな言語モデルをデプロイする場合は、複数の GPU ノードを使用してモデルをサービングします。vLLM はテンソル並列処理とパイプライン並列処理の両方をサポートし、複数の GPU でワークロードを実行します。

テンソル並列処理では、Transformer レイヤの行列乗算を複数の GPU に分割します。ただし、この方法では GPU 間の通信が必要になるため、高速なネットワークが必要であり、ノード間でワークロードを実行する場合は適していません。

パイプライン並列処理では、モデルをレイヤ（垂直方向）で分割します。この方法では、GPU 間の通信を常に行う必要がないため、ノードをまたいでモデルを実行する場合に適しています。

マルチノード サービングでは、どちらの戦略も使用できます。たとえば、それぞれ 8 個の H100 GPU が割り当てられた 2 つのノードを使用する場合、双方向のパイプライン並列処理を使用して 2 つのノードにモデルをシャードし、8 方向テンソル並列処理を使用して各ノード上の 8 個の GPU にモデルをシャードできます。

目標

  1. GKE Standard クラスタを準備する。
  2. クラスタ内の複数のノードに vLLM をデプロイする。
  3. vLLM を使用して、curl を介して Llama3 405b モデルをサービングする。

始める前に

  • Sign in to your Google Cloud account. If you're new to Google Cloud, create an account to evaluate how our products perform in real-world scenarios. New customers also get $300 in free credits to run, test, and deploy workloads.

  • In the Google Cloud console, on the project selector page, select or create a Google Cloud project.

    Go to project selector

  • Make sure that billing is enabled for your Google Cloud project.

  • Enable the required API.

    Enable the API

    •

  • In the Google Cloud console, on the project selector page, select or create a Google Cloud project.

    Go to project selector

  • Make sure that billing is enabled for your Google Cloud project.

  • Enable the required API.

    Enable the API

    •

  • Make sure that you have the following role or roles on the project: roles/container.admin, roles/iam.serviceAccountAdmin

    Check for the roles

    1. In the Google Cloud console, go to the IAM page.

      Go to IAM
    2. Select the project.

    3. In the Principal column, find all rows that identify you or a group that you're included in. To learn which groups you're included in, contact your administrator.

    4. For all rows that specify or include you, check the Role colunn to see whether the list of roles includes the required roles.

    Grant the roles

    1. In the Google Cloud console, go to the IAM page.

      [IAM] に移動
    2. プロジェクトを選択します。
    3. [ アクセスを許可] をクリックします。

    4. [新しいプリンシパル] フィールドに、ユーザー ID を入力します。 これは通常、Google アカウントのメールアドレスです。

    5. [ロールを選択] リストでロールを選択します。
    6. 追加のロールを付与するには、 [別のロールを追加] をクリックして各ロールを追加します。
    7. [保存] をクリックします。

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

アクセス トークンを生成する

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

  1. [Your Profile] > [Settings] > [Access Tokens] の順にクリックします。
  2. [New Token] を選択します。
  3. 任意の名前と、少なくとも Read ロールを指定します。
  4. [Generate a token] を選択します。

環境を準備する

このチュートリアルでは、Cloud Shell を使用して Google Cloud でホストされるリソースを管理します。Cloud Shell には、このチュートリアルに必要な kubectl や gcloud CLI などのソフトウェアがプリインストールされています。

Cloud Shell を使用して環境を設定するには、次の操作を行います。

  1. Google Cloud コンソールで、Google Cloud コンソールの Cloud Shell 有効化アイコン [Cloud Shell をアクティブにする] をクリックして、Cloud Shell セッションを起動します。これにより、Google Cloud コンソールの下部ペインでセッションが起動されます。

  2. デフォルトの環境変数を設定します。

    gcloud config set project PROJECT_ID
export PROJECT_ID=$(gcloud config get project)
export CLUSTER_NAME=CLUSTER_NAME
export ZONE=ZONE
export HF_TOKEN=HUGGING_FACE_TOKEN
export IMAGE_NAME=IMAGE_NAME

    次の値を置き換えます。

    • PROJECT_ID: Google Cloud のプロジェクト ID。
    • CLUSTER_NAME: GKE クラスタの名前。
    • ZONE: H100 をサポートするゾーン。
    • IMAGE_NAME: レイ スクリプトを含む vLLM イメージ。us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-vllm-serve:20240821_1034_RC00 は Google が提供しますが、独自に作成することもできます。

GKE クラスタを作成する

GKE Autopilot クラスタまたは GKE Standard クラスタの複数の GPU ノードにまたがって vLLM を使用してモデルを提供できます。フルマネージドの Kubernetes エクスペリエンスを実現するには、Autopilot クラスタを使用することをおすすめします。ワークロードに最適な GKE の運用モードを選択するには、GKE の運用モードを選択するをご覧ください。

Autopilot

Cloud Shell で、次のコマンドを実行します。

  gcloud container clusters create-auto ${CLUSTER_NAME} \
    --project=${PROJECT_ID} \
    --region=${REGION} \
    --cluster-version=${CLUSTER_VERSION}

スタンダード

  1. 2 つの CPU ノードを含む GKE Standard クラスタを作成します。

    gcloud container clusters create CLUSTER_NAME \
    --project=PROJECT_ID \
    --num-nodes=2 \
    --location=ZONE \
    --machine-type=e2-standard-16

  2. 2 つのノードと 8 つの H100 で構成される A3 ノードプールを作成します。

    gcloud container node-pools create gpu-nodepool \
    --location=ZONE \
    --num-nodes=2 \
    --machine-type=a3-highgpu-8g \
    --accelerator=type=nvidia-h100-80gb,count=8,gpu-driver-version=LATEST \
    --placement-type=COMPACT \
    --cluster=CLUSTER_NAME

クラスタと通信を行うように kubectl を構成します。

gcloud container clusters get-credentials CLUSTER_NAME --location=ZONE

Hugging Face の認証情報用の Kubernetes Secret を作成する

Hugging Face トークンを含む Kubernetes Secret を作成します。

kubectl create secret generic hf-secret \
  --from-literal=hf_api_token=${HF_TOKEN} \
  --dry-run=client -o yaml | kubectl apply -f -

（省略可）独自の vLLM マルチノード イメージを作成する

Docker イメージの内容をより細かく制御し、スクリプトとともに特定の依存関係を含めるには、このオプションを選択します。複数のノードにまたがって vLLM を実行するには、ノード間の通信に Ray を使用できます。Dockerfile を表示できます。このファイルには、LeaderWorkerSet リポジトリで vLLM を使用して Ray を設定する bash スクリプトが含まれています。

コンテナをビルドする

  1. LeaderWorkerSet リポジトリのクローンを作成します。

    git clone https://github.com/kubernetes-sigs/lws.git

  2. イメージを構築する。

    cd lws/docs/examples/vllm/build/ && docker build -f Dockerfile . -t vllm-multihost

イメージを Artifact Registry に push する

Kubernetes Deployment がイメージにアクセスできるようにするには、Google Cloud プロジェクト内の Artifact Registry にイメージを保存します。

gcloud artifacts repositories create vllm-multihost --repository-format=docker --location=REGION_NAME && \
gcloud auth configure-docker REGION_NAME-docker.pkg.dev && \
docker image tag vllm-multihost REGION_NAME-docker.pkg.dev/PROJECT_ID/vllm-multihost/vllm-multihost:latest && \
docker push REGION_NAME-docker.pkg.dev/PROJECT_ID/vllm-multihost/vllm-multihost:latest

LeaderWorkerSet をインストールする

LWS をインストールするには、次のコマンドを実行します。

VERSION=v0.4.0
kubectl apply --server-side -f https://github.com/kubernetes-sigs/lws/releases/download/$VERSION/manifests.yaml

LeaderWorkerSet コントローラが lws-system Namespace で実行されていることを確認します。

kubectl get pod -n lws-system

出力は次のようになります。

NAME                                      READY   STATUS    RESTARTS   AGE
lws-controller-manager-5c4ff67cbd-9jsfc   2/2     Running   0          6d23h

vLLM モデルサーバーをデプロイする

vLLM モデルサーバーをデプロイする手順は次のとおりです。

  1. マニフェスト vllm-llama3-405b-A3.yaml を調べます。

    
apiVersion: leaderworkerset.x-k8s.io/v1
kind: LeaderWorkerSet
metadata:
  name: vllm
spec:
  replicas: 1
  leaderWorkerTemplate:
    size: 2
    restartPolicy: RecreateGroupOnPodRestart
    leaderTemplate:
      metadata:
        labels:
          role: leader
      spec:
        nodeSelector:
          cloud.google.com/gke-accelerator: nvidia-h100-80gb
        containers:
          - name: vllm-leader
            image: IMAGE_NAME
            env:
              - name: HUGGING_FACE_HUB_TOKEN
                valueFrom:
                  secretKeyRef:
                    name: hf-secret
                    key: hf_api_token
            command:
              - sh
              - -c
              - "/workspace/vllm/examples/ray_init.sh leader --ray_cluster_size=$(LWS_GROUP_SIZE); 
                python3 -m vllm.entrypoints.openai.api_server --port 8080 --model meta-llama/Meta-Llama-3.1-405B-Instruct --tensor-parallel-size 8 --pipeline-parallel-size 2"
            resources:
              limits:
                nvidia.com/gpu: "8"
            ports:
              - containerPort: 8080
            readinessProbe:
              tcpSocket:
                port: 8080
              initialDelaySeconds: 15
              periodSeconds: 10
            volumeMounts:
              - mountPath: /dev/shm
                name: dshm
        volumes:
        - name: dshm
          emptyDir:
            medium: Memory
            sizeLimit: 15Gi
    workerTemplate:
      spec:
        nodeSelector:
          cloud.google.com/gke-accelerator: nvidia-h100-80gb
        containers:
          - name: vllm-worker
            image: IMAGE_NAME
            command:
              - sh
              - -c
              - "/workspace/vllm/examples/ray_init.sh worker --ray_address=$(LWS_LEADER_ADDRESS)"
            resources:
              limits:
                nvidia.com/gpu: "8"
            env:
              - name: HUGGING_FACE_HUB_TOKEN
                valueFrom:
                  secretKeyRef:
                    name: hf-secret
                    key: hf_api_token
            volumeMounts:
              - mountPath: /dev/shm
                name: dshm   
        volumes:
        - name: dshm
          emptyDir:
            medium: Memory
            sizeLimit: 15Gi
---
apiVersion: v1
kind: Service
metadata:
  name: vllm-leader
spec:
  ports:
    - name: http
      port: 8080
      protocol: TCP
      targetPort: 8080
  selector:
    leaderworkerset.sigs.k8s.io/name: vllm
    role: leader
  type: ClusterIP

  2. 次のコマンドを実行してマニフェストを適用します。

    kubectl apply -f vllm-llama3-405b-A3.yaml

  3. 実行中のモデルサーバーのログを表示する

    kubectl logs vllm-0 -c vllm-leader

    出力は次のようになります。

    INFO 08-09 21:01:34 api_server.py:297] Route: /detokenize, Methods: POST
INFO 08-09 21:01:34 api_server.py:297] Route: /v1/models, Methods: GET
INFO 08-09 21:01:34 api_server.py:297] Route: /version, Methods: GET
INFO 08-09 21:01:34 api_server.py:297] Route: /v1/chat/completions, Methods: POST
INFO 08-09 21:01:34 api_server.py:297] Route: /v1/completions, Methods: POST
INFO 08-09 21:01:34 api_server.py:297] Route: /v1/embeddings, Methods: POST
INFO:     Started server process [7428]
INFO:     Waiting for application startup.
INFO:     Application startup complete.
INFO:     Uvicorn running on http://0.0.0.0:8080 (Press CTRL+C to quit)

モデルをサービングする

次のコマンドを実行して、モデルへのポート転送を設定します。

kubectl port-forward svc/vllm-leader 8080:8080

curl を使用してモデルを操作する

新しいターミナルで、サーバーにリクエストを送信します。

curl http://localhost:8080/v1/completions \
-H "Content-Type: application/json" \
-d '{
    "model": "meta-llama/Meta-Llama-3.1-405B-Instruct",
    "prompt": "San Francisco is a",
    "max_tokens": 7,
    "temperature": 0
}'

出力例を以下に示します。

{"id":"cmpl-0a2310f30ac3454aa7f2c5bb6a292e6c",
"object":"text_completion","created":1723238375,"model":"meta-llama/Meta-Llama-3.1-405B-Instruct","choices":[{"index":0,"text":" top destination for foodies, with","logprobs":null,"finish_reason":"length","stop_reason":null}],"usage":{"prompt_tokens":5,"total_tokens":12,"completion_tokens":7}}

（省略可）Hyperdisk ML を使用してモデルの読み込み時間を短縮する

vLLM で、新しいレプリカごとに Llama 3.1-405B をダウンロード、読み込み、ウォームアップするまでに最大 90 分かかることがあります。モデルを Hyperdisk ML に直接ダウンロードして各 Pod にマウントすると、この時間を 20 分に短縮できます。

次の YAML ファイルを使用して、Hyperdisk ML で AI/ML データの読み込みを高速化するチュートリアルに沿って操作できます。

  1. 次のマニフェストの例を producer-pvc.yaml として保存します。

    kind: PersistentVolumeClaim
apiVersion: v1
metadata:
  name: producer-pvc
spec:
  storageClassName: hyperdisk-ml
  accessModes:
  - ReadWriteOnce
  resources:
    requests:
      storage: 800Gi

  2. 次のマニフェストの例を producer-job.yaml として保存します。

    apiVersion: batch/v1
kind: Job
metadata:
  name: producer-job
spec:
  template:  # Template for the Pods the Job will create
    spec:
      affinity:
        nodeAffinity:
          requiredDuringSchedulingIgnoredDuringExecution:
            nodeSelectorTerms:
            - matchExpressions:
              - key: cloud.google.com/machine-family
                operator: In
                values:
                - "c3"
            - matchExpressions:
              - key: topology.kubernetes.io/zone
                operator: In
                values:
                - "ZONE"
      containers:
      - name: copy
        resources:
          requests:
            cpu: "32"
          limits:
            cpu: "32"
        image: python:3.11-bookworm
        command:
        - bash
        - -c
        - "pip install 'huggingface_hub==0.24.6' && \
          huggingface-cli download meta-llama/Meta-Llama-3.1-405B-Instruct --local-dir-use-symlinks=False --local-dir=/data/Meta-Llama-3.1-405B-Instruct --include *.safetensors *.json"
        env:
        - name: HUGGING_FACE_HUB_TOKEN
          valueFrom:
            secretKeyRef:
              name: hf-secret
              key: hf_api_token
        volumeMounts:
          - mountPath: "/data"
            name: volume
      restartPolicy: Never
      volumes:
        - name: volume
          persistentVolumeClaim:
            claimName: producer-pvc
  parallelism: 1         # Run 1 Pods concurrently
  completions: 1         # Once 1 Pods complete successfully, the Job is done
  backoffLimit: 4        # Max retries on failure

vLLM モデルサーバーをデプロイする

手順を完了したら、Hyperdisk ML ボリュームを使用する vLLM マルチノード GPU サーバー デプロイメントをデプロイできます。



apiVersion: leaderworkerset.x-k8s.io/v1
kind: LeaderWorkerSet
metadata:
  name: vllm
spec:
  replicas: 1
  leaderWorkerTemplate:
    size: 2
    restartPolicy: RecreateGroupOnPodRestart
    leaderTemplate:
      metadata:
        labels:
          role: leader
      spec:
        containers:
          - name: vllm-leader
            image: IMAGE_NAME
            env:
              - name: HUGGING_FACE_HUB_TOKEN
                valueFrom:
                  secretKeyRef:
                    name: hf-secret
                    key: hf_api_token
            command:
              - sh
              - -c
              - "/workspace/vllm/examples/ray_init.sh leader --ray_cluster_size=$(LWS_GROUP_SIZE); 
                python3 -m vllm.entrypoints.openai.api_server --port 8080 --model /models/Meta-Llama-3.1-405B-Instruct --tensor-parallel-size 8 --pipeline-parallel-size 2"
            resources:
              limits:
                nvidia.com/gpu: "8"
            ports:
              - containerPort: 8080
            readinessProbe:
              tcpSocket:
                port: 8080
              initialDelaySeconds: 15
              periodSeconds: 10
            volumeMounts:
              - mountPath: /dev/shm
                name: dshm
              - mountPath: /models
                name: llama3-405b
        volumes:
        - name: dshm
          emptyDir:
            medium: Memory
        - name: llama3-405b
          persistentVolumeClaim:
            claimName: hdml-static-pvc
    workerTemplate:
      spec:
        containers:
          - name: vllm-worker
            image: IMAGE_NAME
            command:
              - sh
              - -c
              - "/workspace/vllm/examples/ray_init.sh worker --ray_address=$(LWS_LEADER_ADDRESS)"
            resources:
              limits:
                nvidia.com/gpu: "8"
            env:
              - name: HUGGING_FACE_HUB_TOKEN
                valueFrom:
                  secretKeyRef:
                    name: hf-secret
                    key: hf_api_token
            volumeMounts:
              - mountPath: /dev/shm
                name: dshm
              - mountPath: /models
                name: llama3-405b
        volumes:
        - name: dshm
          emptyDir:
            medium: Memory
        - name: llama3-405b
          persistentVolumeClaim:
            claimName: hdml-static-pvc
---
apiVersion: v1
kind: Service
metadata:
  name: vllm-leader
spec:
  ports:
    - name: http
      port: 8080
      protocol: TCP
      targetPort: 8080
  selector:
    leaderworkerset.sigs.k8s.io/name: vllm
    role: leader
  type: ClusterIP

クリーンアップ

このチュートリアルで使用したリソースについて、Google Cloud アカウントに課金されないようにするには、リソースを含むプロジェクトを削除するか、プロジェクトを維持して個々のリソースを削除します。

デプロイされたリソースを削除する

このガイドで作成したリソースについて、Google Cloud アカウントに課金されないようにするには、次のコマンドを実行します。

gcloud container clusters delete CLUSTER_NAME \
  --location=ZONE

次のステップ