Ajustar modelos em TPUs de vários hosts com o Ray e o TorchTPU

Autopilot Standard

Neste tutorial, mostramos como executar cargas de trabalho distribuídas do PyTorch e ajustar um modelo Gemma 4 em frações de Cloud TPU de vários hosts no Google Kubernetes Engine (GKE) usando Ray Core, Ray Train, PyTorch Fully Sharded Data Parallel (FSDP2) e TorchTPU.

Este tutorial é destinado a engenheiros de machine learning (ML), administradores e operadores de plataforma e especialistas em dados e IA que querem orquestrar cargas de trabalho de treinamento distribuídas do PyTorch em frações de TPU de vários hosts no Kubernetes. Para saber mais sobre papéis comuns e tarefas de exemplo referenciados no conteúdo do Google Cloud , consulte Tarefas e papéis de usuário comuns do GKE.

Antes de iniciar este tutorial, você precisa conhecer:

• Sobre os modos de operação do GKE
• Sobre TPUs no GKE
• Pools de nós de fração de TPU de vários hosts
• Sobre o Ray no GKE

Contexto

Esta seção descreve as principais tecnologias usadas neste tutorial.

TPUs

As Unidades de Processamento de Tensor (TPUs) são aceleradores personalizados do Google para cargas de trabalho de machine learning. Este tutorial é compatível com a TPU Trillium (v6e) e a Ironwood (TPU7x). Em uma fração de TPU de vários hosts, vários hosts de TPU se comunicam diretamente por uma rede de interconexão entre chips (ICI) de alta velocidade, fornecendo a memória de alta largura de banda (HBM) agregada e a comunicação coletiva de baixa latência necessárias para treinar modelos que excedem a capacidade de memória de um único chip ou host de TPU.

TorchTPU

O TorchTPU é um back-end de código aberto do PyTorch para Cloud TPU, desenvolvido em colaboração entre o Google e a Meta. Com o TorchTPU, é possível executar modelos do PyTorch em hardware de TPU com mudanças mínimas no código e escalonar cargas de trabalho de treinamento e inferência de um único chip para Pods de TPU de várias frações usando o compilador de álgebra linear acelerada (XLA, na sigla em inglês). O TorchTPU usa um modelo de execução de um processo por dispositivo na rede TPU ICI e é compatível com as APIs padrão do PyTorch, incluindo execução adiantada, torch.compile, torch.distributed, DeviceMesh, DTensor e Fully Sharded Data Parallel (FSDP2).

KubeRay e Ray em TPUs

O KubeRay é um operador do Kubernetes para implantar e gerenciar aplicativos do Ray. Ativar o complemento Ray no GKE implanta o operador gerenciado do KubeRay e o webhook de TPU do KubeRay. Quando o KubeRay cria pods de worker que solicitam recursos google.com/tpu, o webhook configura automaticamente os nomes de host dos pods, o serviço headless e as variáveis de ambiente de TPU necessárias para a coordenação de vários hosts.

O Ray detecta a topologia de TPU em cada nó na inicialização e programa cargas de trabalho distribuídas em frações de TPU. É possível orquestrar cargas de trabalho distribuídas do TorchTPU em dois níveis:

• O Ray Core oferece grupos de posicionamento com reconhecimento de topologia e programação de atores para loops de execução distribuída personalizados, como computação paralela de tensores ou pipelines de aprendizado por reforço (RL).
• O Ray Train oferece treinamento gerenciado distribuído em paralelo de dados e modelos, incluindo o PyTorch Distributed Data Parallel (DDP), o PyTorch FSDP2 e o JAX de programa único e vários dados (SPMD), com checkpoint integrado e recuperação automática de falhas.

Objetivos

Este tutorial mostra como fazer o seguinte:

