Chapter 6. Serving and inferencing language models with Podman using Google TPU AI accelerators

Procedure

1. Open a terminal on your TPU server host, and log in to registry.redhat.io:

   $ podman login registry.redhat.io

2. Pull the Red Hat AI Inference Server image by running the following command:

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

3. Optional: Verify that the TPUs are available in the host.

   1. Open a shell prompt in the Red Hat AI Inference Server container. Run the following command:

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

   2. Verify system TPU access and basic operations by running the following Python code in the container shell prompt:

      $ 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')
      "

      Example output

      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 the shell prompt.

      $ exit

4. Create a volume and mount it into the container. Adjust the container permissions so that the container can use it.

   $ mkdir ./.cache/rhaiis

   $ chmod g+rwX ./.cache/rhaiis

5. Add the HF_TOKEN Hugging Face token to the private.env file.

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

6. Append the HF_HOME variable to the private.env file.

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

   Source the private.env file.

   $ source private.env

7. Start the AI Inference Server container image:

   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.3.0 \
     --model Qwen/Qwen2.5-1.5B-Instruct \
     --tensor-parallel-size 1 \
     --max-model-len=256 \
     --host=0.0.0.0 \
     --port=8000

   Where:

   --tensor-parallel-size 1

   Specifies the number of TPUs to use for tensor parallelism. Set this value to match the number of available TPUs.

   --max-model-len=256

   Specifies the maximum model context length. For optimal performance, set this value as low as your workload allows.