OpenShift AIでマルチノードのLLM推論を試す

 こんにちは、Red Hatでソリューションアーキテクトをしている石川です。 先日OpenShift AIのバージョン2.16がリリースされました。 現在OpenShift AIは非常にハイペースでリリースが行われており、約1ヶ月に一度のペースで機能アップデートが行われています。12月にリリースされたバージョン2.16についても多くの機能が追加されており、その中の一つとしてマルチノードでLLMの推論を行う機能がTech Previewとして追加されました。
docs.redhat.com

 LLMの推論を行う上で大事な要素として、どれだけのGPUのVRAM(メモリ)を利用できるかがあります。例えば、パラメーターが7BのLLMを量子化せずに動かそうとする場合、おおよそ必要なVRAMは以下の式で計算することができます。

7B * 2byte * 1.2= 16.8GB

 パラメーターが16bit表現の場合、パラメーターごとに2byte必要となり、そこにオーバーヘッドとして1.2を乗算して、必要なVRAMは凡そ16.8GBと計算できます。実際には同じパラメーターのモデルであってもコンテクスト長などの要素により、必要なVRAM容量は変わりますが、一つの目安とすることができます。

 7Bよりも大きな70B程度のサイズのモデルを動かそうとする場合、ネックとなるのがVRAMの容量です。70Bのモデルを先程の式に当てはめると、凡そ168GBのVRAM容量が必要だと計算できます。そのため、例えば、VRAM容量が24GBとなるNVIDIAのL4では、7~8台のGPUデバイスが要求されます。
 1台のノード上にこれらのGPUを備えることができれば、シングルノード上で推論を実施することができますが、そうではない場合、マルチノードでの推論を行う必要があります。前者については以前よりOpenShift AIで可能でしたが、後者については今回のバージョンアップで新たに機能が追加され、実行可能となりました。本記事ではこの機能を利用しマルチノードでの推論を実施します。

 なお、Tech Preview機能は商用環境での利用を想定しておらず、サポートの対象外となります。詳しくは以下をご参照下さい。 access.redhat.com

環境の構築

 まずマルチノード推論を行うために必要な環境を準備しましょう。今回はAWS上でクラスタを構成していきます。AWSでのOpenShift AIを使ったLLM推論環境のセットアップ方法については以前書いた以下の記事をご確認下さい。

rheb.hatenablog.com

 上記の記事と異なるポイントを2つご紹介します。

 一つ目は利用するGPUインスタンスの種類です。今回は70Bのモデルを実行するため、以下のインスタンスを2台クラスタに追加します。
・g6.12xlarge(48vCPU,192GiB Memory, L4 GPU 24GB VRAM * 4)

これによりクラスタ内で合計で8つのGPUデバイスを利用することができます。なお、AWSの東京リージョンではオンデマンドインスタンスとして利用する場合、1台あたり約$6.7の費用が発生するため、実際に試す際には利用コストに注意して下さい。  

 二つ目はPVです。OpenShift AIのマルチノード推論ではRWX(ReadWriteMany)のPVが必要となります。本記事ではOpenShift Data Foundationのocs-storagecluster-cephfsStorageClassを利用します。ODFのインストール方法については、公式ドキュメントの記載を参照して下さい。ROSA HCP環境を利用する場合は、必要な権限が付与されたIAMロールの作成とARNの設定が求めらるため、適宜設定を行って下さい。

 これでマルチノード推論を行うのに必要な準備が整いました。

利用対象のLLM

 今回利用するのは先日Meta社より発表されたLlama3.3です。 Meta公式のXのポストによると、モデルのアラインメントチューニングの改善、オンライン強化学習の適用により、70Bというモデルサイズでありながら、Llama3.1の405B相当の性能となっているとのことです。
 モデルはHugging Face上で公開されていますが、Gated Modelとなっているため、アクセスには個人情報の登録や規約への同意が必要です。必要な項目を入力し、しばらく待つと承認のメールが届き、モデルにアクセスすることができます。
huggingface.co

モデルのデプロイ

 それではモデルをデプロイしていきます。流れとしては以下の通りとなります。
(1) Hugging Faceからモデルを取得し、オブジェクトストレージにアップロード
(2) オブジェクトストレージからRWXのPVにモデルをダウンロード
(3) マルチノード推論のためのカスタムリソースを作成
順番に実施していきましょう。

(1) Hugging Faceからモデルを取得し、オブジェクトストレージにアップロード
 OpenShift AIコンソール上で、Data Science Project、Data Connectionを作成し、カスタムNotebookイメージとしてODH Toolsを登録します。モデルの取得とアップロードの手順については、前回のブログの"OpenShift AIにおける各種設定"、"モデルのダウンロード"と同じ手順となるため、具体的な手順についてはそちらを参照して下さい。登録したODH Toolsを用いてHugging FaceのLlama3.3のページからモデルを取得しましょう。 モデルサイズが大きいため完了まで20~30分程度かかります。