1. Crie um cluster do GKE no Autopilot ou no Standard na TPU Trillium (v6e) ou Ironwood (TPU7x).
2. Crie uma imagem de contêiner que sobreponha as bibliotecas de treinamento do Ray e do Hugging Face à imagem de base do TorchTPU.
3. Implante um recurso personalizado do RayCluster para provisionar os nós principais e de worker da TPU.
4. Execute computação distribuída em paralelo de tensores com o PyTorch DeviceMesh e DTensor em uma fração de TPU usando o Ray Core.
5. Ajustar um modelo do Gemma 4 com o PyTorch FSDP2 em uma fração de TPU de vários hosts usando o Ray Train.

Custos

Neste tutorial, você vai usar os seguintes componentes faturáveis do Google Cloud:

• GKE
• Cloud TPU
• Artifact Registry
• Cloud Storage

Para gerar uma estimativa de custo baseada na projeção de uso deste tutorial, use a calculadora de preços.

Ao concluir as tarefas descritas neste tutorial, é possível evitar o faturamento contínuo excluindo os recursos criados. Para mais informações, consulte Limpeza.

Antes de começar

Antes de começar, verifique se você realizou as tarefas a seguir:

• Ative a API Artifact Registry, a API Cloud Storage e a API Google Kubernetes Engine.
Ativar APIs
• Para usar a Google Cloud CLI nesta tarefa, instale e inicialize a CLI gcloud. Se você instalou a CLI gcloud anteriormente, instale a versão mais recente executando o comando gcloud components update. Talvez as versões anteriores da CLI gcloud não sejam compatíveis com a execução dos comandos neste documento.
  Observação:em instalações prévias da CLI gcloud, defina a propriedade compute/region. Se você usa principalmente clusters zonais, defina compute/zone. Ao definir um local padrão, é possível evitar erros na CLI gcloud como: One of [--zone, --region] must be supplied: Please specify location. Talvez seja necessário especificar o local em determinados comandos se o local do cluster for diferente do padrão definido.

• Verifique se o projeto tem cota suficiente para a capacidade de TPU Trillium (v6e) ou Ironwood (TPU7x) na zona selecionada. Para mais informações, consulte Cotas do Cloud TPU.
• Verifique se a conta tem os seguintes papéis do IAM:
  • Administrador do Kubernetes Engine (roles/container.admin)
  • Administrador do Artifact Registry (roles/artifactregistry.admin)
  • Administrador do Storage (roles/storage.admin)

Preparar o ambiente

Neste tutorial, você usará o Cloud Shell para gerenciar recursos hospedados no Google Cloud. O Cloud Shell vem pré-instalado com o software necessário para este tutorial, incluindo kubectl, docker, python3 e a CLI gcloud.

Para configurar o ambiente com o Cloud Shell, siga estas etapas:

1. No console do Google Cloud , inicie uma sessão do Cloud Shell clicando em Ativar o Cloud Shell no console doGoogle Cloud . Isso inicia uma sessão no painel inferior do console do Google Cloud .

2. Clone o repositório kubernetes-engine-samples e mude para o diretório de exemplo deste tutorial:

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

3. Crie e ative um ambiente virtual local do Python com a CLI do Ray instalada:

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

4. Para baixar os pesos do modelo Gemma 4, aceite o contrato de licença do Google no Hugging Face:

   1. Acesse a página do modelo Gemma 4 no Hugging Face.
   2. Faça login e aceite os termos de licença clicando em Concordar e acessar o repositório.
   3. Nas configurações da sua conta do Hugging Face, gere um token de acesso de usuário com permissões Read.

5. Defina as variáveis de ambiente para o projeto do Google Cloud , o cluster, o bucket do Cloud Storage e o token do 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}
   

   Substitua:

   • PROJECT_ID: o ID do projeto Google Cloud .
   • CLUSTER_NAME: o nome do cluster do GKE, como ray-torchtpu-cluster.
   • REGION: a região em que o repositório do Artifact Registry, o bucket do Cloud Storage e a capacidade de TPU estão localizados.
   • ZONE: a zona em que sua capacidade de TPU está disponível. Para ver as regiões e zonas disponíveis por versão de TPU, consulte Regiões e zonas de TPU.
   • ARTIFACT_REGISTRY: o nome do repositório do Docker do Artifact Registry, como ray-torchtpu.
   • BUCKET_NAME: o nome do bucket do Cloud Storage para armazenar checkpoints do Ray Train, como PROJECT_ID-ray-torchtpu.
   • HF_TOKEN: seu token de acesso do Hugging Face.

