Google utiliza tecnología de IA para traducir contenido a tu idioma preferido. Las traducciones realizadas con IA pueden contener errores.

Entrega LLM en TPU de varios hosts con Ray y TorchTPU

Autopilot Standard

En este instructivo, se muestra cómo entregar un modelo de lenguaje grande (LLM) de mezcla de expertos (MoE) Gemma 4 26B-A4B-it con paralelismo de tensores y paralelismo de expertos en porciones de Cloud TPU de varios hosts en Google Kubernetes Engine (GKE) con Ray Serve LLM, vLLM y vLLM-TorchTPU.

Este instructivo está dirigido a ingenieros de aprendizaje automático (AA), administradores y operadores de plataformas, y especialistas en datos y en IA que usan la organización de contenedores de Kubernetes para entregar cargas de trabajo de IA y AA en porciones de TPU distribuidas de varios hosts. Para obtener más información sobre los roles comunes y las tareas de ejemplo a las que se hace referencia en el contenido de Google Cloud , consulta Roles y tareas comunes del usuario de GKE.

Antes de comenzar este instructivo, asegúrate de estar familiarizado con lo siguiente:

Fondo

En esta sección, se describen las tecnologías clave que se usan en este instructivo.

TPU

Las unidades de procesamiento tensorial (TPU) son aceleradores personalizados de Google para cargas de trabajo de aprendizaje automático. Este instructivo admite tanto TPU Trillium (v6e) como Ironwood (TPU7x). En una porción de TPU de varios hosts, varios hosts de TPU se comunican directamente a través de una red de interconexión entre chips (ICI) de alta velocidad, lo que proporciona la memoria de alto ancho de banda (HBM) agregada y la comunicación colectiva de baja latencia necesarias para entregar modelos que superan la capacidad de memoria de un solo chip o host de TPU.

vLLM y TorchTPU

vLLM es un motor de inferencia y entrega de código abierto para LLM. TorchTPU es un backend de PyTorch de código abierto para Cloud TPU que proporciona administración de dispositivos y comunicación colectiva a través de ICI.

vLLM-TorchTPU registra el complemento de la plataforma de TPU con vLLM y compila los pases hacia adelante del modelo de PyTorch para Cloud TPU a través de TorchTPU y el compilador de álgebra lineal acelerada (XLA). Admite el paralelismo de tensores distribuidos y el paralelismo de expertos en porciones de TPU de varios hosts con Ray.

LLM de Ray Serve en GKE

Ray Serve LLM proporciona una entrada HTTP compatible con OpenAI para los extremos /v1/chat/completions, /v1/completions y /v1/models, y organiza las implementaciones de vLLM distribuidas en GKE. Cuando habilitas el complemento de Ray en GKE, el operador de KubeRay y su webhook de TPU insertan metadatos de topología de porción, nombres de host de trabajadores y variables de coordinación de TorchTPU en cada Pod de trabajador TPU de varios hosts.

Ray Serve LLM proporciona las siguientes capacidades para la entrega de TPU de varios hosts:

La ubicación de porciones con reconocimiento de topología reserva una porción de TPU de varios hosts intacta para cada réplica del modelo cuando especificas una topología de TPU en el campo accelerator_config, lo que programa un trabajador de vLLM por dispositivo TPU en toda la porción.
El enrutamiento que reconoce prefijos dirige las solicitudes con instrucciones del sistema compartidas o el historial de conversaciones de varios turnos a la réplica que ya almacena en caché esos bloques de par clave-valor (KV) de prefijo en la HBM de la TPU cuando se entrega en varias réplicas de porciones de TPU.
La transmisión directa de tokens agrupa en lotes las transmisiones de tokens salientes en solicitudes simultáneas y transmite las respuestas directamente desde las réplicas del modelo para reducir el tiempo hasta el primer token (TTFT) y la latencia entre tokens (ITL).
El ajuste de escala automático de réplicas basado en la simultaneidad escala las réplicas de modelos en incrementos de porciones completas según la simultaneidad de las solicitudes en curso.

Objetivos

En este instructivo, se muestra cómo realizar lo siguiente:

Compila una imagen de contenedor que coloque Ray Serve sobre la imagen base precompilada de vLLM-TorchTPU y envíala a Artifact Registry.
Crea un clúster de GKE en modo Autopilot o Standard en TPU Trillium (v6e) o Ironwood (TPU7x).
Implementar un recurso personalizado de RayCluster de varios hosts y verificar la configuración de la topología de TPU
Implementa una aplicación de LLM de Ray Serve para entregar el modelo google/gemma-4-26B-A4B-it con paralelismo de tensores y paralelismo de expertos en una porción de TPU de varios hosts, o implementa el clúster y la aplicación de forma declarativa con un recurso personalizado de RayService.
Enviar solicitudes de descubrimiento de modelos y de finalización de chat compatibles con OpenAI al modelo entregado

Costos

En este instructivo, 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 instructivo, 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

Antes de comenzar, asegúrate de haber realizado las siguientes tareas:

Habilita la API de Artifact Registry y la API de Google Kubernetes Engine.

Habilita las APIs

Si deseas usar Google Cloud CLI para esta tarea, instala y, luego, inicializa gcloud CLI. Si ya instalaste gcloud CLI, ejecuta el comando gcloud components update para obtener la versión más reciente. Es posible que las versiones anteriores de gcloud CLI no admitan la ejecución de los comandos de este documento.
Nota: Para las instalaciones de gcloud CLI existentes, asegúrate de configurar la propiedad compute/region. Si usas principalmente clústeres zonales, configura compute/zone en su lugar. Cuando configuras una ubicación predeterminada, puedes evitar errores en gcloud CLI como el siguiente: One of [--zone, --region] must be supplied: Please specify location. Es posible que debas especificar la ubicación en ciertos comandos si la ubicación de tu clúster difiere de la predeterminada que estableciste.

Verifica que tu proyecto tenga suficiente cuota para la capacidad de TPU Trillium (v6e) o Ironwood (TPU7x) en la zona seleccionada. Para obtener más información, consulta Cuotas de Cloud TPU.
Asegúrate de tener una cuenta de Hugging Face, acepta las condiciones de la licencia en la página del modelo google/gemma-4-26B-A4B-it y genera un token de acceso de Hugging Face con permisos de Read.
Asegúrate de que tu cuenta tenga los siguientes roles de IAM:
- Administrador de Kubernetes Engine (roles/container.admin)
- Administrador de Artifact Registry (roles/artifactregistry.admin)

Prepara el entorno

En este instructivo, usarás Cloud Shell o una estación de trabajo local para administrar recursos alojados en Google Cloud. Cloud Shell viene preinstalado con las herramientas de línea de comandos que se usan en este instructivo, incluidas kubectl, docker, python3 y la CLI de gcloud.

Para configurar tu entorno con Cloud Shell, sigue estos pasos:

En la consola de Google Cloud , haz clic en Activar Cloud Shell en la consola deGoogle Cloud para iniciar una sesión de Cloud Shell. Esta acción inicia una sesión en el panel inferior de la consola de Google Cloud .

