Serving Polish Large Language Model (PLLuM) on TPUs

There are many really good and popular LLMs that support multiple languages. Yet sometimes, you might find yourself in need of a real specialist when it comes to a specific language or culture. One of such specialists is PLLuM (Polish Large Language Model), a model created specifically for handling interactions in the Polish language. Trained on vast sets of native source materials, models like this can handle communication in a given language better than bigger, more universal models.

Deploying a model like PLLuM on GPU architecture is as simple as with any other model. However, when you consider running it on a TPU, you may discover that inference engines like vLLM don’t list it as a supported model. Does it mean you can’t run it on TPUs? No! In this article, I’ll show you how to set up a PLLuM inference service on Trillium TPU nodes in GKE.

Setting up the GKE cluster

We will start with setting up a simple GKE Autopilot cluster. This will allow us to focus on deploying the model, without manual node management.

Let’s begin with some environment variables:

export PROJECT_ID="YOUR_PROJECT_ID"
export CLUSTER_NAME="YOUR_CLUSTER_NAME"
export REGION="CLUSTER_REGION" 
export ZONE="CLUSTER_ZONE"
export RESERVATION="RESERVATION_NAME"
export HF_TOKEN="your-hugging-face-token"

In my example, I will be using a TPU reservation I made using DWS Calendar Mode. This way I won’t have to worry about hardware availability. Other deployment options are also available, but don’t guarantee hardware availability.

Create the GKE cluster, set up kubectl to work on this cluster and create a Kubernetes secret that will hold your Hugging Face read token:

gcloud container clusters create-auto $CLUSTER_NAME \
    --project=$PROJECT_ID \
    --region=$REGION

gcloud container clusters get-credentials $CLUSTER_NAME --location=$REGION --project=$PROJECT_ID
    
kubectl create secret generic hf-secret \
    --from-literal=hf_api_token=${HF_TOKEN} \
    --dry-run=client -o yaml | kubectl apply -f -

Deploying the model on a single TPU v6e chip

We can start with a simple deployment of CYFRAGOVPL/Llama-PLLuM-8B-instruct-2512 — this is a small model (16 GB) that fits just fine in the memory of a single TPU v6e chip (32 GB). That makes our deployment simple and easy, though limited in performance. It is the best option for development and testing. Use the following YAML to start your single-chip deployment (save it as pllum-1.yaml):

apiVersion: apps/v1
kind: Deployment
metadata:
  name: vllm-tpu
spec:
  replicas: 1
  selector:
    matchLabels:
      app: vllm-tpu
  template:
    metadata:
      labels:
        app: vllm-tpu
    spec:
      nodeSelector:
        cloud.google.com/gke-tpu-topology: 1x1
        # Following values can be used to run v6e TPUs with a reservation.
        cloud.google.com/gke-tpu-accelerator: tpu-v6e-slice
        cloud.google.com/reservation-affinity: "specific"
        cloud.google.com/reservation-name: RESERVATION_NAME
      containers:
      - name: vllm-tpu
        image: docker.io/vllm/vllm-tpu:latest
        command: ["python3", "-m", "vllm.entrypoints.openai.api_server"]
        args:
        - --host=0.0.0.0
        - --port=8000
        - --tensor-parallel-size=1
        - --max-model-len=4096
        - --model=CYFRAGOVPL/Llama-PLLuM-8B-instruct-2512
        - --max-num-batched-tokens=512
        - --max-num-seqs=128
        env:
        - name: HUGGING_FACE_HUB_TOKEN
          valueFrom:
            secretKeyRef:
              name: hf-secret
              key: hf_api_token
        ports:
        - containerPort: 8000
        resources:
          limits:
            google.com/tpu: 1
        readinessProbe:
          tcpSocket:
            port: 8000
          initialDelaySeconds: 1800
          periodSeconds: 60
          failureThreshold: 10
          timeoutSeconds: 5
        volumeMounts:
        - name: dshm
          mountPath: /dev/shm
      volumes:
      - name: dshm
        emptyDir:
          medium: Memory
---
apiVersion: v1
kind: Service
metadata:
  name: vllm-service
spec:
  selector:
    app: vllm-tpu
  type: LoadBalancer
  ports:
    - name: http
      protocol: TCP
      port: 8000
      targetPort: 8000

Remember to replace RESERVATION_NAME with a proper value. Then apply the file with kubectl apply -f pllum-1.yaml.

The deployment will take a while. You can watch how first the TPU node is created (kubectl get nodes) and how the vLLM pod starts (kubectl logs -l app=vllm-tpu).

Once the pod is ready to work, you can check the external IP of the created service with: kubectl get svc— just check what’s in the External IP column.

Now, when we know that external IP, we can send a request to our model:

