Google utilise la technologie IA pour traduire le contenu dans votre langue préférée. Les traductions générées par IA peuvent contenir des erreurs.

Affiner des modèles sur des TPU multi-hôtes avec Ray et TorchTPU

Autopilot Standard

Ce tutoriel vous explique comment exécuter des charges de travail PyTorch distribuées et affiner un modèle Gemma 4 sur des tranches Cloud TPU multi-hôtes dans Google Kubernetes Engine (GKE) à l'aide de Ray Core, Ray Train, PyTorch Fully Sharded Data Parallel (FSDP2) et TorchTPU.

Ce tutoriel s'adresse aux ingénieurs en machine learning (ML), aux administrateurs et opérateurs de plate-forme, ainsi qu'aux spécialistes des données et de l'IA qui souhaitent orchestrer des charges de travail d'entraînement PyTorch distribuées sur des tranches de TPU multi-hôtes dans Kubernetes. Pour en savoir plus sur les rôles courants et les exemples de tâches que nous citons dans le contenu Google Cloud , consultez Rôles utilisateur et tâches courantes de GKE.

Avant de commencer ce tutoriel, assurez-vous de connaître les éléments suivants :

Arrière-plan

Cette section décrit les principales technologies utilisées dans ce tutoriel.

TPU

Les Tensor Processing Units (TPU) sont des accélérateurs conçus sur mesure par Google pour les charges de travail de machine learning. Ce tutoriel est compatible avec les TPU Trillium (v6e) et Ironwood (TPU7x). Dans une tranche de TPU multi-hôte, plusieurs hôtes TPU communiquent directement via un réseau d'interconnexion entre puces (ICI) à haut débit. Cela fournit la mémoire à haut débit (HBM) agrégée et la communication collective à faible latence requises pour entraîner des modèles qui dépassent la capacité de mémoire d'une seule puce ou d'un seul hôte TPU.

TorchTPU

TorchTPU est un backend PyTorch Open Source pour Cloud TPU, développé en collaboration par Google et Meta. Avec TorchTPU, vous pouvez exécuter des modèles PyTorch existants sur du matériel TPU avec un minimum de modifications du code, et faire évoluer les charges de travail d'entraînement et d'inférence d'une seule puce à des pods TPU multitranches grâce au compilateur XLA (Accelerated Linear Algebra). TorchTPU utilise un modèle d'exécution à un processus par appareil sur le réseau ICI de TPU et est compatible avec les API PyTorch standards, y compris l'exécution eager, torch.compile, torch.distributed, DeviceMesh, DTensor et Fully Sharded Data Parallel (FSDP2).

KubeRay et Ray sur TPU

KubeRay est un opérateur Kubernetes permettant de déployer et de gérer des applications Ray. L'activation du module complémentaire Ray sur GKE déploie l'opérateur KubeRay géré et le webhook TPU KubeRay. Lorsque KubeRay crée des pods de nœuds de calcul qui demandent des ressources google.com/tpu, le webhook configure automatiquement les noms d'hôte des pods, le service sans adresse IP et les variables d'environnement TPU requises pour la coordination multi-hôte.

Ray détecte la topologie TPU sur chaque nœud au démarrage et planifie les charges de travail distribuées sur les tranches TPU. Vous pouvez orchestrer des charges de travail TorchTPU distribuées à deux niveaux :

Ray Core fournit des groupes d'emplacements tenant compte de la topologie et une planification des acteurs pour les boucles d'exécution distribuées personnalisées, telles que le calcul tensoriel parallèle ou les pipelines d'apprentissage par renforcement (RL).
Ray Train fournit un entraînement géré distribué en parallèle des données et en parallèle des modèles, y compris PyTorch Distributed Data Parallel (DDP), PyTorch FSDP2 et JAX SPMD (single program, multiple data), avec des points de contrôle intégrés et une récupération automatique après défaillance.

Objectifs

Ce tutoriel vous explique comment effectuer les tâches suivantes :

Créez un cluster GKE en mode Autopilot ou Standard sur TPU Trillium (v6e) ou Ironwood (TPU7x).
Créez une image de conteneur qui superpose les bibliothèques d'entraînement Ray et Hugging Face sur l'image de base TorchTPU.
Déployez une ressource personnalisée RayCluster pour provisionner les nœuds principaux et les nœuds de calcul TPU.
Exécutez un calcul distribué de tenseurs parallèles avec PyTorch DeviceMesh et DTensor sur une tranche TPU à l'aide de Ray Core.
Affinez un modèle Gemma 4 avec PyTorch FSDP2 sur une tranche de TPU multi-hôte à l'aide de Ray Train.