(2) オブジェクトストレージからRWXのPVにモデルをダウンロード
 マルチノード推論ではモデルを一度RWXのPVに保存し、それを各Podに接続、vLLMにロードします。そのためまず必要なPVCを作成します。OpenShiftのコンソール画面から以下の画像のようにPVCを作成しましょう。

StorageClass: ocs-storagecluster-cephfs
名前: model-store
アクセスモード: RWX
サイズ: 300GiB

 このPVCを作成するプロジェクトの中で最終的なモデルの推論まで行うため、必ず(1)で作成したData Science Projectの中でPVCを作成するようにして下さい。
 PVCが作成できたら以下のPodを作成します。このPodではオブジェクトストレージに保存したLLMをダウンロードし、PVをアタッチしたディレクトリに保存します。

$ cat << EOF | oc apply -f -
apiVersion: v1
kind: Pod
metadata:
  name: download-llama33-70b-instruct
  labels:
    name: download-llama33-70b-instruct
spec:
  volumes:
    - name: model-volume
      persistentVolumeClaim:
        claimName: model-store
  restartPolicy: Never
  initContainers:
    - name: fix-volume-permissions
      image: quay.io/quay/busybox:latest
      command: ["sh"]
      args: ["-c", "mkdir -p /mnt/models/llama33-70b-instruct && chmod -R 777 /mnt/models"] #①
      volumeMounts:
        - mountPath: "/mnt/models/"
          name: model-volume
  containers:
    - resources:
        requests:
          memory: 8Gi
      name: download-model
      imagePullPolicy: IfNotPresent
      image: quay.io/modh/kserve-storage-initializer:rhoai-2.16
      args:
        - 's3://my-trial-llm/meta-llama/Llama-3.3-70B-Instruct/' #②
        - /mnt/models/llama33-70b-instruct
      env:
        - name: STORAGE_CONFIG
          valueFrom:
            secretKeyRef:
              name: storage-config
              key: my-storage #③
      volumeMounts:
        - mountPath: "/mnt/models/"
          name: model-volume
EOF
  • ①: /mnt/models/llama33-70b-instructというディレクトリを作成し、モデルを保存。
  • ②: モデルが保存されたオブジェクトストレージのバケット、パスを指定。
  • ③: ②にアクセスするための認証情報をSecretより取得。ここでは一つ前の工程で作成したData ConnectionのSecretを指定。

 作成したPodが完了状態になれば、モデルのダウンロードは完了となります。

(3) マルチノード推論のためのカスタムリソースを作成
 それではマルチノード推論のためのカスタムリソースを作成しましょう。マルチノード推論はまだGUIからの作成に対応していないため、マニフェストを作成し実行します。
 最初にモデルの実行環境となるServingRuntimeリソースを作成します。ServingRuntimeでは利用するコンテナイメージや、コマンド引数などの情報が登録されます。作業環境から、以下を実行して、予め準備されているテンプレートからリソースを作成しましょう。

# Project(Namespace)を変更
$ oc project multi-node-serving

#  ServingRuntimeの作成
oc process vllm-multinode-runtime-template -n redhat-ods-applications|oc apply  -f -

 テンプレートから作成されたServingRuntimeを今回のモデルに合わせてカスタマイズします。oc editにより、vLLM実行時に渡す引数を追加します。

$ oc edit servingruntime vllm-multinode-runtime
...
  containers:
  - args:
    - "ray start --head --disable-usage-stats --include-dashboard false \n# wait for
      other node to join\nuntil [[ $(ray status --address ${RAY_ADDRESS} | grep -c
      node_) -eq ${PIPELINE_PARALLEL_SIZE} ]]; do\n  echo \"Waiting...\"\n  sleep
      1\ndone\nray status --address ${RAY_ADDRESS}\n\nexport SERVED_MODEL_NAME=${MODEL_NAME}\nexport
      MODEL_NAME=${MODEL_DIR} \n\nexec python3 -m vllm.entrypoints.openai.api_server
      --port=8080 --distributed-executor-backend ray --model=${MODEL_NAME} --served-model-name=${SERVED_MODEL_NAME}
      --tensor-parallel-size=${TENSOR_PARALLEL_SIZE} --pipeline-parallel-size=${PIPELINE_PARALLEL_SIZE}
      --disable_custom_all_reduce 
      --max-model-len=47104 --gpu-memory-utilization=0.98\n"