Clona el repositorio de kubernetes-engine-samples y cambia al directorio de muestra de este instructivo:

git clone https://github.com/GoogleCloudPlatform/kubernetes-engine-samples.git
cd kubernetes-engine-samples/ai-ml/gke-ray/rayserve/llm/torchtpu

Crea y activa un entorno virtual de Python local con la CLI de Ray instalada:

python3 -m venv ~/ray-env
source ~/ray-env/bin/activate
pip install "ray[default]==2.58.0"

Establece las variables de entorno para tu proyecto y clúster de Google Cloud :
```
export PROJECT_ID=PROJECT_ID
export CLUSTER_NAME=CLUSTER_NAME
export REGION=REGION
export ZONE=ZONE
export ARTIFACT_REGISTRY=ARTIFACT_REGISTRY
export HF_TOKEN=HF_TOKEN
export IMAGE_URI=${REGION}-docker.pkg.dev/${PROJECT_ID}/${ARTIFACT_REGISTRY}/ray-torchtpu-vllm:latest
gcloud config set project ${PROJECT_ID}
```
Reemplaza lo siguiente:
- PROJECT_ID: Es el ID del proyecto de Google Cloud .
- CLUSTER_NAME: Es el nombre del clúster de GKE, como ray-serve-torchtpu-cluster.
- REGION: Es la región en la que se encuentran el repositorio de Artifact Registry y la capacidad de TPU, como us-east5.
- ZONE: Es la zona en la que está disponible tu capacidad de TPU, como us-east5-b. Para obtener más información, consulta Valida la disponibilidad de TPU en GKE.
- ARTIFACT_REGISTRY: Es el nombre de tu repositorio de Docker de Artifact Registry, como ray-torchtpu.
- HF_TOKEN: Es tu token de acceso de Hugging Face con permisos de Read.

Compila la imagen de contenedor personalizada

Compila una imagen de contenedor que instale el paquete ray[default,serve] sobre la imagen base precompilada de vLLM-TorchTPU, que incluye TorchTPU, vLLM y el complemento vLLM-TorchTPU.

Revisa el archivo Dockerfile en el directorio kubernetes-engine-samples/ai-ml/gke-ray/rayserve/llm/torchtpu:

FROM vllm/vllm-torchtpu:v0.29.0

ENV HF_HOME=/dev/shm/huggingface
ENV RAY_SERVE_DEFAULT_HTTP_HOST=0.0.0.0

