第 6 章 使用 Google TPU AI Accelerators 使用 Podman 提供和推断语言模型

流程

1. 在 TPU 服务器主机上打开一个终端，并登录到 registry.redhat.io ：

   $ podman login registry.redhat.io

   Copy to Clipboard Toggle word wrap

2. 运行以下命令拉取 Red Hat AI Inference Server 镜像：

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

   Copy to Clipboard Toggle word wrap

3. 可选：验证主机中是否有 TPU。

   1. 在 Red Hat AI Inference Server 容器中打开 shell 提示符。运行以下命令:

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

      Copy to Clipboard Toggle word wrap

   2. 通过在容器 shell 提示符中运行以下 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')
      "

      Copy to Clipboard Toggle word wrap

      输出示例

      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.

      Copy to Clipboard Toggle word wrap

   3. 退出 shell 提示符。

      $ exit

      Copy to Clipboard Toggle word wrap

4. 创建卷并将其挂载到容器中。调整容器权限，以便容器可以使用它。

   $ mkdir ./.cache/rhaiis

   Copy to Clipboard Toggle word wrap

   $ chmod g+rwX ./.cache/rhaiis

   Copy to Clipboard Toggle word wrap

5. 将 HF_TOKEN Hugging Face 令牌添加到 private.env 文件中。

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

   Copy to Clipboard Toggle word wrap

6. 将 HF_HOME 变量附加到 private.env 文件。

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

   Copy to Clipboard Toggle word wrap

   提供 private.env 文件。

   $ source private.env

   Copy to Clipboard Toggle word wrap

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.4 \
     --model Qwen/Qwen2.5-1.5B-Instruct \
     --tensor-parallel-size 1 \ 

   1

   
     --max-model-len=256 \ 

   2

   
     --host=0.0.0.0 \
     --port=8000

   Copy to Clipboard Toggle word wrap

   1

   set- tensor-parallel-size，以匹配 TPU 的数量。

   2

   为了获得最佳性能，请将 max-model-len 参数设置为 low，因为您的工作负载允许。