Entrega Qwen3-30B-A3B en GKE con vLLM TorchTPU en TPU v6e

En este instructivo, se muestra cómo implementar y entregar el modelo de mezcla de expertos Qwen3-30B-A3B (Qwen/Qwen3-30B-A3B-FP8) en Google Kubernetes Engine (GKE) con un acelerador TPU v6e (Trillium) de un solo host y el framework de entrega de vLLM TPU con TorchTPU (torchtpu). Ejecutarás la carga de trabajo en un solo nodo ct6e-standard-4t, que contiene cuatro chips TPU v6e en una topología 2x2, con el paralelismo de datos y el paralelismo de expertos habilitados.

Para ejecutar un flujo de trabajo de entrega de TorchTPU de vLLM de host único similar en aceleradores TPU7x (Ironwood) con el modelo Qwen3.5-397B más grande, consulta Entrega Qwen3.5-397B en GKE con TorchTPU de vLLM en aceleradores TPU7x.

Objetivos

  • Configura tu clúster de GKE y el grupo de nodos TPU v6e de host único.
  • Implementar el servidor de vLLM TorchTPU de host único en GKE.
  • Envía una solicitud de inferencia al servidor de vLLM TorchTPU.
  • Ejecuta una carga de trabajo de comparativas de entrega.
  • Realizar una limpieza

Costos

En este documento, usarás los siguientes componentes facturables de Google Cloud:

Para generar una estimación de costos en función del uso previsto, usa la calculadora de precios.

Cuando finalices las tareas que se describen en este documento, puedes borrar los recursos que creaste para evitar que continúe la facturación. Para obtener más información, consulta Cómo realizar una limpieza.

Antes de comenzar

Configura tu entorno para el instructivo:

  • Habilita la API de Kubernetes Engine en tu proyecto de Google Cloud :

    Habilita la API de Kubernetes Engine

  • Habilita la API de Artifact Registry en tu proyecto de Google Cloud :

    Habilita la API de Artifact Registry

  • Instala y, luego, inicializa Google Cloud CLI, y, luego, instala kubectl. Si usas Cloud Shell, ambas herramientas están preinstaladas.

  • Asegúrate de que tu proyecto tenga suficiente cuota de SSD y de direcciones IP externas en la región en la que se encuentra tu capacidad de TPU:

    1. En la consola de Google Cloud , ve a la página Cuotas y límites del sistema.

      Ir a Cuotas y límites del sistema

    2. Filtra por SSD de Persistent Disk (GB) y verifica que tengas al menos 500 GB de cuota disponible para el volumen de caché del modelo.

    3. Filtra por Direcciones IPv4 externas regionales en uso y verifica que tengas suficiente cuota de direcciones IP en tu región de destino.

Obtén un token de acceso de Hugging Face

Necesitas un token de acceso de Hugging Face para usar este instructivo. Puedes registrarte para obtener una cuenta gratuita en Hugging Face. Después de tener una cuenta, genera un token de acceso:

  1. En la página Welcome to Hugging Face, haz clic en el avatar de tu cuenta y selecciona Tokens de acceso.
  2. En la página Tokens de acceso, haz clic en Crear token nuevo.
  3. Selecciona el tipo de token Lectura y, luego, ingresa un nombre para el token.
  4. Se mostrará tu token de acceso. Guarda el token en un lugar seguro.

Roles obligatorios

Si quieres obtener los permisos que necesitas para completar este instructivo, pídele a tu administrador que te otorgue los siguientes roles de IAM en tu proyecto:

Para obtener más información sobre cómo otorgar roles, consulta Administra el acceso a proyectos, carpetas y organizaciones.

También puedes obtener los permisos necesarios a través de roles personalizados o cualquier otro rol predefinido.

Configura tus variables de entorno

Configura la autenticación:

gcloud auth login
gcloud auth application-default login

Configura las variables de entorno:

export PROJECT_ID="YOUR_PROJECT_ID"
export REGION="YOUR_REGION"
export ZONE="YOUR_ZONE"
export CLUSTER_NAME="YOUR_CLUSTER_NAME"
export MACHINE_TYPE="tpu7x-standard-4t"
export NODE_POOL_NAME="np-0"
export RESERVATION_NAME="YOUR_RESERVATION_NAME"
export HF_TOKEN="YOUR_HF_TOKEN"
export NETWORK="YOUR_NETWORK_NAME"
export SUBNETWORK="YOUR_SUBNETWORK_NAME"
export VLLM_IMAGE="YOUR_VLLM_TPU_IMAGE"