curl http://${EXTERNAL_IP}:8000/v1/completions \
    -H "Content-Type: application/json" \
    -d '{
        "model": "CYFRAGOVPL/Llama-PLLuM-8B-instruct-2512",
        "prompt": "Jak to jest być skrybą?",
        "max_tokens": 150,
        "temperature": 0.7
    }'

The answer you get is a JSON containing various information, including the completion generated by the model!

Scaling to 8 TPUs

To handle more requests concurrently and increase the efficiency of your service, you may want to scale the vLLM service to work with more than one TPU chip. If you try to simply switch — tensor-parallel-size to 8, you will get an error message similar to:

(EngineCore pid=572) jax._src.sharding.IndivisibleError: One of device_put args was given the sharding NamedSharding(mesh=Mesh(‘data’: 1, ‘model’: 8, axis_types=(Auto, Auto)), spec=P((‘data’, ‘model’), None), memory_kind=device) implies that array axis 0 is partitioned 8 times, but the dimension size is 128259 (full shape: (128259, 4096), per-dimension tiling factors: (8, 1) should evenly divide the shape). This error occurs at source: /workspace/tpu_inference/tpu_inference/layers/common/utils.py:149:23 (general_device_put.<locals>._put)

The model we chose (Llama-PLLuM-8B-instruct) has a vocabulary size of 128259 — that’s not divisible by 8! The way vLLM spreads the work across multiple chips requires the model to have a vocabulary divisible by the number of chips that will be used. In this case, we will need to pad the model’s vocabulary to 128264 with “fake tokens”.

Here’s a script that will do just that. It will download the model locally, modify it, and store it in a new directory:

# pip install torch --index-url https://download.pytorch.org/whl/cpu
# pip install huggingface_hub safetensors

import json
import os
import torch
from huggingface_hub import snapshot_download
from safetensors.torch import load_file, save_file

MODEL_ID = "CYFRAGOVPL/Llama-PLLuM-8B-instruct-2512"
OUTPUT_DIR = "./Llama-PLLuM-8B-instruct-2512-padded"
TARGET_VOCAB_SIZE = 128264  # 16033 * 8

# 1. Download model snapshot
model_path = snapshot_download(repo_id=MODEL_ID)
os.makedirs(OUTPUT_DIR, exist_ok=True)

# 2. Copy and update config.json
with open(os.path.join(model_path, "config.json"), "r") as f:
    config = json.load(f)

old_vocab_size = config["vocab_size"]
pad_amount = TARGET_VOCAB_SIZE - old_vocab_size
config["vocab_size"] = TARGET_VOCAB_SIZE

with open(os.path.join(OUTPUT_DIR, "config.json"), "w") as f:
    json.dump(config, f, indent=2)

# 3. Copy other metadata files (tokenizers, generation configs, etc.)
for file_name in os.listdir(model_path):
    if not file_name.endswith(".safetensors") and file_name != "config.json":
        src = os.path.join(model_path, file_name)
        dst = os.path.join(OUTPUT_DIR, file_name)
        if os.path.isfile(src):
            os.system(f"cp {src} {dst}")

# 4. Pad embedding and lm_head tensors in safetensors
for file_name in os.listdir(model_path):
    if file_name.endswith(".safetensors"):
        tensors = load_file(os.path.join(model_path, file_name))
        modified = False

        for key in ["model.embed_tokens.weight", "lm_head.weight"]:
            if key in tensors:
                tensor = tensors[key]
                if tensor.shape[0] == old_vocab_size:
                    padding = torch.zeros(
                        (pad_amount, *tensor.shape[1:]), dtype=tensor.dtype
                    )
                    tensors[key] = torch.cat([tensor, padding], dim=0)
                    modified = True

        save_file(tensors, os.path.join(OUTPUT_DIR, file_name))

print(f"✅ Successfully padded model saved to: {OUTPUT_DIR}")

To make use of the new model, we need to upload it to a Cloud Storage Bucket and allow the vLLM pod access to it, using the Workload Identity bindings.

# We start with creating a new bucket
export BUCKET_NAME="padded-pllum-model"

gcloud storage buckets create gs://${BUCKET_NAME} \
    --project=${PROJECT_ID} \
    --location=${REGION}

# Upload the modified model to the bucket
gcloud storage cp -r ./Llama-PLLuM-8B-instruct-2512-padded gs://${BUCKET_NAME}/

# Create GCP Service Account and bind with GKE Service Account via Workload Identity for GCS Fuse
export GSA_NAME="vllm-pllum-gcs-sa"
export KSA_NAME="default"
export NAMESPACE="default"

gcloud iam service-accounts create ${GSA_NAME} \
    --project=${PROJECT_ID}

gcloud storage buckets add-iam-policy-binding gs://${BUCKET_NAME} \
    --member="serviceAccount:${GSA_NAME}@${PROJECT_ID}.iam.gserviceaccount.com" \
    --role="roles/storage.objectViewer"