Coûts

Dans ce tutoriel, vous utilisez les composants facturables suivants de Google Cloud :

Obtenez une estimation des coûts en fonction de votre utilisation prévue à l'aide du simulateur de coût.

Une fois que vous avez terminé les tâches décrites dans ce tutoriel, vous pouvez éviter de continuer à payer des frais en supprimant les ressources que vous avez créées. Pour en savoir plus, consultez la section Effectuer un nettoyage.

Avant de commencer

Avant de commencer, effectuez les tâches suivantes :

Activez l'API Artifact Registry, l'API Cloud Storage et l'API Google Kubernetes Engine.

Activer les API

Si vous souhaitez utiliser la Google Cloud CLI pour cette tâche, installez puis initialisez la gcloud CLI. Si vous avez déjà installé la gcloud CLI, obtenez la dernière version en exécutant la commande gcloud components update. Il est possible que les versions antérieures de la gcloud CLI ne permettent pas d'exécuter les commandes de ce document.
Remarque : Pour les installations de la gcloud CLI existantes, veillez à définir la propriété compute/region. Si vous utilisez principalement des clusters zonaux, définissez plutôt compute/zone. En définissant un emplacement par défaut, vous pouvez éviter les erreurs gcloud CLI de ce type : One of [--zone, --region] must be supplied: Please specify location. Vous devrez peut-être spécifier l'emplacement dans certaines commandes si celui de votre cluster diffère de l'emplacement par défaut que vous avez défini.

Vérifiez que votre projet dispose d'un quota suffisant pour la capacité TPU Trillium (v6e) ou Ironwood (TPU7x) dans la zone sélectionnée. Pour en savoir plus, consultez Quotas Cloud TPU.
Assurez-vous que votre compte dispose des rôles IAM suivants :
- Administrateur Kubernetes Engine (roles/container.admin)
- Administrateur Artifact Registry (roles/artifactregistry.admin)
- Administrateur de l'espace de stockage (roles/storage.admin)

Préparer votre environnement

Dans ce tutoriel, vous utilisez Cloud Shell pour gérer les ressources hébergées sur Google Cloud. Cloud Shell est préinstallé avec les logiciels dont vous avez besoin pour ce tutoriel, y compris kubectl, docker, python3 et la CLI gcloud.

Pour configurer votre environnement avec Cloud Shell, procédez comme suit :

Dans la console Google Cloud , lancez une session Cloud Shell en cliquant sur Activer Cloud Shell dans la consoleGoogle Cloud . Une session s'ouvre dans le volet inférieur de la console Google Cloud .

Clonez le dépôt kubernetes-engine-samples et accédez au répertoire de l'exemple pour ce tutoriel :

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

Créez et activez un environnement virtuel Python local avec la CLI Ray installée :

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