Criar um cluster do GKE

É possível executar cargas de trabalho de treinamento distribuídas do TorchTPU em um cluster do GKE Autopilot ou Standard com o complemento gerenciado do operador do Ray ativado.

1. Crie o cluster do GKE e configure as credenciais para a ferramenta de linha de comando kubectl:

   Piloto automático

   1. Crie um cluster regional do GKE Autopilot com o complemento do operador do Ray ativado:

      gcloud container clusters create-auto ${CLUSTER_NAME} \
          --project=${PROJECT_ID} \
          --region=${REGION} \
          --cluster-version=1.36 \
          --enable-ray-operator
      

      Quando você implanta o recurso personalizado RayCluster em uma etapa posterior, o GKE Autopilot provisiona automaticamente o nó principal da CPU e a fração de TPU de vários hosts que correspondem às solicitações de recursos do pod e aos seletores de nós.

   2. Configure a ferramenta de linha de comando kubectl para se comunicar com o cluster:

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

   Padrão

   1. Crie um cluster zonal do GKE Standard com um pool de nós de CPU para hospedar o pod principal do Ray, o streaming de imagens, o GKE Dataplane V2 e a Federação de Identidade da Carga de Trabalho para GKE ativados:

      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
      

      Observação: se você autogerenciar o KubeRay de código aberto em vez de usar o complemento RayOperator do GKE, instale o KubeRay e o webhook de TPU do KubeRay versão 1.4.0 ou mais recente para injeção variável de ambiente do TorchTPU.

   2. Configure a ferramenta de linha de comando kubectl para se comunicar com o cluster:

      gcloud container clusters get-credentials ${CLUSTER_NAME} \
          --project=${PROJECT_ID} \
          --zone=${ZONE}
      

   3. Crie um pool de nós de fração de TPU de vários hosts com 16 dispositivos com uma topologia 4x4 em quatro hosts no TPU Trillium (v6e) ou uma topologia 2x2x2 em dois hosts no 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
      

      Dica: para consumir uma reserva de TPU específica, anexe as flags --reservation-affinity=specific e --reservation=RESERVATION_NAME ao comando node-pools create, em que RESERVATION_NAME é o nome da reserva de TPU. Para usar VMs spot, adicione a flag --spot.

2. Crie um bucket do Cloud Storage para armazenar checkpoints compartilhados do Ray Train nos nós de TPU de vários hosts e conceda ao recurso default da conta de serviço do Kubernetes acesso ao 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"
   

3. Crie um Secret do Kubernetes com seu token de acesso do Hugging Face para que os pods do Ray possam baixar os pesos restritos do Gemma 4:

   kubectl create secret generic hf-secret \
       --from-literal=hf_api_token=${HF_TOKEN}
   

Crie a imagem do contêiner

Crie uma imagem de contêiner que instale o Ray e as bibliotecas de treinamento do Hugging Face sobre a imagem de base pré-criada do TorchTPU, que inclui o TorchTPU e a biblioteca libtpu.

1. Crie um repositório Docker do Artifact Registry e configure a autenticação do Docker:

   gcloud artifacts repositories create ${ARTIFACT_REGISTRY} \
       --repository-format=docker \
       --location=${REGION}
   
   gcloud auth configure-docker ${REGION}-docker.pkg.dev,us-docker.pkg.dev
   

2. Inspecione o Dockerfile no diretório de amostra:

   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

   O treinamento de fração única com a classe TorchTrainer do Ray e a programação por chiplet no Ironwood (TPU7x) exigem o Ray versão 2.58.0 ou mais recente, e o ajuste fino do Gemma 4 exige a versão 5.5.0 ou mais recente da biblioteca transformers.

3. Crie e envie a imagem para o Artifact Registry:

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

Implantar o recurso personalizado RayCluster

