第6章 Google TPU AI アクセラレーターを使用した Podman による言語モデルのサービングと推論

手順

1. TPU サーバーホストでターミナルを開き、registry.redhat.io にログインします。

   $ podman login registry.redhat.io

2. 次のコマンドを実行して、Red Hat AI Inference Server イメージをプルします。

   $ podman pull registry.redhat.io/rhaiis/vllm-tpu-rhel9:3.2.5

3. オプション: ホストで TPU が使用可能であることを確認します。

   1. Red Hat AI Inference Server コンテナーでシェルプロンプトを開きます。以下のコマンドを実行します。

      $ podman run -it --net=host --privileged -e PJRT_DEVICE=TPU --rm --entrypoint /bin/bash registry.redhat.io/rhaiis/vllm-tpu-rhel9:3.2.5

   2. コンテナーシェルプロンプトで次の Python コードを実行して、システム TPU アクセスと基本操作を確認します。

      $ python3 -c "
      import torch
      import torch_xla
      try:
          device = torch_xla.device()
          print(f'')
          print(f'XLA device available: {device}')
          x = torch.randn(3, 3).to(device)
          y = torch.randn(3, 3).to(device)
          z = torch.matmul(x, y)
          import torch_xla.core.xla_model as xm
          torch_xla.sync()
          print(f'Matrix multiplication successful')
          print(f'Result tensor shape: {z.shape}')
          print(f'Result tensor device: {z.device}')
          print(f'Result tensor: {z.data}')
          print('TPU is operational.')
      except Exception as e:
          print(f'TPU test failed: {e}')
          print('Try restarting the container to clear TPU locks')
      "

      出力例

      XLA device available: xla:0
      Matrix multiplication successful
      Result tensor shape: torch.Size([3, 3])
      Result tensor device: xla:0
      Result tensor: tensor([[-1.8161,  1.6359, -3.1301],
              [-1.2205,  0.8985, -1.4422],
              [ 0.0588,  0.7693, -1.5683]], device='xla:0')
      TPU is operational.

   3. シェルプロンプトを終了します。

      $ exit

4. ボリュームを作成してコンテナーにマウントします。コンテナーが使用できるようにコンテナーの権限を調整します。

   $ mkdir ./.cache/rhaiis

   $ chmod g+rwX ./.cache/rhaiis

5. HF_TOKEN Hugging Face トークンを private.env ファイルに追加します。

   $ echo "export HF_TOKEN=<huggingface_token>" > private.env

6. HF_HOME 変数を private.env ファイルに追加します。

   $ echo "export HF_HOME=./.cache/rhaiis" >> private.env

   source コマンドで private.env ファイルを読み込みます。

   $ source private.env

7. AI Inference Server コンテナーイメージを起動します。

   podman run --rm -it \
     --name vllm-tpu \
     --network=host \
     --privileged \
     --shm-size=4g \
     --device=/dev/vfio/vfio \
     --device=/dev/vfio/0 \
     -e PJRT_DEVICE=TPU \
     -e HF_HUB_OFFLINE=0 \
     -v ./.cache/rhaiis:/opt/app-root/src/.cache \
     registry.redhat.io/rhaiis/vllm-tpu-rhel9:3.2.5 \
     --model Qwen/Qwen2.5-1.5B-Instruct \
     --tensor-parallel-size 1 \ 

   1

   
     --max-model-len=256 \ 

   2

   
     --host=0.0.0.0 \
     --port=8000

   1

   TPU の数に合わせて --tensor-parallel-size を設定します。

   2

   最高のパフォーマンスを得るには、ワークロードで許容される限り、max-model-len パラメーターを低く設定してください。