Pour télécharger les pondérations du modèle Gemma 4, acceptez le contrat de licence Google sur Hugging Face :
1. Accédez à la page du modèle Gemma 4 sur Hugging Face.
2. Connectez-vous et acceptez les conditions de licence en cliquant sur Accepter et accéder au dépôt.
3. Dans les paramètres de votre compte Hugging Face, générez un jeton d'accès utilisateur avec des autorisations Read.
Définissez les variables d'environnement pour votre projet Google Cloud , votre cluster, votre bucket Cloud Storage et votre jeton 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}
```
Remplacez les éléments suivants :
- PROJECT_ID : ID de votre projet Google Cloud .
- CLUSTER_NAME : nom de votre cluster GKE, par exemple ray-torchtpu-cluster.
- REGION : région où se trouvent votre dépôt Artifact Registry, votre bucket Cloud Storage et votre capacité TPU.
- ZONE : zone dans laquelle votre capacité de TPU est disponible. Pour connaître les régions et zones disponibles par version de TPU, consultez Régions et zones TPU.
- ARTIFACT_REGISTRY : nom de votre dépôt Docker Artifact Registry, par exemple ray-torchtpu.
- BUCKET_NAME : nom de votre bucket Cloud Storage pour stocker les points de contrôle Ray Train, par exemple PROJECT_ID-ray-torchtpu.
- HF_TOKEN : votre jeton d'accès Hugging Face.

Créer un cluster GKE

Vous pouvez exécuter des charges de travail d'entraînement TorchTPU distribuées dans un cluster GKE Autopilot ou Standard avec le module complémentaire Ray Operator géré activé.

Créez votre cluster GKE et configurez les identifiants pour l'outil de ligne de commande kubectl :
Autopilot
1. Créez un cluster GKE Autopilot régional avec le module complémentaire de l'opérateur Ray activé :
  
  gcloud container clusters create-auto ${CLUSTER_NAME} \ --project=${PROJECT_ID} \ --region=${REGION} \ --cluster-version=1.36 \ --enable-ray-operator
  
  Lorsque vous déploierez la ressource personnalisée RayCluster lors d'une étape ultérieure, GKE Autopilot provisionnera automatiquement le nœud principal du processeur et la tranche TPU multi-hôte qui correspondent à vos requêtes de ressources de pod et à vos sélecteurs de nœuds.
2. Configurez l'outil de ligne de commande kubectl pour communiquer avec votre cluster :
  
  gcloud container clusters get-credentials ${CLUSTER_NAME} \ --project=${PROJECT_ID} \ --region=${REGION}
Standard
1. Créez un cluster GKE Standard zonal avec un pool de nœuds de processeur pour héberger le pod de tête Ray, le streaming d'images, GKE Dataplane V2 et la fédération d'identité de charge de travail pour GKE activés :
  
  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
  Remarque : Si vous gérez vous-même KubeRay Open Source au lieu d'utiliser le module complémentaire GKE RayOperator, installez KubeRay et le webhook TPU KubeRay version 1.4.0 ou ultérieure pour l'injection de la variable d'environnement TorchTPU.
2. Configurez l'outil de ligne de commande kubectl pour communiquer avec votre cluster :
  
  gcloud container clusters get-credentials ${CLUSTER_NAME} \ --project=${PROJECT_ID} \ --zone=${ZONE}
3. Créez un pool de nœuds de tranche TPU multi-hôte de 16 appareils avec une topologie 4x4 sur quatre hôtes sur TPU Trillium (v6e) ou une topologie 2x2x2 sur deux hôtes sur 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
  
  Conseil : Pour utiliser une réservation TPU spécifique, ajoutez les options --reservation-affinity=specific et --reservation=RESERVATION_NAME à la commande node-pools create, où RESERVATION_NAME est le nom de votre réservation TPU. Pour utiliser des VM Spot, ajoutez l'option --spot.

Créez un bucket Cloud Storage pour stocker les points de contrôle Ray Train partagés entre les nœuds TPU multi-hôtes et accordez à la ressource de compte de service Kubernetes default l'accès au 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"

Créez un secret Kubernetes avec votre jeton d'accès Hugging Face afin que les pods Ray puissent télécharger les pondérations Gemma 4 protégées :
```
kubectl create secret generic hf-secret \
    --from-literal=hf_api_token=${HF_TOKEN}
