Chapter 5. 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

   Copy to Clipboard Toggle word wrap

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

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

   Copy to Clipboard Toggle word wrap

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.2.2

      Copy to Clipboard Toggle word wrap

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

      Copy to Clipboard Toggle word wrap

      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.

      Copy to Clipboard Toggle word wrap

   3. Exit the shell prompt.

      $ exit

      Copy to Clipboard Toggle word wrap

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

   $ mkdir ./.cache/rhaiis

   Copy to Clipboard Toggle word wrap

   $ chmod g+rwX ./.cache/rhaiis

   Copy to Clipboard Toggle word wrap

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

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

   Copy to Clipboard Toggle word wrap

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

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

   Copy to Clipboard Toggle word wrap

   Source the private.env file.

   $ source private.env

   Copy to Clipboard Toggle word wrap

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.2.2 \
     --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 to match the number of TPUs.

   2

   For the best performance, set the max-model-len parameter to as low as your workload allows.