# Install wget for KubeRay health probes and pin Ray Serve and transformers
# versions compatible with vLLM 0.29.0 and multi-host TPUs.
RUN apt-get update && apt-get install -y --no-install-recommends wget \
    && rm -rf /var/lib/apt/lists/* \
    && pip install --no-cache-dir \
       "ray[default,serve] @ https://s3-us-west-2.amazonaws.com/ray-wheels/master/6b2b414bb7a5b4e2705ed9f4fdab9787b92f3b08/ray-3.0.0.dev0-cp312-cp312-manylinux2014_x86_64.whl" \
       "transformers>=5.18.0"

La entrega de TPU de varios hosts con Ray Serve LLM requiere la versión 2.56.0 o posterior de Ray en TPU Trillium (v6e) y la versión 2.58.0 o posterior de Ray para la programación por chiplet en Ironwood (TPU7x).

Crea el repositorio de Docker de Artifact Registry, compila la imagen de contenedor y envíala a Artifact Registry:

gcloud artifacts repositories create ${ARTIFACT_REGISTRY} \
    --repository-format=docker \
    --location=${REGION} \
    --description="Ray and vLLM-TorchTPU container images"

gcloud auth configure-docker ${REGION}-docker.pkg.dev,us-docker.pkg.dev
docker build -t ${IMAGE_URI} .
docker push ${IMAGE_URI}

Crea un clúster de GKE

Puedes entregar LLM en porciones de TPU de varios hosts en un clúster de GKE Autopilot o Standard con el complemento administrado Ray operator habilitado.

Crea tu clúster de GKE y configura las credenciales para la herramienta de línea de comandos de kubectl:
Autopilot
1. Crea un clúster regional de GKE Autopilot con el complemento de operador de Ray habilitado:
  
  gcloud container clusters create-auto ${CLUSTER_NAME} \ --project=${PROJECT_ID} \ --region=${REGION} \ --cluster-version=1.36 \ --enable-ray-operator
  
  Cuando implementas el recurso personalizado RayCluster o RayService en un paso posterior, GKE Autopilot aprovisiona automáticamente el nodo principal de CPU y la porción de TPU de varios hosts que coinciden con tus solicitudes de recursos de Pod y selectores de nodos.
2. Configura la herramienta de línea de comandos de kubectl para que se comunique con tu clúster:
  
  gcloud container clusters get-credentials ${CLUSTER_NAME} \ --project=${PROJECT_ID} \ --region=${REGION}
Estándar
1. Crea un clúster zonal de GKE Standard con un grupo de nodos de CPU para alojar el Pod principal de Ray, la transmisión de imágenes y GKE Dataplane V2 habilitados:
  
  gcloud container clusters create ${CLUSTER_NAME} \ --project=${PROJECT_ID} \ --zone=${ZONE} \ --cluster-version=1.36 \ --machine-type=n2-standard-8 \ --num-nodes=1 \ --addons=RayOperator \ --enable-image-streaming \ --enable-dataplane-v2
  Nota: Si autoadministras KubeRay de código abierto en lugar de usar el complemento RayOperator de GKE, instala la versión 1.4.0 o posterior de KubeRay y el webhook de TPU de KubeRay para la inserción de variable de entorno de TorchTPU.
2. Configura la herramienta de línea de comandos de kubectl para que se comunique con tu clúster:
  
  gcloud container clusters get-credentials ${CLUSTER_NAME} \ --project=${PROJECT_ID} \ --zone=${ZONE}
3. Crea un grupo de nodos de porciones de TPU de varios hosts de 16 dispositivos con una topología 4x4 en cuatro hosts en TPU Trillium (v6e) o una topología 2x2x2 en dos hosts en Ironwood (TPU7x):
  
  Trillium (v6e)
  gcloud container node-pools create tpu-multi-host-pool \ --project=${PROJECT_ID} \ --cluster=${CLUSTER_NAME} \ --zone=${ZONE} \ --machine-type=ct6e-standard-4t \ --tpu-topology=4x4 \ --num-nodes=4
  
  Ironwood (TPU7x)
  gcloud container node-pools create tpu-multi-host-pool \ --project=${PROJECT_ID} \ --cluster=${CLUSTER_NAME} \ --zone=${ZONE} \ --machine-type=tpu7x-standard-4t \ --tpu-topology=2x2x2 \ --num-nodes=2
  
  Sugerencia: Para consumir una reserva de TPU específica, agrega las marcas --reservation-affinity=specific y --reservation=RESERVATION_NAME al comando node-pools create, donde RESERVATION_NAME es el nombre de tu reserva de TPU. Para usar VMs Spot, agrega la marca --spot.
Crea un Secret de Kubernetes con tu token de Hugging Face para que los Pods de Ray puedan autenticarse con Hugging Face cuando descarguen los pesos del modelo:
```
kubectl create secret generic hf-secret \
    --from-literal=hf_api_token=${HF_TOKEN}
```

Implementa el recurso personalizado RayCluster de varios hosts

Implementar un recurso personalizado de RayCluster que aprovisione un Pod principal de CPU, un volumen /dev/shm en memoria para el almacén de objetos de memoria compartida de Ray y un grupo de trabajador TPU de varios hosts

Revisa el manifiesto de RayCluster para tu versión de TPU en el repositorio de muestra:

Trillium (v6e)

apiVersion: ray.io/v1
kind: RayCluster
metadata:
  name: ray-serve-torchtpu
spec:
  headGroupSpec:
    rayStartParams:
      num-cpus: "0"
    template:
      spec:
        containers:
        - name: ray-head
          image: ${IMAGE_URI}
          env:
          - name: HF_TOKEN
            valueFrom:
              secretKeyRef:
                name: hf-secret
                key: hf_api_token
          ports:
          - containerPort: 6379
            name: gcs
          - containerPort: 8265
            name: dashboard
          - containerPort: 10001
            name: client
          - containerPort: 8000
            name: serve
          resources:
            limits:
              cpu: "4"
              memory: "16Gi"
            requests:
              cpu: "4"
              memory: "16Gi"
  workerGroupSpecs:
  - groupName: v6e-16
    replicas: 1
    numOfHosts: 4
    template:
      spec:
        containers:
        - name: ray-worker
          image: ${IMAGE_URI}
          env:
          - name: HF_TOKEN
            valueFrom:
              secretKeyRef:
                name: hf-secret
                key: hf_api_token
          resources:
            limits:
              cpu: "64"
              memory: "200Gi"
              google.com/tpu: "4"
            requests:
              cpu: "64"
              memory: "200Gi"
              google.com/tpu: "4"
          volumeMounts:
          - mountPath: /dev/shm
            name: dshm
        volumes:
        - name: dshm
          emptyDir:
            medium: Memory
            sizeLimit: 100Gi
        nodeSelector:
          cloud.google.com/gke-tpu-accelerator: tpu-v6e-slice
          cloud.google.com/gke-tpu-topology: 4x4

Ironwood (TPU7x)

apiVersion: ray.io/v1
kind: RayCluster
metadata:
  name: ray-serve-torchtpu
spec:
  headGroupSpec:
    rayStartParams:
      num-cpus: "0"
    template:
      spec:
        containers:
        - name: ray-head
          image: ${IMAGE_URI}
          env:
          - name: HF_TOKEN
            valueFrom:
              secretKeyRef:
                name: hf-secret
                key: hf_api_token
          - name: RAY_TPU_RESOURCE_PER_CHIP
            value: "2"
          ports:
          - containerPort: 6379
            name: gcs
          - containerPort: 8265
            name: dashboard
          - containerPort: 10001
            name: client
          - containerPort: 8000
            name: serve
          resources:
            limits:
              cpu: "4"
              memory: "16Gi"
            requests:
              cpu: "4"
              memory: "16Gi"
  workerGroupSpecs:
  - groupName: tpu7x-16
    replicas: 1
    numOfHosts: 2
    resources:
      # Advertises 8 logical TPU cores (2 chiplets per physical chip) per host.
      TPU: "8"
    template:
      spec:
        containers:
        - name: ray-worker
          image: ${IMAGE_URI}
          env:
          - name: HF_TOKEN
            valueFrom:
              secretKeyRef:
                name: hf-secret
                key: hf_api_token
          - name: RAY_TPU_RESOURCE_PER_CHIP
            value: "2"
          resources:
            limits:
              cpu: "64"
              memory: "250Gi"
              google.com/tpu: "4"
            requests:
              cpu: "64"
              memory: "250Gi"
              google.com/tpu: "4"
          volumeMounts:
          - mountPath: /dev/shm
            name: dshm
        volumes:
        - name: dshm
          emptyDir:
            medium: Memory
            sizeLimit: 128Gi
        nodeSelector:
          cloud.google.com/gke-tpu-accelerator: tpu7x
          cloud.google.com/gke-tpu-topology: 2x2x2

En Ironwood (TPU7x), cada chip físico contiene dos chiplets, lo que proporciona ocho núcleos de TPU lógicos por host de cuatro chips. El manifiesto establece el valor del campo TPU en "8" en la sección resources y establece la variable de entorno RAY_TPU_RESOURCE_PER_CHIP en "2" para anunciar y programar ocho recursos lógicos de TPU por host.

Aplica el manifiesto para crear el recurso personalizado de RayCluster:

Trillium (v6e)

envsubst < ray-cluster.tpu-v6e-multihost.yaml | kubectl apply -f -

Ironwood (TPU7x)

envsubst < ray-cluster.tpu-v7x-multihost.yaml | kubectl apply -f -

Espera hasta que el Pod principal y todos los Pods de trabajador TPU alcancen el estado Running y 1/1 READY:

kubectl get pods -l ray.io/cluster=ray-serve-torchtpu

El resultado es similar a lo siguiente:

Trillium (v6e)

NAME                                     READY   STATUS    RESTARTS   AGE
ray-serve-torchtpu-head-5cj7g            1/1     Running   0          35s
ray-serve-torchtpu-v6e-16-worker-6x8cg   1/1     Running   0          35s
ray-serve-torchtpu-v6e-16-worker-79j5l   1/1     Running   0          35s
ray-serve-torchtpu-v6e-16-worker-psxbl   1/1     Running   0          35s
ray-serve-torchtpu-v6e-16-worker-rz7dv   1/1     Running   0          35s

Ironwood (TPU7x)

NAME                                       READY   STATUS    RESTARTS   AGE
ray-serve-torchtpu-head-w7b2v              1/1     Running   0          2m30s
ray-serve-torchtpu-tpu7x-16-worker-6zlnj   1/1     Running   0          2m30s
ray-serve-torchtpu-tpu7x-16-worker-cb4zz   1/1     Running   0          2m30s

Verifica que el webhook de TPU de KubeRay haya insertado la porción de varios hosts, la coordinación de trabajadores y las variables de entorno de TorchTPU en los Pods de trabajador:

kubectl get pods -l ray.io/node-type=worker \
    -o=jsonpath='{range .items[0].spec.containers[0].env[*]}{.name}={.value}{"\n"}{end}' \
    | grep -E "TPU_WORKER_ID|TPU_WORKER_HOSTNAMES|TPU_NAME|TORCH_TPU_TOPOLOGY|TORCH_TPU_SLICEBUILDER_ADDRESSES"

El resultado es similar a lo siguiente:

Trillium (v6e)

TPU_WORKER_HOSTNAMES=v6e-16-0-0.ray-serve-torchtpu-headless,v6e-16-0-1.ray-serve-torchtpu-headless,v6e-16-0-2.ray-serve-torchtpu-headless,v6e-16-0-3.ray-serve-torchtpu-headless
TORCH_TPU_TOPOLOGY=4,4,1
TORCH_TPU_SLICEBUILDER_ADDRESSES=v6e-16-0-0.ray-serve-torchtpu-headless:8471,v6e-16-0-0.ray-serve-torchtpu-headless:8472,v6e-16-0-0.ray-serve-torchtpu-headless:8473,v6e-16-0-0.ray-serve-torchtpu-headless:8474,v6e-16-0-1.ray-serve-torchtpu-headless:8471,v6e-16-0-1.ray-serve-torchtpu-headless:8472,v6e-16-0-1.ray-serve-torchtpu-headless:8473,v6e-16-0-1.ray-serve-torchtpu-headless:8474,v6e-16-0-2.ray-serve-torchtpu-headless:8471,v6e-16-0-2.ray-serve-torchtpu-headless:8472,v6e-16-0-2.ray-serve-torchtpu-headless:8473,v6e-16-0-2.ray-serve-torchtpu-headless:8474,v6e-16-0-3.ray-serve-torchtpu-headless:8471,v6e-16-0-3.ray-serve-torchtpu-headless:8472,v6e-16-0-3.ray-serve-torchtpu-headless:8473,v6e-16-0-3.ray-serve-torchtpu-headless:8474
TPU_WORKER_ID=1
TPU_NAME=v6e-16-0

Ironwood (TPU7x)

TPU_WORKER_HOSTNAMES=tpu7x-16-0-0.ray-serve-torchtpu-headless,tpu7x-16-0-1.ray-serve-torchtpu-headless
TORCH_TPU_TOPOLOGY=2,2,2,2
TORCH_TPU_SLICEBUILDER_ADDRESSES=tpu7x-16-0-0.ray-serve-torchtpu-headless:8471,tpu7x-16-0-0.ray-serve-torchtpu-headless:8472,tpu7x-16-0-0.ray-serve-torchtpu-headless:8473,tpu7x-16-0-0.ray-serve-torchtpu-headless:8474,tpu7x-16-0-0.ray-serve-torchtpu-headless:8475,tpu7x-16-0-0.ray-serve-torchtpu-headless:8476,tpu7x-16-0-0.ray-serve-torchtpu-headless:8477,tpu7x-16-0-0.ray-serve-torchtpu-headless:8478,tpu7x-16-0-1.ray-serve-torchtpu-headless:8471,tpu7x-16-0-1.ray-serve-torchtpu-headless:8472,tpu7x-16-0-1.ray-serve-torchtpu-headless:8473,tpu7x-16-0-1.ray-serve-torchtpu-headless:8474,tpu7x-16-0-1.ray-serve-torchtpu-headless:8475,tpu7x-16-0-1.ray-serve-torchtpu-headless:8476,tpu7x-16-0-1.ray-serve-torchtpu-headless:8477,tpu7x-16-0-1.ray-serve-torchtpu-headless:8478
TPU_WORKER_ID=1
TPU_NAME=tpu7x-16-0

Puedes inspeccionar los nodos trabajador TPU y sus etiquetas de ray.io/tpu-* detectadas automáticamente en la pestaña Clúster del panel de Ray.

Reenvía los puertos locales al panel de Ray en el puerto 8265 y al extremo HTTP de Ray Serve en el puerto 8000:
```
kubectl port-forward svc/ray-serve-torchtpu-head-svc 8265:8265 8000:8000 >/dev/null 2>&1 &
```

Crea e implementa la aplicación de LLM de Ray Serve

Implementa la aplicación de LLM de Ray Serve en tu recurso personalizado de RayCluster para entregar el modelo google/gemma-4-26B-A4B-it en la porción de TPU de varios hosts.

Revisa la secuencia de comandos serve_vllm_torchtpu.py en el repositorio de ejemplo:

from ray.serve.deployment import Application
from ray.serve.llm import (
    LLMConfig,
    LLMServingArgs,
    ModelLoadingConfig,
    build_openai_app,
)


def app(args: dict[str, str]) -> Application:
    llm_config = LLMConfig(
        model_loading_config=ModelLoadingConfig(
            model_id=args.get("model_id", "google/gemma-4-26B-A4B-it"),
        ),
        accelerator_type=args.get("accelerator_type", "TPU-V6E"),
        accelerator_config={
            "kind": "tpu",
            "topology": args.get("topology", "4x4"),
        },
        placement_group_config={
            "bundle_per_worker": {"TPU": 1},
        },
        engine_kwargs={
            "tensor_parallel_size": int(args.get("tensor_parallel_size", "16")),
            "enable_expert_parallel": True,
            "enable_prefix_caching": True,
            "language_model_only": True,
            "disable_hybrid_kv_cache_manager": True,
            "max_model_len": int(args.get("max_model_len", "8192")),
            "max_num_batched_tokens": int(
                args.get("max_num_batched_tokens", "4096")
            ),
        },
        experimental_configs={
            "stream_batching_interval_ms": 10,
        },
        runtime_env={
            "env_vars": {
                "TPU_MULTIHOST_BACKEND": "ray",
                # Execute Gemma 4 routed expert layers on TPU TensorCores instead of SparseCore.
                "USE_MOE_SPARSE_CORE": "0",
                "VLLM_RAY_EXTRA_ENV_VARS_TO_COPY": "USE_MOE_SPARSE_CORE",
            }
        },
    )
    return build_openai_app(LLMServingArgs(llm_configs=[llm_config]))

La secuencia de comandos configura el objeto LLMConfig con los siguientes parámetros para la entrega de TPU de varios hosts:

Los campos accelerator_config y placement_group_config reservan un grupo de posiciones en una sola porción de TPU y asignan un núcleo de TPU por trabajador de vLLM en toda la porción.
El campo engine_kwargs fragmenta el modelo en los 16 núcleos de TPU con el paralelismo de tensores y el paralelismo de expertos, y habilita el almacenamiento en caché de prefijos. También inhabilita las entradas multimodales para reservar la HBM de la TPU para la inferencia de texto, usa un diseño de caché de KV unificado para las capas de atención de Gemma 4 y limita la longitud de contexto a 8,192 tokens y el tamaño del lote de precompletado a 4,096 tokens.
El campo experimental_configs agrupa en lotes los flujos de tokens salientes cada 10 ms para reducir la sobrecarga de transmisión en las solicitudes simultáneas.
El campo runtime_env habilita el backend de varios hosts de Ray para vLLM-TorchTPU (TPU_MULTIHOST_BACKEND: "ray"), ejecuta capas de expertos de Gemma 4 MoE en TensorCores de TPU en lugar de SparseCore y propaga ese parámetro de configuración a los 16 actores trabajadores.

Implementa la aplicación de Ray Serve en el recurso personalizado de RayCluster con el comando serve run. El comando transmite los registros de implementación hasta que se terminan de cargar los pesos del modelo, se completa la compilación del grafo de XLA en los 16 dispositivos TPU y la aplicación de Serve alcanza el estado RUNNING:

Trillium (v6e)

ray job submit \
    --address http://localhost:8265 \
    --working-dir . \
    --runtime-env-json '{"excludes": ["ray-env", ".git", "*.whl"]}' \
    -- serve run --non-blocking serve_vllm_torchtpu:app accelerator_type=TPU-V6E topology=4x4

Ironwood (TPU7x)

ray job submit \
    --address http://localhost:8265 \
    --working-dir . \
    --runtime-env-json '{"excludes": ["ray-env", ".git", "*.whl"]}' \
    -- serve run --non-blocking serve_vllm_torchtpu:app accelerator_type=TPU-V7X topology=2x2x2

El resultado es similar a lo siguiente:

Trillium (v6e)

(ServeController pid=1637) INFO 2026-10-13 10:23:05,160 controller 1637 -- Deploying new version of Deployment(name='LLMServer:google--gemma-4-26B-A4B-it', app='default') (initial target replicas: 1).
(ServeController pid=1637) INFO 2026-10-13 10:23:05,161 controller 1637 -- Deploying new version of Deployment(name='OpenAiIngress', app='default') (initial target replicas: 1).
(ServeReplica:default:LLMServer:google--gemma-4-26B-A4B-it pid=515, ip=10.72.3.21) INFO 10-13 10:24:06 [__init__.py:318] TPU info: node_name=v6e-16-0 | tpu_type=v6e-16 | worker_id=1 | num_chips=4 | num_cores_per_chip=1
(ServeReplica:default:LLMServer:google--gemma-4-26B-A4B-it pid=515, ip=10.72.3.21) INFO 10-13 10:24:06 [__init__.py:238] Platform plugin torchtpu is activated
(ServeReplica:default:LLMServer:google--gemma-4-26B-A4B-it pid=515, ip=10.72.3.21) (EngineCore pid=1041) INFO 10-13 10:24:08 [core.py:112] Initializing a V1 LLM engine (v0.22.1) with config: model='google/gemma-4-26B-A4B-it', speculative_config=None, tokenizer='google/gemma-4-26B-A4B-it', skip_tokenizer_init=False, tokenizer_mode=auto, revision=None, tokenizer_revision=None, trust_remote_code=False, dtype=torch.bfloat16, max_seq_len=8192, download_dir=None, load_format=auto, tensor_parallel_size=16, pipeline_parallel_size=1, data_parallel_size=1, decode_context_parallel_size=1, dcp_comm_backend=ag_rs, disable_custom_all_reduce=True, quantization=None, quantization_config=None, enforce_eager=False, enable_return_routed_experts=False, kv_cache_dtype=auto, device_config=tpu, ...
(ServeReplica:default:LLMServer:google--gemma-4-26B-A4B-it pid=515, ip=10.72.3.21) (EngineCore pid=1041) (RayWorkerWrapper pid=1017, ip=10.72.14.21) INFO 10-13 10:24:30 [expert_map_manager.py:245] [EP Rank 4/16] Expert parallelism is enabled. Expert placement strategy: linear. Local/global number of experts: 8/128. Experts local to global index map: 0->32, 1->33, 2->34, 3->35, 4->36, 5->37, 6->38, 7->39.
(ServeReplica:default:LLMServer:google--gemma-4-26B-A4B-it pid=515, ip=10.72.3.21) (EngineCore pid=1041) (RayWorkerWrapper pid=856, ip=10.72.6.21) INFO 10-13 10:27:55 [default_loader.py:397] Loading weights took 12.76 seconds
(ServeReplica:default:LLMServer:google--gemma-4-26B-A4B-it pid=515, ip=10.72.3.21) (EngineCore pid=1041) (RayWorkerWrapper pid=856, ip=10.72.6.21) INFO 10-13 10:32:53 [tpu_worker.py:376] Memory statistics | total_hbm_limit_gb=31.25GiB | total_hbm_limit_cap_gb=28.75GiB | total_hbm_used_gb=3.69GiB | kv_cache_headroom_gb=5.0GiB | kv_connector_hbm_reserve_gb=0.0GiB | total_hbm_avail_gb=20.06GiB
(ServeReplica:default:LLMServer:google--gemma-4-26B-A4B-it pid=515, ip=10.72.3.21) (EngineCore pid=1041) INFO 10-13 10:32:54 [kv_cache_utils.py:1733] GPU KV cache size: 600,832 tokens
(ServeReplica:default:LLMServer:google--gemma-4-26B-A4B-it pid=515, ip=10.72.3.21) (EngineCore pid=1041) INFO 10-13 10:32:54 [kv_cache_utils.py:1734] Maximum concurrency for 8,192 tokens per request: 73.34x
(ServeReplica:default:LLMServer:google--gemma-4-26B-A4B-it pid=515, ip=10.72.3.21) INFO 2026-10-13 10:36:59,851 default_LLMServer:google--gemma-4-26B-A4B-it b2grwv9c -- Started vLLM engine.
INFO 2026-10-13 10:37:05,997 serve 1456 -- Application 'default' is ready at http://127.0.0.1:8000/.
2026-10-13 10:37:11,447 SUCC cli.py:66 -- ------------------------------------------
2026-10-13 10:37:11,447 SUCC cli.py:67 -- Job 'raysubmit_94DTwVzR2pvGQVwR' succeeded
2026-10-13 10:37:11,447 SUCC cli.py:68 -- ------------------------------------------

Ironwood (TPU7x)

(ServeController pid=2042) INFO 2026-10-13 10:24:38,092 controller 2042 -- Deploying new version of Deployment(name='LLMServer:google--gemma-4-26B-A4B-it', app='default') (initial target replicas: 1).
(ServeController pid=2042) INFO 2026-10-13 10:24:38,093 controller 2042 -- Deploying new version of Deployment(name='OpenAiIngress', app='default') (initial target replicas: 1).
(ServeReplica:default:LLMServer:google--gemma-4-26B-A4B-it pid=316, ip=172.16.7.5) INFO 10-13 10:24:49 [__init__.py:318] TPU info: node_name=tpu7x-16-0 | tpu_type=tpu7x-16 | worker_id=0 | num_chips=8 | num_cores_per_chip=2
(ServeReplica:default:LLMServer:google--gemma-4-26B-A4B-it pid=316, ip=172.16.7.5) INFO 10-13 10:24:49 [__init__.py:238] Platform plugin torchtpu is activated
(ServeReplica:default:LLMServer:google--gemma-4-26B-A4B-it pid=316, ip=172.16.7.5) (EngineCore pid=719) INFO 10-13 10:25:23 [core.py:112] Initializing a V1 LLM engine (v0.22.1) with config: model='google/gemma-4-26B-A4B-it', speculative_config=None, tokenizer='google/gemma-4-26B-A4B-it', skip_tokenizer_init=False, tokenizer_mode=auto, revision=None, tokenizer_revision=None, trust_remote_code=False, dtype=torch.bfloat16, max_seq_len=8192, download_dir=None, load_format=auto, tensor_parallel_size=16, pipeline_parallel_size=1, data_parallel_size=1, decode_context_parallel_size=1, dcp_comm_backend=ag_rs, disable_custom_all_reduce=True, quantization=None, quantization_config=None, enforce_eager=False, enable_return_routed_experts=False, kv_cache_dtype=auto, device_config=tpu, ...
(ServeReplica:default:LLMServer:google--gemma-4-26B-A4B-it pid=316, ip=172.16.7.5) (EngineCore pid=719) (RayWorkerWrapper pid=602, ip=172.16.10.6) INFO 10-13 10:25:44 [expert_map_manager.py:245] [EP Rank 8/16] Expert parallelism is enabled. Expert placement strategy: linear. Local/global number of experts: 8/128. Experts local to global index map: 0->64, 1->65, 2->66, 3->67, 4->68, 5->69, 6->70, 7->71.
(ServeReplica:default:LLMServer:google--gemma-4-26B-A4B-it pid=316, ip=172.16.7.5) (EngineCore pid=719) (RayWorkerWrapper pid=762) INFO 10-13 10:29:06 [default_loader.py:397] Loading weights took 11.69 seconds
(ServeReplica:default:LLMServer:google--gemma-4-26B-A4B-it pid=316, ip=172.16.7.5) (EngineCore pid=719) (RayWorkerWrapper pid=604, ip=172.16.10.6) INFO 10-13 10:34:51 [tpu_worker.py:376] Memory statistics | total_hbm_limit_gb=94.75GiB | total_hbm_limit_cap_gb=87.17GiB | total_hbm_used_gb=3.74GiB | kv_cache_headroom_gb=5.0GiB | kv_connector_hbm_reserve_gb=0.0GiB | total_hbm_avail_gb=78.43GiB
(ServeReplica:default:LLMServer:google--gemma-4-26B-A4B-it pid=316, ip=172.16.7.5) (EngineCore pid=719) INFO 10-13 10:34:58 [kv_cache_utils.py:1733] GPU KV cache size: 2,349,568 tokens
(ServeReplica:default:LLMServer:google--gemma-4-26B-A4B-it pid=316, ip=172.16.7.5) (EngineCore pid=719) INFO 10-13 10:34:58 [kv_cache_utils.py:1734] Maximum concurrency for 8,192 tokens per request: 286.81x
(ServeReplica:default:LLMServer:google--gemma-4-26B-A4B-it pid=316, ip=172.16.7.5) INFO 2026-10-13 10:40:13,586 default_LLMServer:google--gemma-4-26B-A4B-it a15k5cv3 -- Started vLLM engine.
INFO 2026-10-13 10:40:18,415 serve 1953 -- Application 'default' is ready at http://127.0.0.1:8000/.
2026-10-13 10:40:24,372 SUCC cli.py:66 -- ------------------------------------------
2026-10-13 10:40:24,372 SUCC cli.py:67 -- Job 'raysubmit_mzWQDQGDM9QSgcMN' succeeded
2026-10-13 10:40:24,372 SUCC cli.py:68 -- ------------------------------------------

Consulta el modelo entregado

Después de que la aplicación de Ray Serve alcance el estado RUNNING, verifica el extremo del modelo registrado, envía una solicitud /v1/chat/completions compatible con OpenAI a través del puerto 8000 y, de manera opcional, realiza una comparativa de la capacidad de procesamiento y la latencia.

Verifica que el modelo esté registrado en el extremo /v1/models:

curl http://localhost:8000/v1/models

El resultado es similar a lo siguiente:

{"data":[{"id":"google/gemma-4-26B-A4B-it","object":"model","owned_by":"organization-owner","permission":[],"metadata":{"model_id":"google/gemma-4-26B-A4B-it","input_modality":"text","max_request_context_length":8192}}],"object":"list"}

Envía una instrucción de finalización de chat al extremo /v1/chat/completions:

curl -X POST http://localhost:8000/v1/chat/completions \
    -H "Content-Type: application/json" \
    -d '{
        "model": "google/gemma-4-26B-A4B-it",
        "messages": [
            {"role": "system", "content": "You are a helpful AI assistant."},
            {"role": "user", "content": "Explain in one sentence how tensor parallelism and expert parallelism accelerate Mixture-of-Experts LLM inference on TPUs."}
        ],
        "max_tokens": 128,
        "temperature": 0.0
    }'

El resultado en Ironwood (TPU7x) es similar al siguiente:

{"id":"chatcmpl-ec30e61c-011d-49cd-989a-9ae9f59448a8","object":"chat.completion","created":1791887860,"model":"google/gemma-4-26B-A4B-it","choices":[{"index":0,"message":{"role":"assistant","content":"Tensor parallelism accelerates inference by partitioning large weight matrices across TPU cores to reduce per-device memory and computation, while expert parallelism distributes the specialized MoE layers across different devices to enable massive model scaling and parallelize the processing of different tokens through different experts.","refusal":null,"annotations":null,"audio":null,"function_call":null,"tool_calls":[],"reasoning":null},"logprobs":null,"finish_reason":"stop","stop_reason":106,"token_ids":null,"routed_experts":null}],"service_tier":null,"system_fingerprint":"vllm-0.22.1-tp16-ep-aaf76e1a","usage":{"prompt_tokens":49,"total_tokens":100,"completion_tokens":51,"prompt_tokens_details":null},"prompt_logprobs":null,"prompt_token_ids":null,"prompt_text":null,"kv_transfer_params":null}

Opcional: Realiza una comparativa de la capacidad de procesamiento y la latencia de entrega con el comando vllm bench serve. Con el siguiente comando, se ejecuta el cliente generador de carga HTTP vllm bench serve desde un Pod de trabajador TPU de Ray y se envían 64 solicitudes simultáneas al proxy HTTP de Ray Serve en el puerto 8000:

WORKER_POD=$(kubectl get pods -l ray.io/node-type=worker -o jsonpath='{.items[0].metadata.name}')
kubectl exec -it ${WORKER_POD} -- vllm bench serve \
    --backend openai-chat \
    --endpoint /v1/chat/completions \
    --model google/gemma-4-26B-A4B-it \
    --dataset-name random \
    --random-input-len 1024 \
    --random-output-len 512 \
    --max-concurrency 32 \
    --num-prompts 64 \
    --num-warmups 4 \
    --temperature 0.0 \
    --ignore-eos \
    --host ray-serve-torchtpu-head-svc \
    --port 8000

El resultado es similar a lo siguiente:

Trillium (v6e)

============ Serving Benchmark Result ============
Successful requests:                     64
Failed requests:                         0
Maximum request concurrency:             32
Benchmark duration (s):                  8.41
Total input tokens:                      66392
Total generated tokens:                  32768
Request throughput (req/s):              7.61
Output token throughput (tok/s):         3894.26
Peak output token throughput (tok/s):    4409.00
Peak concurrent requests:                64.00
Total token throughput (tok/s):          11784.50
---------------Time to First Token----------------
Mean TTFT (ms):                          414.57
Median TTFT (ms):                        348.94
P99 TTFT (ms):                           640.26
-----Time per Output Token (excl. 1st token)------
Mean TPOT (ms):                          7.37
Median TPOT (ms):                        7.29
P99 TPOT (ms):                           8.24
---------------Inter-token Latency----------------
Mean ITL (ms):                           7.37
Median ITL (ms):                         0.00
P99 ITL (ms):                            25.94
==================================================

Ironwood (TPU7x)

============ Serving Benchmark Result ============
Successful requests:                     64
Failed requests:                         0
Maximum request concurrency:             32
Benchmark duration (s):                  9.65
Total input tokens:                      66392
Total generated tokens:                  32768
Request throughput (req/s):              6.64
Output token throughput (tok/s):         3397.33
Peak output token throughput (tok/s):    4574.00
Peak concurrent requests:                64.00
Total token throughput (tok/s):          10280.73
---------------Time to First Token----------------
Mean TTFT (ms):                          654.11
Median TTFT (ms):                        522.30
P99 TTFT (ms):                           1288.23
-----Time per Output Token (excl. 1st token)------
Mean TPOT (ms):                          8.13
Median TPOT (ms):                        8.08
P99 TPOT (ms):                           9.69
---------------Inter-token Latency----------------
Mean ITL (ms):                           8.12
Median ITL (ms):                         9.31
P99 ITL (ms):                            26.85
==================================================

Implementa de forma declarativa con un recurso personalizado de RayService

Para las implementaciones de producción, puedes administrar el clúster de Ray de varios hosts y la aplicación de LLM de Ray Serve como un único recurso personalizado declarativo de RayService en lugar de enviar un trabajo de forma manual. Debido a que el compilador de LLM de Ray Serve acepta la configuración del modelo directamente en la configuración de Serve, no necesitas una secuencia de comandos de Python independiente. KubeRay aprovisiona el clúster, implementa la aplicación de Serve, supervisa el estado de las réplicas y crea un Service ClusterIP de Kubernetes en el puerto 8000 cuando la aplicación está lista para entregar tráfico.

Si el recurso personalizado RayCluster de la sección anterior aún se está ejecutando, detén la redirección de puertos y borra el recurso personalizado RayCluster para liberar la porción de TPU:
```
pkill -f "kubectl port-forward" || true
kubectl delete raycluster ray-serve-torchtpu --ignore-not-found
```

Revisa el manifiesto de RayService para tu versión de TPU en el repositorio de muestra:

Trillium (v6e)

apiVersion: ray.io/v1
kind: RayService
metadata:
  name: ray-serve-torchtpu
spec:
  serveConfigV2: |
    applications:
    - name: default
      route_prefix: /
      import_path: ray.serve.llm:build_openai_app
      args:
        llm_configs:
        - model_loading_config:
            model_id: google/gemma-4-26B-A4B-it
          accelerator_type: TPU-V6E
          accelerator_config:
            kind: tpu
            topology: 4x4
          placement_group_config:
            bundle_per_worker:
              TPU: 1
          engine_kwargs:
            tensor_parallel_size: 16
            enable_expert_parallel: true
            enable_prefix_caching: true
            language_model_only: true
            disable_hybrid_kv_cache_manager: true
            max_model_len: 8192
            max_num_batched_tokens: 4096
          experimental_configs:
            stream_batching_interval_ms: 10
          runtime_env:
            env_vars:
              TPU_MULTIHOST_BACKEND: "ray"
              # Execute Gemma 4 routed expert layers on TPU TensorCores instead of SparseCore.
              USE_MOE_SPARSE_CORE: "0"
              VLLM_RAY_EXTRA_ENV_VARS_TO_COPY: "USE_MOE_SPARSE_CORE"
  rayClusterConfig:
    headGroupSpec:
      rayStartParams:
        num-cpus: "0"
      template:
        spec:
          containers:
          - name: ray-head
            image: ${IMAGE_URI}
            env:
            - name: HF_TOKEN
              valueFrom:
                secretKeyRef:
                  name: hf-secret
                  key: hf_api_token
            ports:
            - containerPort: 6379
              name: gcs
            - containerPort: 8265
              name: dashboard
            - containerPort: 10001
              name: client
            - containerPort: 8000
              name: serve
            resources:
              limits:
                cpu: "4"
                memory: "16Gi"
              requests:
                cpu: "4"
                memory: "16Gi"
    workerGroupSpecs:
    - groupName: v6e-16
      replicas: 1
      numOfHosts: 4
      template:
        spec:
          containers:
          - name: ray-worker
            image: ${IMAGE_URI}
            env:
            - name: HF_TOKEN
              valueFrom:
                secretKeyRef:
                  name: hf-secret
                  key: hf_api_token
            resources:
              limits:
                cpu: "64"
                memory: "200Gi"
                google.com/tpu: "4"
              requests:
                cpu: "64"
                memory: "200Gi"
                google.com/tpu: "4"
            volumeMounts:
            - mountPath: /dev/shm
              name: dshm
          volumes:
          - name: dshm
            emptyDir:
              medium: Memory
              sizeLimit: 100Gi
          nodeSelector:
            cloud.google.com/gke-tpu-accelerator: tpu-v6e-slice
            cloud.google.com/gke-tpu-topology: 4x4

Ironwood (TPU7x)

apiVersion: ray.io/v1
kind: RayService
metadata:
  name: ray-serve-torchtpu
spec:
  serveConfigV2: |
    applications:
    - name: default
      route_prefix: /
      import_path: ray.serve.llm:build_openai_app
      args:
        llm_configs:
        - model_loading_config:
            model_id: google/gemma-4-26B-A4B-it
          accelerator_type: TPU-V7X
          accelerator_config:
            kind: tpu
            topology: 2x2x2
          placement_group_config:
            bundle_per_worker:
              TPU: 1
          engine_kwargs:
            tensor_parallel_size: 16
            enable_expert_parallel: true
            enable_prefix_caching: true
            language_model_only: true
            disable_hybrid_kv_cache_manager: true
            max_model_len: 8192
            max_num_batched_tokens: 4096
          experimental_configs:
            stream_batching_interval_ms: 10
          runtime_env:
            env_vars:
              TPU_MULTIHOST_BACKEND: "ray"
              # Execute Gemma 4 routed expert layers on TPU TensorCores instead of SparseCore.
              USE_MOE_SPARSE_CORE: "0"
              VLLM_RAY_EXTRA_ENV_VARS_TO_COPY: "USE_MOE_SPARSE_CORE"
  rayClusterConfig:
    headGroupSpec:
      rayStartParams:
        num-cpus: "0"
      template:
        spec:
          containers:
          - name: ray-head
            image: ${IMAGE_URI}
            env:
            - name: HF_TOKEN
              valueFrom:
                secretKeyRef:
                  name: hf-secret
                  key: hf_api_token
            - name: RAY_TPU_RESOURCE_PER_CHIP
              value: "2"
            ports:
            - containerPort: 6379
              name: gcs
            - containerPort: 8265
              name: dashboard
            - containerPort: 10001
              name: client
            - containerPort: 8000
              name: serve
            resources:
              limits:
                cpu: "4"
                memory: "16Gi"
              requests:
                cpu: "4"
                memory: "16Gi"
    workerGroupSpecs:
    - groupName: tpu7x-16
      replicas: 1
      numOfHosts: 2
      resources:
        # Advertises 8 logical TPU cores (2 chiplets per physical chip) per host.
        TPU: "8"
      template:
        spec:
          containers:
          - name: ray-worker
            image: ${IMAGE_URI}
            env:
            - name: HF_TOKEN
              valueFrom:
                secretKeyRef:
                  name: hf-secret
                  key: hf_api_token
            - name: RAY_TPU_RESOURCE_PER_CHIP
              value: "2"
            resources:
              limits:
                cpu: "64"
                memory: "250Gi"
                google.com/tpu: "4"
              requests:
                cpu: "64"
                memory: "250Gi"
                google.com/tpu: "4"
            volumeMounts:
            - mountPath: /dev/shm
              name: dshm
          volumes:
          - name: dshm
            emptyDir:
              medium: Memory
              sizeLimit: 128Gi
          nodeSelector:
            cloud.google.com/gke-tpu-accelerator: tpu7x
            cloud.google.com/gke-tpu-topology: 2x2x2

Aplica el manifiesto de RayService:

Trillium (v6e)

envsubst < ray-service.tpu-v6e-multihost.yaml | kubectl apply -f -

Ironwood (TPU7x)

envsubst < ray-service.tpu-v7x-multihost.yaml | kubectl apply -f -

Espera hasta que el recurso personalizado de RayService termine de cargar los pesos del modelo y compilar los gráficos de XLA en la porción de TPU y alcance el estado Running:

kubectl get rayservice ray-serve-torchtpu
kubectl get svc ray-serve-torchtpu-serve-svc

El resultado es similar a lo siguiente:

NAME                 SERVICE STATUS   NUM SERVE ENDPOINTS
ray-serve-torchtpu   Running          3
NAME                           TYPE        CLUSTER-IP       EXTERNAL-IP   PORT(S)    AGE
ray-serve-torchtpu-serve-svc   ClusterIP   34.118.239.213   <none>        8000/TCP   29s

Reenvía el puerto local 8000 al Service ray-serve-torchtpu-serve-svc y consulta el modelo con los comandos curl en Consulta el modelo entregado:
```
kubectl port-forward svc/ray-serve-torchtpu-serve-svc 8000:8000 >/dev/null 2>&1 &
```

Observa el rendimiento del modelo

Para supervisar el uso de TensorCore de TPU, el uso de HBM y las métricas del clúster de Ray en Cloud Monitoring, consulta Recopila y visualiza registros y métricas para clústeres de Ray en GKE y Visualiza métricas de TPU en el panel de Ray.

Accede al panel de Ray

Para inspeccionar el estado de tus implementaciones de Ray Serve, los actores de réplica, las etiquetas de nodo y las latencias de las consultas en el panel de Ray integrado, sigue estos pasos:

Si la redirección de puertos al servicio principal de Ray aún no se está ejecutando, iníciala en segundo plano:
```
kubectl port-forward svc/ray-serve-torchtpu-head-svc 8265:8265 8000:8000 >/dev/null 2>&1 &
```
En Cloud Shell, haz clic en Vista previa en la Web y selecciona Vista previa en el puerto 8265, o bien abre http://localhost:8265 en tu navegador local.
Haz clic en las pestañas Entrega, Clúster y Métricas para ver tus implementaciones de vLLM, el estado de las réplicas, las etiquetas de nodo de ray.io/tpu-* detectadas automáticamente, las métricas de latencia de las solicitudes y el gráfico Uso de TPU por nodo.

Realiza una limpieza

Para evitar que se apliquen cargos a tu cuenta de Google Cloud por los recursos que usaste en este instructivo, borra el recurso personalizado RayService o RayCluster, el clúster de GKE y el repositorio de Artifact Registry:

Detén el proceso de redirección de puertos en segundo plano y borra los recursos personalizados de Ray y el Secret de Kubernetes:

pkill -f "kubectl port-forward" || true
kubectl delete rayservice ray-serve-torchtpu --ignore-not-found
kubectl delete raycluster ray-serve-torchtpu --ignore-not-found
kubectl delete secret hf-secret --ignore-not-found

Borra el clúster de GKE:

Autopilot

gcloud container clusters delete ${CLUSTER_NAME} \
    --region=${REGION}

Estándar

gcloud container clusters delete ${CLUSTER_NAME} \
    --zone=${ZONE}

Borra el repositorio de Artifact Registry:

gcloud artifacts repositories delete ${ARTIFACT_REGISTRY} \
    --location=${REGION}

¿Qué sigue?

Obtén más información sobre PyTorch en Cloud TPU.
Explora la documentación de TorchTPU.
Aprende a ajustar modelos en TPU de varios hosts con Ray y TorchTPU.
Aprende a entregar modelos abiertos de Gemma con TPU de varios hosts en GKE con Ray.
Aprende a administrar implementaciones de Ray Serve de forma declarativa con un recurso personalizado de RayService de KubeRay.
Lee Acerca de las TPU en GKE.
Explora la documentación de Ray Serve LLM.

Entrega LLM en TPU de varios hosts con Ray y TorchTPU Organiza tus páginas con colecciones Guarda y categoriza el contenido según tus preferencias.

Fondo

TPU

vLLM y TorchTPU

LLM de Ray Serve en GKE

Objetivos

Costos

Antes de comenzar

Prepara el entorno

Compila la imagen de contenedor personalizada

Crea un clúster de GKE

Autopilot

Estándar

Trillium (v6e)

Ironwood (TPU7x)

Implementa el recurso personalizado RayCluster de varios hosts

Trillium (v6e)

Ironwood (TPU7x)

Trillium (v6e)

Ironwood (TPU7x)

Trillium (v6e)

Ironwood (TPU7x)

Trillium (v6e)

Ironwood (TPU7x)

Crea e implementa la aplicación de LLM de Ray Serve

Trillium (v6e)

Ironwood (TPU7x)

Trillium (v6e)

Ironwood (TPU7x)

Consulta el modelo entregado

Trillium (v6e)

Ironwood (TPU7x)

Implementa de forma declarativa con un recurso personalizado de RayService

Trillium (v6e)

Ironwood (TPU7x)

Trillium (v6e)

Ironwood (TPU7x)

Observa el rendimiento del modelo

Accede al panel de Ray

Realiza una limpieza

Autopilot

Estándar

¿Qué sigue?

Entrega LLM en TPU de varios hosts con Ray y TorchTPU