gcloud config set project "${PROJECT_ID}"
gcloud config set billing/quota_project "${PROJECT_ID}"

Reemplaza lo siguiente:

  • YOUR_PROJECT_ID: Es el ID de tu proyecto de Google Cloud.
  • YOUR_REGION: Es la región en la que deseas implementar tu clúster, como us-central1.
  • YOUR_ZONE: Es la zona en la que deseas implementar tu grupo de nodos de TPU v6e, como us-central1-c.
  • YOUR_CLUSTER_NAME: Es el nombre de tu clúster de GKE, como tpu-v6e-tutorial-cluster.
  • YOUR_RESERVATION_NAME: Si usas una reserva, es el nombre de tu reserva de TPU v6e.
  • YOUR_HF_TOKEN: Es tu token de acceso de Hugging Face.
  • YOUR_NETWORK_NAME: Es el nombre de tu red deGoogle Cloud . Si creaste una red personalizada, especifica esa red. De lo contrario, especifica default.
  • YOUR_SUBNETWORK_NAME: Es el nombre de tu subred deGoogle Cloud . Si creaste una subred personalizada, especifícala. De lo contrario, especifica default.
  • YOUR_VLLM_TPU_IMAGE: Es la imagen de TorchTPU de vLLM que se usará. Para las imágenes públicas, consulta el repositorio de vLLM TorchTPU.

Crea el clúster de GKE

Para crear un clúster de GKE en modo estándar y conectar un grupo de nodos TPU v6e de host único, completa los siguientes pasos:

  1. Crea un clúster de GKE regional:

    gcloud container clusters create "${CLUSTER_NAME}" \
        --project="${PROJECT_ID}" \
        --location="${REGION}" \
        --node-locations="${ZONE}" \
        --network="${NETWORK}" \
        --subnetwork="${SUBNETWORK}"

    El tipo de máquina predeterminado para los clústeres de GKE (e2-medium) no tiene suficiente capacidad de CPU para los controladores de cargas de trabajo, por lo que este comando especifica e2-standard-4.

    El resultado es similar a lo siguiente:

    NAME                      LOCATION     MASTER_VERSION      MASTER_IP      MACHINE_TYPE   NODE_VERSION        NUM_NODES  STATUS
    tpu-v6e-tutorial-cluster  us-central1  1.35.7-gke.1027000  34.63.222.246  e2-standard-4  1.35.7-gke.1027000  3          RUNNING
    
  2. Crea un grupo de nodos TPU v6e de host único con un nodo ct6e-standard-4t (cuatro chips TPU v6e en una topología 2x2) con tu reserva:

    gcloud container node-pools create "${NODE_POOL_NAME}" \
        --cluster="${CLUSTER_NAME}" \
        --machine-type="${MACHINE_TYPE}" \
        --location="${REGION}" \
        --node-locations="${ZONE}" \
        --project="${PROJECT_ID}" \
        --num-nodes=1 \
        --reservation-affinity=specific \
        --reservation="${RESERVATION_NAME}"

    En lugar de usar una reserva específica, también puedes crear el grupo de nodos con capacidad de Spot o VMs de inicio flexible pasando las siguientes marcas:

    • VMs Spot: Reemplaza las marcas --reservation-affinity y --reservation por --spot para usar la capacidad de Spot. Las instancias que usan la capacidad de Spot pueden tardar en crearse después de que se crea el grupo de nodos y pueden interrumpirse durante las ejecuciones de las cargas de trabajo. Si tus Pods tardan en crearse o las cargas de trabajo se detienen en medio de la ejecución, esta podría ser la causa.
    • VMs de inicio flexible: Para aprovisionar el grupo de nodos con la opción de consumo Flex-start, pasa la marca --flex-start. Para obtener más información, consulta Ejecuta una carga de trabajo por lotes pequeña con TPU y VMs de inicio flexible.

    El resultado es similar a lo siguiente:

    NAME  MACHINE_TYPE      DISK_SIZE_GB  NODE_VERSION
    np-0  ct6e-standard-4t  100           1.35.7-gke.1027000
    
  3. Configura las credenciales de kubectl para comunicarte con tu clúster:

    gcloud container clusters get-credentials "${CLUSTER_NAME}" \
        --location="${REGION}" \
        --project="${PROJECT_ID}"