Implante um recurso personalizado do RayCluster que provisiona um pod principal de CPU, um volume /dev/shm na memória para o repositório de objetos de memória compartilhada do Ray e um grupo de workers da TPU de vários hosts.

1. Revise o manifesto do RayCluster para sua versão de TPU no repositório de amostra:

   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

   No Ironwood (TPU7x), cada chip físico contém dois chiplets, fornecendo oito núcleos lógicos de TPU por host de quatro chips. O manifesto define o valor do campo TPU como "8" na seção resources e define a variável de ambiente RAY_TPU_RESOURCE_PER_CHIP como "2" para anunciar e programar oito recursos lógicos TPU por host.

2. Aplique o manifesto para criar o recurso personalizado do 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 -
   

3. Aguarde até que o pod principal do RayCluster e todos os pods de worker da TPU alcancem o estado Running:

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

   O resultado será o seguinte:

   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
   

4. Inspecione as variáveis de ambiente em um dos pods de worker para confirmar se o webhook da TPU do KubeRay injetou as variáveis de coordenação da TPU de vários hosts e de topologia do 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"
   

   O resultado será o seguinte:

   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
   

5. Estabeleça uma sessão de encaminhamento de portas para o serviço de cabeça do Ray:

   kubectl port-forward svc/ray-torchtpu-cluster-head-svc 8265:8265 >/dev/null 2>&1 &
   

   É possível inspecionar os nós de trabalho da TPU e os rótulos ray.io/tpu-* detectados automaticamente na guia Cluster do painel do Ray.

Executar o paralelismo de tensor distribuído com o Ray Core

O Ray Core fornece a função slice_placement_group no módulo ray.util.tpu para reservar uma fração de TPU de vários hosts intacta e a classe PlacementGroupSchedulingStrategy para programar um ator do Ray por dispositivo de TPU lógico na fração.

1. Revise o script ray_core_collectives.py no repositório de amostra:

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

   O script executa o paralelismo de tensor distribuído na fração de TPU em três etapas:

   1. A função slice_placement_group reserva uma fração de TPU de 16 dispositivos e programa um ator TorchTPUSliceWorker em cada dispositivo de TPU.
   2. Cada ator calcula a classificação global com base no índice do host e no ID do dispositivo local e, em seguida, inicializa o grupo de processos tpu_dist.
   3. As funções init_device_mesh e parallelize_module fragmentam as camadas de feed-forward em todos os 16 dispositivos de TPU, verificam a saída em relação a uma referência de CPU de dispositivo único e executam um coletivo all_reduce sobre a ICI.

2. Envie o job do Ray Core para sua versão da 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
   

   O resultado será o seguinte:

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

Ajustar um modelo do Gemma 4 com o Ray Train e o PyTorch FSDP2

A classe TorchTrainer do Ray Train automatiza a reserva de frações de TPU, o posicionamento de workers e a inicialização de torch.distributed em uma fração de TPU de vários hosts.

1. Analise o script ray_train_llm_finetune.py no repositório de amostra, que ajusta o modelo google/gemma-4-E2B-it no conjunto de dados de instruções 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()

   O script configura a classe TorchTrainer com os seguintes componentes:

   • A classe ScalingConfig configura a topologia de worker distribuído para o job de treinamento. Definir os parâmetros use_tpu=True, num_workers=16, topology e accelerator_type reserva uma fração de TPU de 16 dispositivos e programa um worker por núcleo de TPU.
   • A classe TorchConfig configura o grupo de processos torch.distributed nos workers. Definir o parâmetro backend="tpu_dist" conecta todos os 16 workers pela rede ICI da TPU antes que a função train_func seja executada.
   • A função create_dataloader faz o download e a tokenização do conjunto de dados de instruções, agrupa as sequências de tokens com um objeto DataLoader do PyTorch e chama a função ray.train.torch.prepare_data_loader para que cada worker receba um fragmento separado de cada lote global no dispositivo de TPU atribuído.
   • A função train_func carrega o modelo com precisão bfloat16 e chama a função fully_shard do PyTorch FSDP2 em cada bloco de feed-forward e no modelo raiz para fragmentar parâmetros, gradientes e estados do otimizador em todos os 16 núcleos de TPU.
   • A classe Checkpoint e a função ray.train.report coletam o dicionário de estado do modelo completo na classificação 0 ao final de cada época e fazem upload do checkpoint para o bucket do Cloud Storage.