...

 ここでは最終行の--max-model-len=47104 --gpu-memory-utilization=0.98を新たに加えています。 Llama3.3は最大で128kのコンテクスト長に対応していますが、それを実現するにはより多くのVRAM容量が必要となるため、今回の環境で利用可能な最大の長さにコンテクスト長を制限します。
  最後にモデルを表すリソースであるInferenceServiceをデプロイします。

$ cat << EOF | oc apply -f -
apiVersion: serving.kserve.io/v1beta1
kind: InferenceService
metadata:
  annotations: #①
    serving.kserve.io/deploymentMode: RawDeployment 
    serving.kserve.io/autoscalerClass: external
  name: llama33-70b-instruct
spec:
  predictor:
    model:
      modelFormat:
        name: vLLM
      runtime: vllm-multinode-runtime #②
      storageUri: pvc://model-store/llama33-70b-instruct #③
    workerSpec: #④
      tensorParallelSize: 4
      pipelineParallelSize: 2
EOF
  • ①: annotationsでKserveのRawDeployment(Serverlessを利用しないモード)を設定。
  • ②: 先ほど作成したServingRuntimeを指定。
  • ③: モデルを保存先のPVCとパスを指定。
  • ④: tensorParallelSizeにはノードあたりのGPU数、pipelineParallelSizeにはGPUノード数をそれぞれ指定。(参照)

 InferenceServiceを作成すると推論用のPodが各GPUノードで起動します。上記で設定した通り、モデル分割が行われ、それぞれのPodで起動するvLLM上にロードされます。
 vLLMでマルチノードでの推論を実行するにはオープンソースのRayが用いられており、各Podが一つのRay Clusterとして機能します。

 モデルが無事起動したかどうかは、Podのログを確認しましょう。

$ oc logs $(oc get pods -o name | grep predictor | grep -v worker) | less

# 以下が出力されていればOK
...
INFO:     Uvicorn running on socket ('0.0.0.0', 8080) (Press CTRL+C to quit)
INFO:     127.0.0.1:60290 - "GET /health HTTP/1.1" 200 OK
INFO:     127.0.0.1:60298 - "GET /health HTTP/1.1" 200 OK
INFO 12-16 06:21:24 metrics.py:349] Avg prompt throughput: 0.0 tokens/s, Avg generation throughput: 0.0 tokens/s, Running: 0 reqs, Swapped: 0 reqs, Pending: 0 reqs, GPU KV cache usage: 0.0%, CPU KV cache usage: 0.0%.
...

 上記のログが出力され、入力トークンを待ち受けている状態となっていれば正しく動作しています。Routeを作成し、APIのアクセスエンドポイントを作成しましょう。

$ cat << EOF | oc apply -f -
apiVersion: route.openshift.io/v1
kind: Route
metadata:
  name: llama33-70b-instruct-predictor
  namespace: multi-node-serving
spec:
  to:
    name: llama33-70b-instruct-predictor
    kind: Service
  tls:
    insecureEdgeTerminationPolicy: Redirect
    termination: edge
  port:
    targetPort: http
EOF

# URLの表示
$ echo https://$(oc get route llama33-70b-instruct-predictor -o "jsonpath={.spec.host}")

モデルの実行

 モデルがデプロイできたので、APIを実行してみましょう。今回はフロントエンドのUIにOpen WebUIを利用します。以下の手順に従い、アプリケーションをデプロイしてください。

$ git clone https://github.com/JPishikawa/open-webui-openshift

$ cd open-webui-openshift

$ vi manifests/openwebui-cm.yaml
...
data:
  WEBUI_AUTH: 'False'
  # 以下を変更
  OPENAI_API_BASE_URLS: 'https://CHANGE_HERE/v1' 
  OPENAI_API_KEYS: na
...

$ chmod +x deploy.sh
$ ./deploy.sh

# アプリURLの表示
$ echo https://$(oc get route open-webui -o "jsonpath={.spec.host}" -n app-webui)

 ブラウザからアプリにアクセスし、LLMを実行してみましょう。画面左上で、呼び出すモデルを指定し、クエリを実行します。

 APIが実行され、結果が返ってくることが確認できました。体感として若干のトークン生成の遅さを感じますが、こちらについては量子化モデルを利用したり、オンプレミスであればGPUノード間のRDMA設定を行うことで改善が可能だと考えられます。

まとめ

 本記事では、OpenShift AIでマルチノードによるLLM推論機能をご紹介しました。マルチノードを利用することで、今回のようにサイズの大きいモデルをホストしたり、より長いコンテクストで利用し易くなります。現在はまだTPであるため、機能として洗練されていない点もありますが、気になった方はぜひ試してみて下さい。

* 各記事は著者の見解によるものでありその所属組織を代表する公式なものではありません。その内容については非公式見解を含みます。