Configura las credenciales

Crea un Secret de Kubernetes llamado hf-secret para almacenar tu token de acceso de Hugging Face:

kubectl create secret generic hf-secret \
    --from-literal=hf_api_token="${HF_TOKEN}" \
    --dry-run=client -o yaml | kubectl apply -f -

Implementa el servidor de vLLM TorchTPU

El manifiesto de la carga de trabajo (qwen3_5_server.yaml) define los siguientes recursos:

  • Un StorageClass (hyperdisk-balanced-qwen3) y un PersistentVolumeClaim de 500 GiB (hd-claim-qwen3) respaldados por el almacenamiento de Google Cloud Hyperdisk Balanced para almacenar los pesos del modelo descargado, la caché de compilación de XLA y la caché de compilación de PyTorch.

  • Una Deployment (vllm-qwen3) programada en un solo nodo de TPU v6e con una topología 2x2 (google.com/tpu: '4'). El contenedor ejecuta vllm serve para Qwen/Qwen3-30B-A3B-FP8 con las siguientes marcas:

    • --data-parallel-size=4
    • --tensor-parallel-size=1
    • --enable-expert-parallel
    • --kv-cache-dtype=fp8
    • --max-model-len=9236
    • --max-num-seqs=64
  • Un Service ClusterIP (vllm-qwen3-service) que expone el puerto 8000.

