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

Ajusta modelos en TPU de varios hosts con Ray y TorchTPU

Autopilot Standard

En este instructivo, se muestra cómo ejecutar cargas de trabajo de PyTorch distribuidas y ajustar un modelo de Gemma 4 en porciones de Cloud TPU de varios hosts en Google Kubernetes Engine (GKE) con Ray Core, Ray Train, PyTorch Fully Sharded Data Parallel (FSDP2) y 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 quieran organizar cargas de trabajo de entrenamiento de PyTorch distribuidas en porciones de TPU de varios hosts en Kubernetes. 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 entrenar modelos que superan la capacidad de memoria de un solo chip o host de TPU.

TorchTPU

TorchTPU es un backend de PyTorch de código abierto para Cloud TPU, desarrollado en colaboración entre Google y Meta. Con TorchTPU, puedes ejecutar modelos de PyTorch existentes en hardware de TPU con cambios mínimos en el código y escalar las cargas de trabajo de entrenamiento y de inferencia desde un solo chip hasta pods de TPU de varias porciones a través del compilador de álgebra lineal acelerada (XLA). TorchTPU usa un modelo de ejecución de un proceso por dispositivo en la red de ICI de TPU y admite las APIs de PyTorch estándar, incluidas la ejecución inmediata, torch.compile, torch.distributed, DeviceMesh, DTensor y Fully Sharded Data Parallel (FSDP2).

KubeRay y Ray en TPU

KubeRay es un operador de Kubernetes para implementar y administrar aplicaciones de Ray. Cuando habilitas el complemento de Ray en GKE, se implementan el operador de KubeRay administrado y el webhook de TPU de KubeRay. Cuando KubeRay crea Pods de trabajador que solicitan recursos google.com/tpu, el webhook configura automáticamente los nombres de host de los Pods, el Service sin interfaz gráfica y las variables de entorno de TPU necesarias para la coordinación de varios hosts.

Ray detecta la topología de TPU en cada nodo durante el inicio y programa las cargas de trabajo distribuidas en las porciones de TPU. Puedes organizar las cargas de trabajo distribuidas de TorchTPU en dos niveles:

Ray Core proporciona grupos de ubicación que tienen en cuenta la topología y la programación de actores para bucles de ejecución distribuidos personalizados, como el procesamiento paralelo de tensores o las canalizaciones de aprendizaje por refuerzo (RL).
Ray Train proporciona entrenamiento administrado de datos paralelos y modelos paralelos distribuidos, incluidos PyTorch Distributed Data Parallel (DDP), PyTorch FSDP2 y JAX de programa único y datos múltiples (SPMD), con puntos de control integrados y recuperación automática de fallas.

Objetivos

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

Crea un clúster de GKE en modo Autopilot o Standard en TPU Trillium (v6e) o Ironwood (TPU7x).
Compila una imagen de contenedor que coloque las bibliotecas de entrenamiento de Ray y Hugging Face sobre la imagen base de TorchTPU.
Implementar un recurso personalizado de RayCluster para aprovisionar los nodos principales y de trabajador TPU
Ejecuta un procesamiento paralelo de tensores distribuido con PyTorch DeviceMesh y DTensor en una porción de TPU con Ray Core.
Ajustar un modelo de Gemma 4 con PyTorch FSDP2 en una porción de TPU de varios hosts con Ray Train

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, la API de Cloud Storage 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 que tu cuenta tenga los siguientes roles de IAM:
- Administrador de Kubernetes Engine (roles/container.admin)
- Administrador de Artifact Registry (roles/artifactregistry.admin)
- Administrador de almacenamiento (roles/storage.admin)

Prepara el entorno

En este instructivo, usarás Cloud Shell para administrar recursos alojados en Google Cloud. Cloud Shell viene preinstalado con el software que necesitas para este instructivo, incluidos 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/raytrain/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"

Para descargar los pesos del modelo de Gemma 4, confirma el contrato de licencia de Google en Hugging Face:
1. Ve a la página del modelo Gemma 4 en Hugging Face.
2. Accede y acepta las condiciones de la licencia haciendo clic en Aceptar y acceder al repositorio.
3. En la configuración de tu cuenta de Hugging Face, genera un token de acceso de usuario con permisos de Read.
Establece las variables de entorno para tu proyecto de Google Cloud , el clúster, el bucket de Cloud Storage y el token de Hugging Face:
```
export PROJECT_ID=PROJECT_ID
export CLUSTER_NAME=CLUSTER_NAME
export REGION=REGION
export ZONE=ZONE
export ARTIFACT_REGISTRY=ARTIFACT_REGISTRY
export BUCKET_NAME=BUCKET_NAME
export HF_TOKEN=HF_TOKEN
export IMAGE_URI=${REGION}-docker.pkg.dev/${PROJECT_ID}/${ARTIFACT_REGISTRY}/ray-torchtpu: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-torchtpu-cluster.
- REGION: Es la región en la que se encuentran el repositorio de Artifact Registry, el bucket de Cloud Storage y la capacidad de TPU.
- ZONE: Es la zona en la que está disponible tu capacidad de TPU. Para conocer las regiones y zonas disponibles por versión de TPU, consulta Regiones y zonas de TPU.
- ARTIFACT_REGISTRY: Es el nombre de tu repositorio de Docker de Artifact Registry, como ray-torchtpu.
- BUCKET_NAME: Es el nombre de tu bucket de Cloud Storage para almacenar puntos de control de Ray Train, como PROJECT_ID-ray-torchtpu.
- HF_TOKEN: Es tu token de acceso de Hugging Face.

Crea un clúster de GKE

Puedes ejecutar cargas de trabajo de entrenamiento distribuidas de TorchTPU 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 implementes el recurso personalizado RayCluster en un paso posterior, GKE Autopilot aprovisionará automáticamente el nodo principal de CPU y la porción de TPU de varios hosts que coincidan con tus solicitudes de recursos de Pods 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, GKE Dataplane V2 y la federación de identidades para cargas de trabajo para GKE habilitada:
  
  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 \ --workload-pool=${PROJECT_ID}.svc.id.goog
  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 bucket de Cloud Storage para almacenar puntos de control de Ray Train compartidos en los nodos TPU de varios hosts y otorga al recurso ServiceAccount de Kubernetes default acceso al bucket:

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

PROJECT_NUMBER=$(gcloud projects describe ${PROJECT_ID} --format="value(projectNumber)")
gcloud storage buckets add-iam-policy-binding gs://${BUCKET_NAME} \
    --member="principal://iam.googleapis.com/projects/${PROJECT_NUMBER}/locations/global/workloadIdentityPools/${PROJECT_ID}.svc.id.goog/subject/ns/default/sa/default" \
    --role="roles/storage.objectAdmin"

Crea un Secret de Kubernetes con tu token de acceso de Hugging Face para que los Pods de Ray puedan descargar los pesos de Gemma 4 restringidos:
```
kubectl create secret generic hf-secret \
    --from-literal=hf_api_token=${HF_TOKEN}