```

Créer l'image de conteneur

Créez une image de conteneur qui installe Ray et les bibliothèques d'entraînement Hugging Face sur l'image de base TorchTPU prédéfinie, qui inclut TorchTPU et la bibliothèque libtpu.

Créez un dépôt Docker Artifact Registry et configurez l'authentification Docker :

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

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

Inspectez le fichier Dockerfile dans le répertoire de l'exemple :

FROM us-docker.pkg.dev/ml-oss-artifacts-transient/torch-tpu-docker-container/torch-tpu:nightly

# 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

L'entraînement sur une seule tranche avec la classe Ray TorchTrainer et la planification par chiplet sur Ironwood (TPU7x) nécessitent Ray version 2.58.0 ou ultérieure, et l'ajustement de Gemma 4 nécessite la bibliothèque transformers version 5.5.0 ou ultérieure.

Créez et transférez l'image vers Artifact Registry :

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

Déployer la ressource personnalisée RayCluster

Déployez une ressource personnalisée RayCluster qui provisionne un pod principal de processeur, un volume /dev/shm en mémoire pour le magasin d'objets en mémoire partagée Ray et un groupe de nœuds de calcul TPU multi-hôtes.

Examinez le fichier manifeste RayCluster pour votre version de TPU dans l'exemple de dépôt :

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

Sur Ironwood (TPU7x), chaque puce physique contient deux chiplets, ce qui fournit huit cœurs TPU logiques par hôte à quatre puces. Le fichier manifeste définit la valeur du champ TPU sur "8" dans la section resources et définit la variable d'environnement RAY_TPU_RESOURCE_PER_CHIP sur "2" pour annoncer et planifier huit ressources TPU logiques par hôte.

Appliquez le fichier manifeste pour créer la ressource personnalisée 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 -

Attendez que le pod principal RayCluster et tous les pods de nœuds de calcul TPU atteignent l'état Running :

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

Le résultat ressemble à ce qui suit :

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

Inspectez les variables d'environnement sur l'un des pods de nœuds de calcul pour vérifier que le webhook TPU KubeRay a injecté les variables de coordination TPU multi-hôte et de topologie 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"

Le résultat ressemble à ce qui suit :

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

Établissez une session de transfert de port vers le service principal Ray :
```
kubectl port-forward svc/ray-torchtpu-cluster-head-svc 8265:8265 >/dev/null 2>&1 &
```
Vous pouvez inspecter les nœuds de calcul TPU et leurs libellés ray.io/tpu-* détectés automatiquement dans l'onglet Cluster du tableau de bord Ray.

Exécuter le parallélisme de tenseur distribué avec Ray Core

Ray Core fournit la fonction slice_placement_group dans le module ray.util.tpu pour réserver une tranche de TPU multi-hôte intacte et la classe PlacementGroupSchedulingStrategy pour planifier un acteur Ray par appareil TPU logique sur la tranche.

Examinez le script ray_core_collectives.py dans l'exemple de dépôt :

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()

Le script exécute le parallélisme Tensor distribué sur la tranche TPU en trois étapes :

La fonction slice_placement_group réserve une tranche de TPU de 16 appareils et planifie un acteur TorchTPUSliceWorker sur chaque appareil TPU.
Chaque acteur calcule son rang global à partir de son index d'hôte et de l'ID de l'appareil local, puis initialise le groupe de processus tpu_dist.
Les fonctions init_device_mesh et parallelize_module segmentent les couches de feed-forward sur les 16 appareils TPU, vérifient le résultat par rapport à une référence de processeur sur un seul appareil et exécutent un collectif all_reduce sur ICI.

Envoyez le job Ray Core pour votre version 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

Le résultat ressemble à ce qui suit :

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

Affiner un modèle Gemma 4 avec Ray Train et PyTorch FSDP2

La classe TorchTrainer de Ray Train automatise la réservation de tranches TPU, l'emplacement des nœuds de calcul et l'initialisation de torch.distributed sur une tranche TPU multi-hôte.

Examinez le script ray_train_llm_finetune.py dans l'exemple de dépôt, qui affine le modèle google/gemma-4-E2B-it sur l'ensemble de données d'instructions 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()

Le script configure la classe TorchTrainer avec les composants suivants :

La classe ScalingConfig configure la topologie des nœuds de calcul distribués pour le job d'entraînement. La définition des paramètres use_tpu=True, num_workers=16, topology et accelerator_type réserve une tranche de TPU de 16 appareils et planifie un nœud de calcul par cœur de TPU.
La classe TorchConfig configure le groupe de processus torch.distributed sur les nœuds de calcul. La définition du paramètre backend="tpu_dist" connecte les 16 workers sur le réseau TPU ICI avant l'exécution de la fonction train_func.
La fonction create_dataloader télécharge et tokenise l'ensemble de données d'instructions, regroupe les séquences de jetons avec un objet PyTorch DataLoader et appelle la fonction ray.train.torch.prepare_data_loader pour que chaque nœud de calcul reçoive un segment disjoint de chaque lot global sur l'appareil TPU qui lui est attribué.
La fonction train_func charge le modèle avec une précision bfloat16 et appelle la fonction PyTorch FSDP2 fully_shard sur chaque bloc de propagation avant et sur le modèle racine pour segmenter les paramètres, les gradients et les états de l'optimiseur sur les 16 cœurs de TPU.
La classe Checkpoint et la fonction ray.train.report collectent le dictionnaire d'état complet du modèle au rang 0 à la fin de chaque epoch et importent le point de contrôle dans votre bucket Cloud Storage.

Envoyez le job d'affinage TorchTrainer FSDP2 pour votre version 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

Le résultat ressemble à ce qui suit :

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

Le résultat de l'entraînement affiche les éléments suivants :

L'epoch 1 inclut la compilation initiale du graphe XLA. Les époques 2 et 3 réutilisent les exécutables mis en cache de /dev/shm/torch_tpu_cache, atteignant plus de 20 000 jetons par seconde sur TPU Trillium (v6e) et 31 000 jetons par seconde sur Ironwood (TPU7x).
La perte d'entropie croisée globale synchronisée sur les 16 rangs de nœuds de calcul FSDP2 diminue de 3,14 à l'époque 1 à 2,20 à l'époque 3 sur les trois époques de réglage sur l'ensemble de données d'instructions databricks/databricks-dolly-15k.
À la fin de chaque époque, le rang 0 rassemble le dictionnaire d'état complet du modèle et importe le point de contrôle dans votre bucket Cloud Storage, en conservant le point de contrôle avec la perte la plus faible, comme configuré dans l'objet CheckpointConfig(num_to_keep=1, checkpoint_score_attribute="loss", checkpoint_score_order="min").

Observer les performances des charges de travail

Pour surveiller l'utilisation des TensorCores TPU, l'utilisation de la mémoire HBM et les métriques du cluster Ray dans Cloud Monitoring, consultez Collecter et afficher les journaux et les métriques des clusters Ray sur GKE et Afficher les métriques TPU sur le tableau de bord Ray.

Accéder au tableau de bord Ray

Pour inspecter l'état de vos acteurs Ray, de vos groupes d'emplacement et de vos jobs d'entraînement dans le tableau de bord Ray intégré, procédez comme suit :

Si le transfert de port vers le service principal Ray n'est pas déjà en cours d'exécution, démarrez-le en arrière-plan :
```
kubectl port-forward svc/ray-torchtpu-cluster-head-svc 8265:8265 >/dev/null 2>&1 &
```
Dans Cloud Shell, cliquez sur Aperçu sur le Web et sélectionnez Prévisualiser sur le port 8265, ou ouvrez http://localhost:8265 dans votre navigateur local.
Cliquez sur les onglets Jobs (Jobs), Cluster et Metrics (Métriques) pour inspecter les journaux des nœuds de calcul, l'état du groupe d'emplacement, les étiquettes des nœuds TPU et le graphique TPU Utilization by Node (Utilisation des TPU par nœud).

Effectuer un nettoyage

Pour éviter que les ressources utilisées dans ce tutoriel ne soient facturées sur votre compte Google Cloud , supprimez la ressource personnalisée RayCluster, le cluster GKE, le bucket Cloud Storage et le dépôt Artifact Registry :

Arrêtez le processus de transfert de port en arrière-plan, puis supprimez la ressource personnalisée RayCluster et le secret Kubernetes :

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

Supprimez le cluster GKE :

Autopilot

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

Standard

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

Supprimez le bucket Cloud Storage et son contenu de point de contrôle :
```
gcloud storage rm --recursive gs://${BUCKET_NAME}
```

Supprimez le dépôt Artifact Registry :

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

Étapes suivantes

En savoir plus sur PyTorch sur Cloud TPU
Consultez la documentation TorchTPU.
Découvrez comment diffuser des LLM sur des TPU multi-hôtes avec Ray et TorchTPU.
Découvrez comment optimiser l'entraînement de l'IA sur les TPU avec DWS et Kueue.
Consultez À propos des TPU dans GKE.
Consultez la documentation de Ray Train.

Affiner des modèles sur des TPU multi-hôtes avec Ray et TorchTPU Restez organisé à l'aide des collections Enregistrez et classez les contenus selon vos préférences.

Arrière-plan

TPU

TorchTPU

KubeRay et Ray sur TPU

Objectifs

Coûts

Avant de commencer

Préparer votre environnement

Créer un cluster GKE

Autopilot

Standard

Trillium (v6e)

Ironwood (TPU7x)

Créer l'image de conteneur

Déployer la ressource personnalisée RayCluster

Trillium (v6e)

Ironwood (TPU7x)

Trillium (v6e)

Ironwood (TPU7x)

Trillium (v6e)

Ironwood (TPU7x)

Trillium (v6e)

Ironwood (TPU7x)

Exécuter le parallélisme de tenseur distribué avec Ray Core

Trillium (v6e)

Ironwood (TPU7x)

Trillium (v6e)

Ironwood (TPU7x)

Affiner un modèle Gemma 4 avec Ray Train et PyTorch FSDP2

Trillium (v6e)

Ironwood (TPU7x)

Trillium (v6e)

Ironwood (TPU7x)

Observer les performances des charges de travail

Accéder au tableau de bord Ray

Effectuer un nettoyage

Autopilot

Standard

Étapes suivantes

Affiner des modèles sur des TPU multi-hôtes avec Ray et TorchTPU