gcloud iam service-accounts add-iam-policy-binding ${GSA_NAME}@${PROJECT_ID}.iam.gserviceaccount.com \
    --project=${PROJECT_ID} \
    --role="roles/iam.workloadIdentityUser" \
    --member="serviceAccount:${PROJECT_ID}.svc.id.goog[${NAMESPACE}/${KSA_NAME}]"

kubectl annotate serviceaccount ${KSA_NAME} \
    --namespace=${NAMESPACE} \
    iam.gke.io/gcp-service-account=${GSA_NAME}@${PROJECT_ID}.iam.gserviceaccount.com \
    --overwrite

At this point, all that’s left is to modify the pllum-1.yaml file to use the new, modified model and 8 TPUs. Let’s save it as pllum-2.yaml:

apiVersion: apps/v1
kind: Deployment
metadata:
  name: vllm-tpu
spec:
  replicas: 1
  selector:
    matchLabels:
      app: vllm-tpu
  template:
    metadata:
      labels:
        app: vllm-tpu
      annotations:
        gke-gcsfuse/volumes: "true"
    spec:
      nodeSelector:
        cloud.google.com/gke-tpu-topology: 2x4
        # Following values can be used to run v6e TPUs with a reservation.
        cloud.google.com/gke-tpu-accelerator: tpu-v6e-slice
        cloud.google.com/reservation-affinity: "specific"
        cloud.google.com/reservation-name: RESERVATION_NAME
      containers:
      - name: vllm-tpu
        image: docker.io/vllm/vllm-tpu:latest
        command: ["python3", "-m", "vllm.entrypoints.openai.api_server"]
        args:
        - --host=0.0.0.0
        - --port=8000
        - --tensor-parallel-size=8
        - --max-model-len=4096
        - --model=/models/Llama-PLLuM-8B-instruct-2512-padded
        - --served-model-name=CYFRAGOVPL/Llama-PLLuM-8B-instruct-2512
        - --max-num-batched-tokens=512
        - --max-num-seqs=128
        env:
        - name: HUGGING_FACE_HUB_TOKEN
          valueFrom:
            secretKeyRef:
              name: hf-secret
              key: hf_api_token
        ports:
        - containerPort: 8000
        resources:
          limits:
            google.com/tpu: 8
        readinessProbe:
          tcpSocket:
            port: 8000
          initialDelaySeconds: 1800
          periodSeconds: 60
          failureThreshold: 10
          timeoutSeconds: 5
        volumeMounts:
        - name: dshm
          mountPath: /dev/shm
        - name: gcs-fuse-csi-ephemeral
          mountPath: /models
          readOnly: true
      volumes:
      - name: dshm
        emptyDir:
          medium: Memory
      - name: gcs-fuse-csi-ephemeral
        csi:
          driver: gcsfuse.csi.storage.gke.io
          readOnly: true
          volumeAttributes:
            bucketName: padded-pllum-model
            mountOptions: "implicit-dirs"
---
apiVersion: v1
kind: Service
metadata:
  name: vllm-service
spec:
  selector:
    app: vllm-tpu
  type: LoadBalancer
  ports:
    - name: http
      protocol: TCP
      port: 8000
      targetPort: 8000

Once again, remember to replace RESERVATION_NAME with the proper name of your TPU reservation. Then apply the new YAML with kubectl apply -f pllum-2.yaml. Make sure you trigger the recreation of the vLLM pod by manually deleting it with: kubectl delete pod -l app=vllm-tpu.

You can verify that it works by again using a simple curl command:

# Get the external ip from: kubectl get svc

curl http://${EXTERNAL_IP}:8000/v1/completions \
    -H "Content-Type: application/json" \
    -d '{
        "model": "CYFRAGOVPL/Llama-PLLuM-8B-instruct-2512",
        "prompt": "Jak to jest być skrybą?",
        "max_tokens": 150,
        "temperature": 0.7
    }'

When you’re done with your experiments, you can delete the deployment, the cluster, and the storage bucket with the following commands:

# Delete the deployment
kubectl delete -f pllum-2.yaml

# Delete the cluster
gcloud container clusters delete $CLUSTER_NAME --location=$REGION --project=$PROJECT_ID

# Delete the storage bucket
gcloud storage rm --recursive gs://BUCKET_NAME

Serving Polish Large Language Model (PLLuM) on TPUs was originally published in Google Cloud – Community on Medium, where people are continuing the conversation by highlighting and responding to this story.

Source Credit: https://medium.com/google-cloud/serving-polish-large-language-model-pllum-on-tpus-877478f77dda?source=rss—-e52cf94d98af—4

About the Author

Deven Goratela

Administrator

Visit Website View All Posts