2. Envie o job de ajuste fino do FSDP2 TorchTrainer para sua versão 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
   

   Dica :por padrão, o script treina em um subconjunto de 1.024 exemplos por três períodos para que o job seja concluído em alguns minutos. Para treinar em todos os 15.011 exemplos ou alterar o número de períodos, transmita a sinalização --dataset-split=train ou --num-epochs para o script ray_train_llm_finetune.py.

   O resultado será o seguinte:

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

   A saída do treinamento mostra o seguinte:

   • A época 1 inclui a compilação inicial do gráfico XLA. As épocas 2 e 3 reutilizam os executáveis em cache de /dev/shm/torch_tpu_cache, alcançando mais de 20.000 tokens por segundo no TPU Trillium (v6e) e 31.000 tokens por segundo no Ironwood (TPU7x).
   • A perda global de entropia cruzada sincronizada em todas as 16 classificações de worker do FSDP2 diminui de 3,14 na época 1 para 2,20 na época 3 nas três épocas de ajuste no conjunto de dados de instruções databricks/databricks-dolly-15k.
   • No final de cada época, a classificação 0 coleta o dicionário de estado do modelo completo e faz upload do checkpoint para o bucket do Cloud Storage, mantendo o checkpoint de menor perda, conforme configurado no objeto CheckpointConfig(num_to_keep=1, checkpoint_score_attribute="loss", checkpoint_score_order="min").

Sucesso :você executou uma computação distribuída em paralelo de tensor com o Ray Core e ajustou um modelo Gemma 4 com o PyTorch FSDP2 em TPUs no GKE usando o Ray Train e o TorchTPU.

Observar o desempenho da carga de trabalho

Para monitorar a utilização do TensorCore da TPU, o uso da HBM e as métricas do cluster do Ray no Cloud Monitoring, consulte Coletar e visualizar registros e métricas para clusters do Ray no GKE e Visualizar métricas da TPU no painel do Ray.

Acessar o painel do Ray

Para inspecionar o status dos atores, grupos de posicionamento e jobs de treinamento do Ray no painel integrado do Ray, siga estas etapas:

1. Se o encaminhamento de portas para o serviço de cabeça do Ray ainda não estiver em execução, inicie-o em segundo plano:

   kubectl port-forward svc/ray-torchtpu-cluster-head-svc 8265:8265 >/dev/null 2>&1 &
   

2. No Cloud Shell, clique em Visualização da Web e selecione Visualizar na porta 8265 ou abra http://localhost:8265 no navegador local.

3. Clique nas guias Jobs, Cluster e Métricas para inspecionar os registros de workers, o estado do grupo de posicionamento, os rótulos dos nós de TPU e o gráfico Uso de TPU por nó.

Limpar

Para evitar cobranças na sua conta do Google Cloud pelos recursos usados neste tutorial, exclua o recurso personalizado RayCluster, o cluster do GKE, o bucket do Cloud Storage e o repositório do Artifact Registry:

1. Interrompa o processo de encaminhamento de portas em segundo plano e exclua o recurso personalizado RayCluster e o secret do Kubernetes:

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

2. Exclua o cluster do GKE:

   Piloto automático

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

   Padrão

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

3. Exclua o bucket do Cloud Storage e o conteúdo do checkpoint:

   gcloud storage rm --recursive gs://${BUCKET_NAME}
   

4. Exclua o repositório do Artifact Registry:

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

A seguir

• Saiba mais sobre o PyTorch no Cloud TPU.
• Consulte a documentação do TorchTPU.
• Saiba como disponibilizar LLMs em TPUs de vários hosts com o Ray e o TorchTPU.
• Saiba como otimizar o treinamento de IA em TPUs com o DWS e o Kueue.
• Leia Sobre TPUs no GKE.
• Confira a documentação do Ray Train.