Para implementar y supervisar la carga de trabajo, completa los siguientes pasos:

  1. Recupera el manifiesto del servidor:

    curl -O https://raw.githubusercontent.com/GoogleCloudSamples/AIHypercomputerSamples/refs/heads/main/tpu/inference/qwen3-5-397b/qwen3_5_server.yaml

    Como alternativa, puedes crear un archivo llamado qwen3_5_server.yaml con el siguiente contenido:

    apiVersion: storage.k8s.io/v1
    kind: StorageClass
    metadata:
      name: hyperdisk-balanced-qwen3-5
    provisioner: pd.csi.storage.gke.io
    parameters:
      type: hyperdisk-balanced
    reclaimPolicy: Delete
    volumeBindingMode: WaitForFirstConsumer
    allowVolumeExpansion: true
    ---
    apiVersion: v1
    kind: PersistentVolumeClaim
    metadata:
      name: hd-claim-qwen3-5
      namespace: default
    spec:
      storageClassName: hyperdisk-balanced-qwen3-5
      accessModes:
        - ReadWriteOnce
      resources:
        requests:
          storage: 1500Gi
    ---
    apiVersion: apps/v1
    kind: Deployment
    metadata:
      name: vllm-qwen3-5
      namespace: default
    spec:
      replicas: 1
      selector:
        matchLabels:
          app: vllm-qwen3-5
      template:
        metadata:
          labels:
            app: vllm-qwen3-5
        spec:
          nodeSelector:
            cloud.google.com/gke-tpu-accelerator: tpu7x
            cloud.google.com/gke-tpu-topology: 2x2x1
          tolerations:
          - key: "cloud.google.com/gke-spot"
            operator: "Equal"
            value: "true"
            effect: "NoSchedule"
          containers:
          - name: vllm-tpu
            image: $VLLM_IMAGE
            command: ["vllm", "serve"]
            args:
            - --host=0.0.0.0
            - --port=8000
            - --model=Qwen/Qwen3.5-397B-A17B-FP8
            - --tokenizer=Qwen/Qwen3.5-397B-A17B-FP8
            - --max-model-len=9236
            - --max-num-batched-tokens=1024
            - --max-num-seqs=64
            - --no-enable-prefix-caching
            - --gpu-memory-utilization=0.9
            - --data-parallel-size=8
            - --tensor-parallel-size=1
            - --async-scheduling
            - --language-model-only
            - --enable-auto-tool-choice
            - --tool-call-parser=qwen3_coder
            - --reasoning-parser=qwen3
            - '--limit-mm-per-prompt={"image": 0, "video": 0}'
            - --kv-cache-dtype=fp8
            - --enable-expert-parallel
            - --download-dir=/data
            - --block-size=128
            env:
            - name: MODEL_IMPL_TYPE
              value: "vllm"
            - name: DP_SCHED_ENABLED
              value: "1"
            - name: ONEHOT_MOE_PERMUTE_THRESHOLD
              value: "32768"
            - name: VLLM_ENGINE_READY_TIMEOUT_S
              value: "9600"
            - name: VLLM_ENGINE_ITERATION_TIMEOUT_S
              value: "1800"
            - name: USE_MOE_SPARSE_CORE
              value: "1"
            - name: TPU_MOE_SKIP_PADDED_TOKENS
              value: "1"
            - name: VLLM_CACHE_ROOT
              value: "/data/vllm_cache"
            - name: VLLM_TORCH_COMPILE_CACHE_DIR
              value: "/data/torch_compile_cache"
            - name: XLA_CACHE_DIR
              value: "/data/xla_cache"
            - name: HF_HOME
              value: /data
            - name: HF_TOKEN
              valueFrom:
                secretKeyRef:
                  name: hf-secret
                  key: hf_api_token
            ports:
            - containerPort: 8000
            resources:
              limits:
                google.com/tpu: '4'
              requests:
                google.com/tpu: '4'
            readinessProbe:
              tcpSocket:
                port: 8000
              initialDelaySeconds: 15
              periodSeconds: 10
            volumeMounts:
            - mountPath: "/data"
              name: data-volume
            - mountPath: /dev/shm
              name: dshm
          volumes:
          - emptyDir:
              medium: Memory
            name: dshm
          - name: data-volume
            persistentVolumeClaim:
              claimName: hd-claim-qwen3-5
    ---
    apiVersion: v1
    kind: Service
    metadata:
      name: vllm-qwen3-5-service
      namespace: default
    spec:
      selector:
        app: vllm-qwen3-5
      type: ClusterIP
      ports:
      - name: http
        protocol: TCP
        port: 8000
        targetPort: 8000
  2. Aplica el manifiesto del servidor de vLLM TorchTPU a tu clúster:

    envsubst '$VLLM_IMAGE' < "qwen3_5_server.yaml" | kubectl apply -f -
  3. Supervisa los registros de implementación del servidor hasta que se complete el inicio de la aplicación:

    kubectl wait --for=jsonpath='{.status.phase}'=Running pod -l app=vllm-qwen3-5 --timeout=900s
    kubectl logs deployment/vllm-qwen3-5 -c vllm-tpu -f

    El inicio inicial del servicio puede tardar entre 15 y 30 minutos, mientras que el Pod descarga los pesos del modelo y lo compila.

    Si usas capacidad interrumpible, como Spot o inicio flexible, es posible que los nodos de porción de TPU tarden más en crearse después de la creación del grupo de nodos, y se pueden interrumpir durante las ejecuciones de la carga de trabajo. Si tus Pods tardan más de lo esperado en iniciarse o tus cargas de trabajo se detienen en el medio, esta podría ser la causa. Para obtener más información, consulta Interrupción de VMs Spot. Para verificar si se crearon los nodos de porción de TPU, consulta Verifica el estado y las condiciones del nodo.

    Espera hasta que los registros informen Application startup complete antes de enviar tráfico:

    (APIServer pid=1) INFO:     Started server process [1]
    (APIServer pid=1) INFO:     Waiting for application startup.
    (APIServer pid=1) INFO:     Application startup complete.
    
  4. Abre una segunda sesión de terminal.

  5. En la segunda sesión de terminal, reenvía el puerto local 8000 a vllm-qwen3-service:

    kubectl port-forward service/vllm-qwen3-5-service 8000:8000

    Deja la redirección de puertos en ejecución en esta sesión de terminal.

Envía solicitudes de inferencia

Regresa a tu primera sesión de terminal y envía una solicitud de finalización de chat al extremo /v1/chat/completions:

curl http://localhost:8000/v1/chat/completions \
    -H "Content-Type: application/json" \
    -d '{
        "model": "Qwen/Qwen3.5-397B-A17B-FP8",
        "messages": [
            {"role": "user", "content": "Hello! Tell me a joke."}
        ],
        "max_tokens": 512,
        "temperature": 0.7
    }' | jq .

Ejecuta una carga de trabajo de comparativas

Para comparar el modelo Qwen3-30B-A3B implementado con una carga de trabajo de entrada de 8,192 tokens y salida de 1,024 tokens (8k/1k) con el paquete de comparativas de InferenceX, completa los siguientes pasos:

  1. Recupera el manifiesto de la comparativa:

    curl -O https://raw.githubusercontent.com/GoogleCloudSamples/AIHypercomputerSamples/refs/heads/main/tpu/inference/qwen3-5-397b/qwen3_5-benchmark-8k1k.yaml

    Como alternativa, crea un archivo llamado qwen3_5-benchmark-8k1k.yaml con el siguiente contenido:

    apiVersion: v1
    kind: Pod
    metadata:
      name: qwen3-5-bench-8k1k
      namespace: default
    spec:
      restartPolicy: OnFailure
      terminationGracePeriodSeconds: 60
      containers:
      - name: vllm-bench
        image: $VLLM_IMAGE
        command: ["/bin/bash", "-c"]
        args:
        - |
          while ! curl http://vllm-qwen3-5-service:8000/ping; do sleep 30 && echo 'Waiting for server...'; done
          apt-get update && apt-get install -y git && \
          git clone https://github.com/SemiAnalysisAI/InferenceX.git /ubench/inferencex && \
          cd /ubench/inferencex && \
          git checkout 89ce6098ef2bc4576a735c43f39c7d972b091cfc && \
          python3 /ubench/inferencex/utils/bench_serving/benchmark_serving.py \
            --backend=vllm \
            --percentile-metrics='ttft,tpot,itl,e2el' \
            --host=vllm-qwen3-5-service \
            --port=8000 \
            --model=Qwen/Qwen3.5-397B-A17B-FP8 \
            --tokenizer=Qwen/Qwen3.5-397B-A17B-FP8 \
            --dataset-name=random \
            --random-input-len=8192 \
            --random-output-len=1024 \
            --random-range-ratio=1.0 \
            --num-prompts=2560 \
            --num-warmups=512 \
            --max-concurrency=256 \
            --ignore-eos
        env:
        - name: HUGGING_FACE_HUB_TOKEN
          valueFrom:
            secretKeyRef:
              key: hf_api_token
              name: hf-secret
  2. Aplica el manifiesto del Pod de comparativas:

    envsubst '$VLLM_IMAGE' < "qwen3_5-benchmark-8k1k.yaml" | kubectl apply -f -
  3. Supervisa los registros de la comparativa:

    kubectl wait --for=jsonpath='{.status.phase}'=Running pod --field-selector=metadata.name=qwen3-5-bench-8k1k --timeout=600s
    kubectl logs -f qwen3-5-bench-8k1k

    La comparativa tarda aproximadamente entre 25 y 30 minutos en completarse.

  4. Una vez que se complete la comparativa, busca un resultado que contenga métricas como las siguientes:

    • Output token throughput
    • Total token throughput
    • TTFT
    • TPOT

Realiza una limpieza

Para evitar que se apliquen cargos adicionales continuos a tu cuenta de Google Cloud por los recursos que se usaron en este instructivo, borra las cargas de trabajo de Kubernetes, el grupo de nodo TPU y el clúster de GKE:

  1. Borra el Deployment, el Service, la PersistentVolumeClaim y el Secret de Hugging Face del servidor de vLLM TorchTPU:

    kubectl delete -f qwen3_5_server.yaml --ignore-not-found || true
    kubectl delete secret hf-secret --ignore-not-found || true
  2. Borra el grupo de nodos TPU v6e y el clúster de GKE:

    gcloud container node-pools delete "${NODE_POOL_NAME}" \
        --cluster="${CLUSTER_NAME}" \
        --location="${REGION}" \
        --project="${PROJECT_ID}" \
        --quiet || true
    
    gcloud container clusters delete "${CLUSTER_NAME}" \
        --location="${REGION}" \
        --project="${PROJECT_ID}" \
        --quiet || true

¿Qué sigue?