```

Compila la imagen del contenedor

Compila una imagen de contenedor que instale Ray y las bibliotecas de entrenamiento de Hugging Face sobre la imagen base de TorchTPU compilada previamente, que incluye TorchTPU y la biblioteca libtpu.

Crea un repositorio de Docker de Artifact Registry y configura la autenticación de Docker:

gcloud artifacts repositories create ${ARTIFACT_REGISTRY} \
    --repository-format=docker \
    --location=${REGION}

gcloud auth configure-docker ${REGION}-docker.pkg.dev,us-docker.pkg.dev

Inspecciona el archivo Dockerfile en el directorio de muestra:

FROM us-docker.pkg.dev/ml-oss-artifacts-published/torch-tpu-public-release-container/torch-tpu:latest

# Install wget for KubeRay health probes, Ray (Core, Dashboard, and Train), and Hugging Face libraries for fine-tuning
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,train]==2.58.0" \
    "transformers>=5.5.0" \
    datasets

El entrenamiento de porción única con la clase TorchTrainer de Ray y la programación por chiplet en Ironwood (TPU7x) requieren la versión 2.58.0 o posterior de Ray, y el ajuste de Gemma 4 requiere la versión 5.5.0 o posterior de la biblioteca transformers.

Compila y envía la imagen a Artifact Registry:

docker build -t ${IMAGE_URI} .
docker push ${IMAGE_URI}

Implementa el recurso personalizado de RayCluster

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-torchtpu-cluster
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
          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-torchtpu-cluster
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
          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 de RayCluster y todos los Pods de trabajador TPU alcancen el estado Running:

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

El resultado es similar a lo siguiente:

Trillium (v6e)

NAME                                         READY   STATUS    RESTARTS   AGE
ray-torchtpu-cluster-head-pdpk6              1/1     Running   0          2m30s
ray-torchtpu-cluster-v6e-16-worker-brwpc     1/1     Running   0          2m30s
ray-torchtpu-cluster-v6e-16-worker-dx4w7     1/1     Running   0          2m30s
ray-torchtpu-cluster-v6e-16-worker-h4d4w     1/1     Running   0          2m30s
ray-torchtpu-cluster-v6e-16-worker-t2g24     1/1     Running   0          2m30s

Ironwood (TPU7x)

NAME                                         READY   STATUS    RESTARTS   AGE
ray-torchtpu-cluster-head-k9x2m              1/1     Running   0          2m30s
ray-torchtpu-cluster-tpu7x-16-worker-7cwpq   1/1     Running   0          2m30s
ray-torchtpu-cluster-tpu7x-16-worker-m98r9   1/1     Running   0          2m30s

Inspecciona las variables de entorno en uno de los Pods de trabajador para confirmar que el webhook de TPU de KubeRay insertó las variables de coordinación de TPU de varios hosts y de topología de TorchTPU:

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-torchtpu-cluster-headless,v6e-16-0-1.ray-torchtpu-cluster-headless,v6e-16-0-2.ray-torchtpu-cluster-headless,v6e-16-0-3.ray-torchtpu-cluster-headless
TORCH_TPU_TOPOLOGY=4,4,1
TORCH_TPU_SLICEBUILDER_ADDRESSES=v6e-16-0-0.ray-torchtpu-cluster-headless:8471,v6e-16-0-0.ray-torchtpu-cluster-headless:8472,v6e-16-0-0.ray-torchtpu-cluster-headless:8473,v6e-16-0-0.ray-torchtpu-cluster-headless:8474,v6e-16-0-1.ray-torchtpu-cluster-headless:8471,v6e-16-0-1.ray-torchtpu-cluster-headless:8472,v6e-16-0-1.ray-torchtpu-cluster-headless:8473,v6e-16-0-1.ray-torchtpu-cluster-headless:8474,v6e-16-0-2.ray-torchtpu-cluster-headless:8471,v6e-16-0-2.ray-torchtpu-cluster-headless:8472,v6e-16-0-2.ray-torchtpu-cluster-headless:8473,v6e-16-0-2.ray-torchtpu-cluster-headless:8474,v6e-16-0-3.ray-torchtpu-cluster-headless:8471,v6e-16-0-3.ray-torchtpu-cluster-headless:8472,v6e-16-0-3.ray-torchtpu-cluster-headless:8473,v6e-16-0-3.ray-torchtpu-cluster-headless:8474
TPU_WORKER_ID=3
TPU_NAME=v6e-16-0

Ironwood (TPU7x)

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

Establece una sesión de redirección de puertos al servicio principal de Ray:
```
kubectl port-forward svc/ray-torchtpu-cluster-head-svc 8265:8265 >/dev/null 2>&1 &
```
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.

Ejecuta el paralelismo de tensores distribuido con Ray Core

Ray Core proporciona la función slice_placement_group en el módulo ray.util.tpu para reservar una porción de TPU de varios hosts intacta y la clase PlacementGroupSchedulingStrategy para programar un actor de Ray por dispositivo de TPU lógico en la porción.

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

import argparse
import os

import ray
from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy
from ray.util.tpu import (
    get_tpu_version_from_type,
    slice_placement_group,
)
import torch
import torch.distributed as dist
from torch.distributed.device_mesh import init_device_mesh
from torch.distributed.tensor.parallel import (
    ColwiseParallel,
    RowwiseParallel,
    parallelize_module,
)
import torch.nn as nn
import torch_tpu  # noqa: F401

RANDOM_SEED = 42


class FeedForward(nn.Module):
    """Two-layer feed-forward block for tensor-parallel sharding across a TPU slice."""

    def __init__(self, hidden_dim: int):
        super().__init__()
        self.w1 = nn.Linear(hidden_dim, 4 * hidden_dim, bias=False)
        self.act = nn.SiLU()
        self.w2 = nn.Linear(4 * hidden_dim, hidden_dim, bias=False)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return x + self.w2(self.act(self.w1(x)))


@ray.remote
class TorchTPUSliceWorker:
    """Ray actor bound to a single logical TPU device within a slice placement group."""

    def __init__(self, devices_per_host: int):
        self.node_ip = ray.util.get_node_ip_address()
        hosts = os.environ["TPU_WORKER_HOSTNAMES"].split(",")
        self.worker_id = int(os.environ.get("TPU_WORKER_ID", "0"))
        self.local_rank = int(
            ray.get_runtime_context().get_accelerator_ids()["TPU"][0]
        )
        self.rank = self.worker_id * devices_per_host + self.local_rank
        # Set per-process PyTorch env:// variables from the webhook's host-level
        # TPU metadata and Ray's assigned local device ID (TorchTrainer sets these
        # automatically in the Ray Train section below).
        os.environ.update(
            {
                "MASTER_ADDR": hosts[0],
                "MASTER_PORT": "29500",
                "WORLD_SIZE": str(len(hosts) * devices_per_host),
                "RANK": str(self.rank),
                "LOCAL_RANK": str(self.local_rank),
                "LOCAL_WORLD_SIZE": str(devices_per_host),
            }
        )

    def run_tensor_parallel_step(
        self,
        batch_size: int = 64,
        hidden_dim: int = 128,
    ) -> dict:
        dist.init_process_group(backend="tpu_dist")
        world_size = dist.get_world_size()
        device = torch.device("tpu")

        # 1. Compute single-device reference output on CPU with identical seed.
        torch.manual_seed(RANDOM_SEED)
        ref_model = FeedForward(hidden_dim).eval()
        inputs_cpu = torch.randn(batch_size, hidden_dim)
        with torch.no_grad():
            ref_output = ref_model(inputs_cpu)

        # 2. Shard the model across all TPU devices in the slice with DTensor TP.
        torch.manual_seed(RANDOM_SEED)
        tp_model = FeedForward(hidden_dim).to(device).eval()
        mesh = init_device_mesh("tpu", (world_size,))
        tp_model = parallelize_module(
            tp_model,
            mesh,
            {"w1": ColwiseParallel(), "w2": RowwiseParallel()},
        )

        with torch.no_grad():
            tp_output = tp_model(inputs_cpu.to(device)).cpu()

        torch.testing.assert_close(tp_output, ref_output, rtol=1e-3, atol=1e-2)
        max_abs_diff = float((tp_output - ref_output).abs().max().item())

        # 3. Verify direct ICI all_reduce across all ranks in the slice.
        rank_tensor = torch.full(
            (4,), float(self.rank + 1), dtype=torch.float32, device=device
        )
        dist.all_reduce(rank_tensor, op=dist.ReduceOp.SUM)
        all_reduce_sum = rank_tensor.cpu().tolist()

        dist.destroy_process_group()
        return {
            "rank": self.rank,
            "local_rank": self.local_rank,
            "worker_id": self.worker_id,
            "node_ip": self.node_ip,
            "all_reduce_sum": all_reduce_sum,
            "max_abs_diff": max_abs_diff,
        }


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Run distributed TorchTPU tensor parallelism across a Ray TPU slice."
    )
    parser.add_argument(
        "--topology",
        type=str,
        default="4x4",
        help="TPU physical topology.",
    )
    parser.add_argument(
        "--accelerator-type",
        type=str,
        default="TPU-V6E",
        help="Ray TPU accelerator type (TPU-V6E, TPU-V7X, etc.).",
    )
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    ray.init()

    slice_handle = slice_placement_group(
        topology=args.topology,
        accelerator_version=get_tpu_version_from_type(args.accelerator_type),
        num_slices=1,
        resources_per_bundle={"TPU": 1},
    )
    ray.get(slice_handle.placement_group.ready())

    devices_per_host = slice_handle.num_bundles // slice_handle.num_hosts
    workers = [
        TorchTPUSliceWorker.options(
            resources={"TPU": 1},
            scheduling_strategy=PlacementGroupSchedulingStrategy(
                placement_group=slice_handle.placement_group,
                placement_group_bundle_index=bundle_idx,
            ),
        ).remote(devices_per_host)
        for bundle_idx in range(slice_handle.num_bundles)
    ]
    results = sorted(
        ray.get([w.run_tensor_parallel_step.remote() for w in workers]),
        key=lambda r: r["rank"],
    )

    expected_sum = sum(range(1, slice_handle.num_bundles + 1))
    for r in results:
        print(
            f"Rank {r['rank']:2d} (worker_id={r['worker_id']}, "
            f"local_rank={r['local_rank']}, ip={r['node_ip']}): "
            f"all_reduce_sum={r['all_reduce_sum']}"
        )
        assert r["all_reduce_sum"] == [float(expected_sum)] * 4

    print(
        f"All {len(results)} TorchTPU ranks verified DTensor TP parity "
        f"(max_abs_diff={results[0]['max_abs_diff']:.6f}) and "
        f"all_reduce sum = {expected_sum:.1f}"
    )
    for w in workers:
        ray.kill(w)
    slice_handle.shutdown()


if __name__ == "__main__":
    main()

La secuencia de comandos ejecuta el paralelismo tensorial distribuido en la porción de TPU en tres pasos:

La función slice_placement_group reserva una porción de TPU de 16 dispositivos y programa un actor TorchTPUSliceWorker en cada dispositivo TPU.
Cada actor calcula su rango global a partir de su índice de host y el ID de dispositivo local, y, luego, inicializa el grupo de procesos tpu_dist.
Las funciones init_device_mesh y parallelize_module fragmentan las capas de retroalimentación en los 16 dispositivos TPU, verifican el resultado con una referencia de CPU de un solo dispositivo y ejecutan un colectivo all_reduce sobre ICI.

Envía el trabajo de Ray Core para tu versión de TPU:

Trillium (v6e)

ray job submit \
    --address http://localhost:8265 \
    --working-dir . \
    --runtime-env-json '{"excludes": ["ray-env", ".git", "*.whl"]}' \
    -- python3 ray_core_collectives.py --topology=4x4 --accelerator-type=TPU-V6E

Ironwood (TPU7x)

ray job submit \
    --address http://localhost:8265 \
    --working-dir . \
    --runtime-env-json '{"excludes": ["ray-env", ".git", "*.whl"]}' \
    -- python3 ray_core_collectives.py --topology=2x2x2 --accelerator-type=TPU-V7X

El resultado es similar a lo siguiente:

Trillium (v6e)

Rank  0 (worker_id=0, local_rank=0, ip=10.72.3.20): all_reduce_sum=[136.0, 136.0, 136.0, 136.0]
Rank  1 (worker_id=0, local_rank=1, ip=10.72.3.20): all_reduce_sum=[136.0, 136.0, 136.0, 136.0]
Rank  2 (worker_id=0, local_rank=2, ip=10.72.3.20): all_reduce_sum=[136.0, 136.0, 136.0, 136.0]
Rank  3 (worker_id=0, local_rank=3, ip=10.72.3.20): all_reduce_sum=[136.0, 136.0, 136.0, 136.0]
Rank  4 (worker_id=1, local_rank=0, ip=10.72.14.20): all_reduce_sum=[136.0, 136.0, 136.0, 136.0]
Rank  5 (worker_id=1, local_rank=1, ip=10.72.14.20): all_reduce_sum=[136.0, 136.0, 136.0, 136.0]
Rank  6 (worker_id=1, local_rank=2, ip=10.72.14.20): all_reduce_sum=[136.0, 136.0, 136.0, 136.0]
Rank  7 (worker_id=1, local_rank=3, ip=10.72.14.20): all_reduce_sum=[136.0, 136.0, 136.0, 136.0]
Rank  8 (worker_id=2, local_rank=0, ip=10.72.7.20): all_reduce_sum=[136.0, 136.0, 136.0, 136.0]
Rank  9 (worker_id=2, local_rank=1, ip=10.72.7.20): all_reduce_sum=[136.0, 136.0, 136.0, 136.0]
Rank 10 (worker_id=2, local_rank=2, ip=10.72.7.20): all_reduce_sum=[136.0, 136.0, 136.0, 136.0]
Rank 11 (worker_id=2, local_rank=3, ip=10.72.7.20): all_reduce_sum=[136.0, 136.0, 136.0, 136.0]
Rank 12 (worker_id=3, local_rank=0, ip=10.72.6.20): all_reduce_sum=[136.0, 136.0, 136.0, 136.0]
Rank 13 (worker_id=3, local_rank=1, ip=10.72.6.20): all_reduce_sum=[136.0, 136.0, 136.0, 136.0]
Rank 14 (worker_id=3, local_rank=2, ip=10.72.6.20): all_reduce_sum=[136.0, 136.0, 136.0, 136.0]
Rank 15 (worker_id=3, local_rank=3, ip=10.72.6.20): all_reduce_sum=[136.0, 136.0, 136.0, 136.0]
All 16 TorchTPU ranks verified DTensor TP parity (max_abs_diff=0.002470) and all_reduce sum = 136.0
------------------------------------------
Job 'raysubmit_HFTkKrjSnfza3ZFD' succeeded
------------------------------------------

Ironwood (TPU7x)

Rank  0 (worker_id=0, local_rank=0, ip=172.16.14.6): all_reduce_sum=[136.0, 136.0, 136.0, 136.0]
Rank  1 (worker_id=0, local_rank=1, ip=172.16.14.6): all_reduce_sum=[136.0, 136.0, 136.0, 136.0]
Rank  2 (worker_id=0, local_rank=2, ip=172.16.14.6): all_reduce_sum=[136.0, 136.0, 136.0, 136.0]
Rank  3 (worker_id=0, local_rank=3, ip=172.16.14.6): all_reduce_sum=[136.0, 136.0, 136.0, 136.0]
Rank  4 (worker_id=0, local_rank=4, ip=172.16.14.6): all_reduce_sum=[136.0, 136.0, 136.0, 136.0]
Rank  5 (worker_id=0, local_rank=5, ip=172.16.14.6): all_reduce_sum=[136.0, 136.0, 136.0, 136.0]
Rank  6 (worker_id=0, local_rank=6, ip=172.16.14.6): all_reduce_sum=[136.0, 136.0, 136.0, 136.0]
Rank  7 (worker_id=0, local_rank=7, ip=172.16.14.6): all_reduce_sum=[136.0, 136.0, 136.0, 136.0]
Rank  8 (worker_id=1, local_rank=0, ip=172.16.15.7): all_reduce_sum=[136.0, 136.0, 136.0, 136.0]
Rank  9 (worker_id=1, local_rank=1, ip=172.16.15.7): all_reduce_sum=[136.0, 136.0, 136.0, 136.0]
Rank 10 (worker_id=1, local_rank=2, ip=172.16.15.7): all_reduce_sum=[136.0, 136.0, 136.0, 136.0]
Rank 11 (worker_id=1, local_rank=3, ip=172.16.15.7): all_reduce_sum=[136.0, 136.0, 136.0, 136.0]
Rank 12 (worker_id=1, local_rank=4, ip=172.16.15.7): all_reduce_sum=[136.0, 136.0, 136.0, 136.0]
Rank 13 (worker_id=1, local_rank=5, ip=172.16.15.7): all_reduce_sum=[136.0, 136.0, 136.0, 136.0]
Rank 14 (worker_id=1, local_rank=6, ip=172.16.15.7): all_reduce_sum=[136.0, 136.0, 136.0, 136.0]
Rank 15 (worker_id=1, local_rank=7, ip=172.16.15.7): all_reduce_sum=[136.0, 136.0, 136.0, 136.0]
All 16 TorchTPU ranks verified DTensor TP parity (max_abs_diff=0.002470) and all_reduce sum = 136.0
------------------------------------------
Job 'raysubmit_MrJbAe3hvARBcBN2' succeeded
------------------------------------------

Ajusta un modelo de Gemma 4 con Ray Train y PyTorch FSDP2

La clase TorchTrainer de Ray Train automatiza la reserva de porciones de TPU, la ubicación de los trabajadores y la inicialización de torch.distributed en una porción de TPU de varios hosts.

Revisa la secuencia de comandos ray_train_llm_finetune.py en el repositorio de muestra, que ajusta el modelo google/gemma-4-E2B-it en el conjunto de datos de instrucciones databricks/databricks-dolly-15k:

import argparse
import os
import tempfile
import time

from datasets import load_dataset
import ray
import ray.train
from ray.train import (
    Checkpoint,
    CheckpointConfig,
    RunConfig,
    ScalingConfig,
)
from ray.train.torch import TorchConfig, TorchTrainer
import torch
from torch.distributed import ReduceOp, all_reduce
from torch.distributed.checkpoint.state_dict import (
    StateDictOptions,
    get_model_state_dict,
)
from torch.distributed.device_mesh import init_device_mesh
from torch.distributed.fsdp import MixedPrecisionPolicy, fully_shard
from torch.utils.data import DataLoader
import torch_tpu  # noqa: F401
from transformers import (
    AutoModelForCausalLM,
    AutoTokenizer,
    DataCollatorForLanguageModeling,
)


def create_dataloader(
    tokenizer: AutoTokenizer,
    batch_size: int,
    dataset_split: str = "train[:1024]",
    max_seq_len: int = 512,
) -> DataLoader:
    """Loads and tokenizes the Dolly 15k dataset for causal LM fine-tuning."""
    dataset = load_dataset(
        "databricks/databricks-dolly-15k", split=dataset_split
    )
    tokenized = dataset.map(
        lambda x: tokenizer(
            f"Instruction: {x['instruction']}\nResponse: {x['response']}",
            truncation=True,
            max_length=max_seq_len,
            padding="max_length",
        ),
        remove_columns=dataset.column_names,
    )
    dataloader = DataLoader(
        tokenized,
        batch_size=batch_size,
        shuffle=True,
        drop_last=True,
        collate_fn=DataCollatorForLanguageModeling(
            tokenizer=tokenizer, mlm=False
        ),
    )
    return ray.train.torch.prepare_data_loader(dataloader)


def train_func(config: dict) -> None:
    """Per-worker distributed FSDP2 training loop executed on each logical TPU core."""
    ctx = ray.train.get_context()
    world_rank = ctx.get_world_rank()
    world_size = ctx.get_world_size()
    device = ray.train.torch.get_device()

    model_id = config["model_id"]
    dataset_split = config["dataset_split"]
    batch_size = config["batch_size_per_device"]
    num_epochs = config["num_epochs"]
    max_seq_len = 512

    tokenizer = AutoTokenizer.from_pretrained(model_id)
    dataloader = create_dataloader(
        tokenizer=tokenizer,
        batch_size=batch_size,
        dataset_split=dataset_split,
        max_seq_len=max_seq_len,
    )

    torch.manual_seed(42)
    model = AutoModelForCausalLM.from_pretrained(
        model_id,
        dtype=torch.bfloat16,
        attn_implementation="eager",
        vision_config=None,
        audio_config=None,
    ).to(device)
    model.config.get_text_config().use_cache = False

    mesh = init_device_mesh("tpu", (world_size,))
    mp_policy = MixedPrecisionPolicy(
        param_dtype=torch.bfloat16,
        reduce_dtype=torch.float32,
    )
    text_model = model.model.language_model
    text_model.embed_tokens_per_layer.requires_grad_(False)
    for layer in text_model.layers:
        fully_shard(layer.mlp, mesh=mesh, mp_policy=mp_policy)
    fully_shard(
        model,
        mesh=mesh,
        mp_policy=mp_policy,
        ignored_params=set(text_model.embed_tokens_per_layer.parameters()),
    )

    optimizer = torch.optim.AdamW(
        model.parameters(), lr=2e-5, weight_decay=0.01
    )
    global_batch_size = batch_size * world_size

    for epoch in range(num_epochs):
        dataloader.sampler.set_epoch(epoch)
        model.train()
        epoch_start = time.perf_counter()
        running_loss = torch.zeros(1, device=device, dtype=torch.float32)

        for batch in dataloader:
            optimizer.zero_grad()
            outputs = model(
                input_ids=batch["input_ids"],
                labels=batch["labels"],
            )
            outputs.loss.backward()
            optimizer.step()
            running_loss += outputs.loss.detach().float()

        all_reduce(running_loss, op=ReduceOp.AVG)
        avg_loss = running_loss.item() / len(dataloader)
        epoch_time = time.perf_counter() - epoch_start
        tokens_per_sec = (
            len(dataloader) * global_batch_size * max_seq_len
        ) / epoch_time

        metrics = {
            "epoch": epoch + 1,
            "loss": round(avg_loss, 4),
            "epoch_time_s": round(epoch_time, 2),
            "tokens_per_s": round(tokens_per_sec, 1),
        }
        if world_rank == 0:
            print(
                f"Epoch {epoch + 1}/{num_epochs} | loss={metrics['loss']:.4f} | "
                f"epoch_time={metrics['epoch_time_s']:.2f}s | "
                f"tokens/s={metrics['tokens_per_s']:.1f}"
            )

        full_state_dict = get_model_state_dict(
            model,
            options=StateDictOptions(
                full_state_dict=True,
                cpu_offload=True,
                ignore_frozen_params=True,
            ),
        )
        if world_rank == 0:
            with tempfile.TemporaryDirectory() as tmpdir:
                torch.save(full_state_dict, os.path.join(tmpdir, "model.pt"))
                ray.train.report(
                    metrics, checkpoint=Checkpoint.from_directory(tmpdir)
                )
        else:
            ray.train.report(metrics)


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Fine-tune a Hugging Face CausalLM on TPUs with Ray Train and TorchTPU FSDP2."
    )
    parser.add_argument(
        "--model-id",
        type=str,
        default="google/gemma-4-E2B-it",
        help="Hugging Face model ID to fine-tune.",
    )
    parser.add_argument(
        "--topology",
        type=str,
        default="4x4",
        help="TPU slice topology.",
    )
    parser.add_argument(
        "--accelerator-type",
        type=str,
        default="TPU-V6E",
        help="Ray TPU accelerator type (TPU-V6E, TPU-V7X, etc.).",
    )
    parser.add_argument(
        "--num-workers",
        type=int,
        default=16,
        help="Total TPU worker processes across the slice.",
    )
    parser.add_argument(
        "--dataset-split",
        type=str,
        default="train[:1024]",
        help="Hugging Face dataset split to load.",
    )
    parser.add_argument(
        "--num-epochs",
        type=int,
        default=3,
        help="Number of fine-tuning epochs.",
    )
    parser.add_argument(
        "--batch-size-per-device",
        type=int,
        default=4,
        help="Per-device training batch size.",
    )
    parser.add_argument(
        "--storage-path",
        type=str,
        default=os.environ.get("STORAGE_PATH"),
        required="STORAGE_PATH" not in os.environ,
        help="Cloud Storage or shared file system path for checkpoints.",
    )
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    ray.init()
    run_name = f"torchtpu_gemma4_fsdp2_sft_{int(time.time())}"

    print(
        f"Starting Ray Train FSDP2 SFT job: model={args.model_id}, "
        f"topology={args.topology}, accelerator_type={args.accelerator_type}, "
        f"num_workers={args.num_workers}"
    )

    trainer = TorchTrainer(
        train_loop_per_worker=train_func,
        train_loop_config={
            "model_id": args.model_id,
            "dataset_split": args.dataset_split,
            "num_epochs": args.num_epochs,
            "batch_size_per_device": args.batch_size_per_device,
        },
        torch_config=TorchConfig(backend="tpu_dist"),
        scaling_config=ScalingConfig(
            use_tpu=True,
            num_workers=args.num_workers,
            topology=args.topology,
            accelerator_type=args.accelerator_type,
            resources_per_worker={"TPU": 1},
        ),
        run_config=RunConfig(
            name=run_name,
            storage_path=args.storage_path,
            checkpoint_config=CheckpointConfig(
                num_to_keep=1,
                checkpoint_score_attribute="loss",
                checkpoint_score_order="min",
            ),
        ),
    )

    result = trainer.fit()
    print(f"Training completed! Final metrics: {result.metrics}")
    print(f"Saved checkpoint: {result.checkpoint}")


if __name__ == "__main__":
    main()

La secuencia de comandos configura la clase TorchTrainer con los siguientes componentes:

La clase ScalingConfig configura la topología de trabajadores distribuidos para el trabajo de entrenamiento. Si configuras los parámetros use_tpu=True, num_workers=16, topology y accelerator_type, se reserva una porción de TPU de 16 dispositivos y se programa un trabajador por núcleo de TPU.
La clase TorchConfig configura el grupo de procesos torch.distributed en los trabajadores. Si configuras el parámetro backend="tpu_dist", se conectarán los 16 trabajadores a través de la red de ICI de TPU antes de que se ejecute la función train_func.
La función create_dataloader descarga y asigna tokens al conjunto de datos de instrucciones, agrupa las secuencias de tokens con un objeto DataLoader de PyTorch y llama a la función ray.train.torch.prepare_data_loader para que cada trabajador reciba un fragmento separado de cada lote global en su dispositivo TPU asignado.
La función train_func carga el modelo con precisión bfloat16 y llama a la función fully_shard de FSDP2 de PyTorch en cada bloque de avance y en el modelo raíz para fragmentar los parámetros, los gradientes y los estados del optimizador en los 16 núcleos de TPU.
La clase Checkpoint y la función ray.train.report recopilan el diccionario de estado completo del modelo en el rango 0 al final de cada época y suben el punto de control a tu bucket de Cloud Storage.

Envía el trabajo de ajuste de FSDP2 de TorchTrainer para tu versión de TPU:

Trillium (v6e)

ray job submit \
    --address http://localhost:8265 \
    --working-dir . \
    --runtime-env-json '{"excludes": ["ray-env", ".git", "*.whl"]}' \
    -- python3 ray_train_llm_finetune.py \
        --topology=4x4 \
        --accelerator-type=TPU-V6E \
        --storage-path=gs://${BUCKET_NAME}/ray_train_torchtpu

Ironwood (TPU7x)

ray job submit \
    --address http://localhost:8265 \
    --working-dir . \
    --runtime-env-json '{"excludes": ["ray-env", ".git", "*.whl"]}' \
    -- python3 ray_train_llm_finetune.py \
        --topology=2x2x2 \
        --accelerator-type=TPU-V7X \
        --storage-path=gs://${BUCKET_NAME}/ray_train_torchtpu

El resultado es similar a lo siguiente:

Trillium (v6e)

Starting Ray Train FSDP2 SFT job: model=google/gemma-4-E2B-it, topology=4x4, accelerator_type=TPU-V6E, num_workers=16
(TrainController pid=5407) Attempting to start training worker group of size 16 with the following resources: [{'TPU': 1, 'accelerator_type:TPU-V6E': 0.001}] * 16
(RayTrainWorker pid=19733, ip=10.72.3.20) Setting up process group for: env:// [rank=0, world_size=16]
(RayTrainWorker pid=19733, ip=10.72.3.20) Epoch 1/3 | loss=3.1402 | epoch_time=28.17s | tokens/s=18613.5
(RayTrainWorker pid=19733, ip=10.72.3.20) Checkpoint successfully created at: Checkpoint(filesystem=gcs, path=my-project-ray-torchtpu/ray_train_torchtpu/torchtpu_gemma4_fsdp2_sft_1791864427/checkpoint_2026-10-13_04-08-10.762175)
(RayTrainWorker pid=19733, ip=10.72.3.20) Epoch 2/3 | loss=2.3318 | epoch_time=26.07s | tokens/s=20112.8
(RayTrainWorker pid=19733, ip=10.72.3.20) Checkpoint successfully created at: Checkpoint(filesystem=gcs, path=my-project-ray-torchtpu/ray_train_torchtpu/torchtpu_gemma4_fsdp2_sft_1791864427/checkpoint_2026-10-13_04-10-02.548088)
(RayTrainWorker pid=19733, ip=10.72.3.20) Epoch 3/3 | loss=2.2050 | epoch_time=25.52s | tokens/s=20545.2
(RayTrainWorker pid=19733, ip=10.72.3.20) Checkpoint successfully created at: Checkpoint(filesystem=gcs, path=my-project-ray-torchtpu/ray_train_torchtpu/torchtpu_gemma4_fsdp2_sft_1791864427/checkpoint_2026-10-13_04-11-48.701853)
Training completed! Final metrics: {'epoch': 3, 'loss': 2.205, 'epoch_time_s': 25.52, 'tokens_per_s': 20545.2}
Saved checkpoint: Checkpoint(filesystem=gcs, path=my-project-ray-torchtpu/ray_train_torchtpu/torchtpu_gemma4_fsdp2_sft_1791864427/checkpoint_2026-10-13_04-11-48.701853)
------------------------------------------
Job 'raysubmit_zbm4pnFZitA1YECA' succeeded
------------------------------------------

Ironwood (TPU7x)

Starting Ray Train FSDP2 SFT job: model=google/gemma-4-E2B-it, topology=2x2x2, accelerator_type=TPU-V7X, num_workers=16
(TrainController pid=1162) Attempting to start training worker group of size 16 with the following resources: [{'TPU': 1, 'accelerator_type:TPU-V7X': 0.001}] * 16
(RayTrainWorker pid=257, ip=172.16.26.6) Setting up process group for: env:// [rank=0, world_size=16]
(RayTrainWorker pid=257, ip=172.16.26.6) Epoch 1/3 | loss=3.1407 | epoch_time=26.28s | tokens/s=19953.3
(RayTrainWorker pid=257, ip=172.16.26.6) Checkpoint successfully created at: Checkpoint(filesystem=gcs, path=my-project-ray-torchtpu/ray_train_torchtpu/torchtpu_gemma4_fsdp2_sft_1791867004/checkpoint_2026-10-13_04-52-00.026930)
(RayTrainWorker pid=257, ip=172.16.26.6) Epoch 2/3 | loss=2.3235 | epoch_time=16.68s | tokens/s=31429.9
(RayTrainWorker pid=257, ip=172.16.26.6) Checkpoint successfully created at: Checkpoint(filesystem=gcs, path=my-project-ray-torchtpu/ray_train_torchtpu/torchtpu_gemma4_fsdp2_sft_1791867004/checkpoint_2026-10-13_04-53-33.533341)
(RayTrainWorker pid=257, ip=172.16.26.6) Epoch 3/3 | loss=2.1964 | epoch_time=16.98s | tokens/s=30884.6
(RayTrainWorker pid=257, ip=172.16.26.6) Checkpoint successfully created at: Checkpoint(filesystem=gcs, path=my-project-ray-torchtpu/ray_train_torchtpu/torchtpu_gemma4_fsdp2_sft_1791867004/checkpoint_2026-10-13_04-55-01.508219)
Training completed! Final metrics: {'epoch': 3, 'loss': 2.1964, 'epoch_time_s': 16.98, 'tokens_per_s': 30884.6}
Saved checkpoint: Checkpoint(filesystem=gcs, path=my-project-ray-torchtpu/ray_train_torchtpu/torchtpu_gemma4_fsdp2_sft_1791867004/checkpoint_2026-10-13_04-55-01.508219)
------------------------------------------
Job 'raysubmit_zSNnCXhHkwKiSmJh' succeeded
------------------------------------------

El resultado del entrenamiento muestra lo siguiente:

El ciclo de entrenamiento 1 incluye la compilación inicial del grafo de XLA. Las épocas 2 y 3 reutilizan los ejecutables almacenados en caché de /dev/shm/torch_tpu_cache, lo que permite alcanzar más de 20,000 tokens por segundo en TPU Trillium (v6e) y 31,000 tokens por segundo en Ironwood (TPU7x).
La pérdida de entropía cruzada global sincronizada en los 16 rangos de trabajadores de FSDP2 disminuye de 3.14 en la época 1 a 2.20 en la época 3 en las tres épocas de ajuste en el conjunto de datos de instrucciones de databricks/databricks-dolly-15k.
Al final de cada ciclo, el rango 0 recopila el diccionario de estado completo del modelo y sube el punto de control a tu bucket de Cloud Storage, lo que conserva el punto de control con la pérdida más baja, según se configuró en el objeto CheckpointConfig(num_to_keep=1, checkpoint_score_attribute="loss", checkpoint_score_order="min").

Observa el rendimiento de las cargas de trabajo

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 actores de Ray, grupos de posición y trabajos de entrenamiento 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-torchtpu-cluster-head-svc 8265:8265 >/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 Trabajos, Clúster y Métricas para inspeccionar los registros de los trabajadores, el estado del grupo de posiciones, las etiquetas de nodo TPU 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 RayCluster, el clúster de GKE, el bucket de Cloud Storage y el repositorio de Artifact Registry:

Detén el proceso de redirección de puertos en segundo plano y borra el recurso personalizado RayCluster y el Secret de Kubernetes:

pkill -f "kubectl port-forward" || true
kubectl delete raycluster ray-torchtpu-cluster
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 bucket de Cloud Storage y su contenido de puntos de control:
```
gcloud storage rm --recursive gs://${BUCKET_NAME}
```

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 entregar LLM en TPU de varios hosts con Ray y TorchTPU.
Obtén información para optimizar el entrenamiento de IA en TPU con DWS y Kueue.
Lee Acerca de las TPU en GKE.
Explora la documentación de Ray Train.

Ajusta modelos 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

TorchTPU

KubeRay y Ray en TPU

Objetivos

Costos

Antes de comenzar

Prepara el entorno

Crea un clúster de GKE

Autopilot

Estándar

Trillium (v6e)

Ironwood (TPU7x)

Compila la imagen del contenedor

Implementa el recurso personalizado de RayCluster

Trillium (v6e)

Ironwood (TPU7x)

Trillium (v6e)

Ironwood (TPU7x)

Trillium (v6e)

Ironwood (TPU7x)

Trillium (v6e)

Ironwood (TPU7x)

Ejecuta el paralelismo de tensores distribuido con Ray Core

Trillium (v6e)

Ironwood (TPU7x)

Trillium (v6e)

Ironwood (TPU7x)

Ajusta un modelo de Gemma 4 con Ray Train y PyTorch FSDP2

Trillium (v6e)

Ironwood (TPU7x)

Trillium (v6e)

Ironwood (TPU7x)

Observa el rendimiento de las cargas de trabajo

Accede al panel de Ray

Realiza una limpieza

Autopilot

Estándar

¿Qué sigue?

Ajusta modelos en TPU de varios hosts con Ray y TorchTPU