Scalare i workload ML utilizzando Ray

Questo documento fornisce dettagli su come eseguire workload di machine learning (ML) con Ray e JAX sulle TPU. Esistono due diverse modalità per utilizzare le TPU con Ray: modalità incentrata sul dispositivo (PyTorch/XLA) e modalità incentrata sull'host (JAX).

Questo documento presuppone che tu abbia già configurato un ambiente TPU. Per maggiori informazioni, consulta le seguenti risorse:

• Cloud TPU: configura l'ambiente Cloud TPU e gestisci le risorse TPU
• Google Kubernetes Engine (GKE): Esegui il deployment dei carichi di lavoro TPU in GKE Autopilot o Esegui il deployment dei carichi di lavoro TPU in GKE Standard

Modalità incentrata sul dispositivo (PyTorch/XLA)

La modalità incentrata sul dispositivo mantiene gran parte dello stile programmatico di PyTorch classico. In questa modalità, aggiungi un nuovo tipo di dispositivo XLA, che funziona come qualsiasi altro dispositivo PyTorch. Ogni singolo processo interagisce con un dispositivo XLA.

Questa modalità è ideale se hai già familiarità con PyTorch con le GPU e vuoi utilizzare astrazioni di codifica simili.

Le sezioni seguenti descrivono come eseguire un carico di lavoro PyTorch/XLA su uno o più dispositivi senza utilizzare Ray, quindi come eseguire lo stesso carico di lavoro su più host utilizzando Ray.

Crea una TPU

Nota:questa sezione si applica all'API Cloud TPU. Per Ironwood (TPU7x), devi utilizzare Google Kubernetes Engine (GKE). Per maggiori informazioni, consulta Informazioni sulle TPU in GKE.

1. Crea variabili di ambiente per i parametri di creazione della TPU.

   export PROJECT_ID=your-project-id
   export TPU_NAME=your-tpu-name
   export ZONE=europe-west4-b
   export ACCELERATOR_TYPE=v5p-8
   export RUNTIME_VERSION=v2-alpha-tpuv5

   Descrizioni delle variabili di ambiente

   Variabile	Descrizione
   PROJECT_ID	L'ID progetto Google Cloud . Utilizza un progetto esistente o creane uno nuovo.
   TPU_NAME	Il nome della TPU.
   ZONE	La zona in cui creare la VM TPU. Per saperne di più sulle zone supportate, consulta Regioni e zone TPU.
   ACCELERATOR_TYPE	Il tipo di acceleratore specifica la versione e le dimensioni della Cloud TPU che vuoi creare. Per saperne di più sui tipi di acceleratore supportati per ogni versione di TPU, consulta la sezione Versioni di TPU.
   RUNTIME_VERSION	La versione software di Cloud TPU.

2. Utilizza il seguente comando per creare una VM TPU v5p con 8 core:

   gcloud compute tpus tpu-vm create $TPU_NAME \
      --zone=$ZONE \
      --accelerator-type=$ACCELERATOR_TYPE  \
      --version=$RUNTIME_VERSION

3. Connettiti alla VM TPU utilizzando il seguente comando:

   gcloud compute tpus tpu-vm ssh $TPU_NAME --zone=$ZONE

Se utilizzi GKE, consulta la guida KubeRay su GKE per informazioni sulla configurazione.

Requisiti di installazione

Esegui questi comandi sulla VM TPU per installare le dipendenze richieste:

1. Salva quanto segue in un file. Ad esempio, requirements.txt.

   --find-links https://storage.googleapis.com/libtpu-releases/index.html
   --find-links https://storage.googleapis.com/libtpu-wheels/index.html
   torch~=2.6.0
   torch_xla[tpu]~=2.6.0
   ray[default]==2.40.0
   

2. Per installare le dipendenze richieste, esegui:

   pip install -r requirements.txt
   

Se esegui il tuo workload su GKE, ti consigliamo di creare un Dockerfile che installi le dipendenze richieste. Per un esempio, consulta Esegui il tuo carico di lavoro sui nodi slice TPU nella documentazione di GKE.

Esegui un workload PyTorch/XLA su un singolo dispositivo

L'esempio seguente mostra come creare un tensore XLA su un singolo dispositivo, ovvero un chip TPU. Questo è simile al modo in cui PyTorch gestisce altri tipi di dispositivi.

1. Salva il seguente snippet di codice in un file. Ad esempio, workload.py.

   import torch
   import torch_xla
   import torch_xla.core.xla_model as xm
   
   t = torch.randn(2, 2, device=xm.xla_device())
   print(t.device)
   print(t)
   

   L'istruzione di importazione import torch_xla inizializza PyTorch/XLA e la funzione xm.xla_device() restituisce il dispositivo XLA corrente, un chip TPU.

2. Imposta la variabile di ambiente PJRT_DEVICE su TPU.

   export PJRT_DEVICE=TPU
   

3. Esegui lo script.

   python workload.py
   

   L'output è simile al seguente. Assicurati che l'output indichi che il dispositivo XLA è stato trovato.

   xla:0
   tensor([[ 0.6220, -1.4707],
           [-1.2112,  0.7024]], device='xla:0')
   

Esegui PyTorch/XLA su più dispositivi

1. Aggiorna lo snippet di codice della sezione precedente in modo che venga eseguito su più dispositivi.

   import torch
   import torch_xla
   import torch_xla.core.xla_model as xm
   
   def _mp_fn(index):
       t = torch.randn(2, 2, device=xm.xla_device())
       print(t.device)
       print(t)
   
   if __name__ == '__main__':
       torch_xla.launch(_mp_fn, args=())
   

2. Esegui lo script.

   python workload.py
   

   Se esegui lo snippet di codice su una TPU v5p-8, l'output è simile al seguente:

   xla:0
   xla:0
   xla:0
   tensor([[ 1.2309,  0.9896],
           [ 0.5820, -1.2950]], device='xla:0')
   xla:0
   tensor([[ 1.2309,  0.9896],
           [ 0.5820, -1.2950]], device='xla:0')
   tensor([[ 1.2309,  0.9896],
           [ 0.5820, -1.2950]], device='xla:0')
   tensor([[ 1.2309,  0.9896],
           [ 0.5820, -1.2950]], device='xla:0')
   

torch_xla.launch() accetta due argomenti: una funzione e un elenco di parametri. Crea un processo per ogni dispositivo XLA disponibile e chiama la funzione specificata negli argomenti. In questo esempio, sono disponibili 4 dispositivi TPU, quindi torch_xla.launch() crea 4 processi e chiama _mp_fn() su ogni dispositivo. Ogni processo ha accesso a un solo dispositivo, quindi ogni dispositivo ha l'indice 0 e xla:0 viene stampato per tutti i processi.

Esegui PyTorch/XLA su più host con Ray

Nota:questa sezione si applica all'API Cloud TPU. Per Ironwood (TPU7x), devi utilizzare Google Kubernetes Engine (GKE). Per maggiori informazioni, consulta Informazioni sulle TPU in GKE.

Le sezioni seguenti mostrano come eseguire lo stesso snippet di codice su una sezione TPU multihost più grande. Per ulteriori informazioni sull'architettura TPU multi-host, vedi Architettura di sistema.

In questo esempio, configuri manualmente Ray. Se hai già dimestichezza con la configurazione di Ray, puoi passare all'ultima sezione, Esegui un carico di lavoro Ray. Per ulteriori informazioni sulla configurazione di Ray per un ambiente di produzione, consulta le seguenti risorse:

• Avvio di cluster Ray
• Avvia Google Cloud il cluster GKE con TPU per KubeRay

Crea una VM TPU multi-host

1. Crea variabili di ambiente per i parametri di creazione della TPU.

   export PROJECT_ID=your-project-id
   export TPU_NAME=your-tpu-name
   export ZONE=europe-west4-b
   export ACCELERATOR_TYPE=v5p-16
   export RUNTIME_VERSION=v2-alpha-tpuv5

   Descrizioni delle variabili di ambiente

   Variabile	Descrizione
   PROJECT_ID	L'ID progetto Google Cloud . Utilizza un progetto esistente o creane uno nuovo.
   TPU_NAME	Il nome della TPU.
   ZONE	La zona in cui creare la VM TPU. Per saperne di più sulle zone supportate, consulta Regioni e zone TPU.
   ACCELERATOR_TYPE	Il tipo di acceleratore specifica la versione e le dimensioni della Cloud TPU che vuoi creare. Per saperne di più sui tipi di acceleratore supportati per ogni versione di TPU, consulta la sezione Versioni di TPU.
   RUNTIME_VERSION	La versione software di Cloud TPU.

2. Crea una TPU v5p multi-host con 2 host (una v5p-16, con 4 chip TPU su ogni host) utilizzando il seguente comando:

   gcloud compute tpus tpu-vm create $TPU_NAME \
      --zone=$ZONE \
      --accelerator-type=$ACCELERATOR_TYPE \
      --version=$RUNTIME_VERSION

Configurare Ray

Una TPU v5p-16 ha due host TPU, ognuno con 4 chip TPU. In questo esempio, avvierai il nodo head di Ray su un host e aggiungerai il secondo host come nodo worker al cluster Ray.

1. Connettiti al primo host utilizzando SSH.

   gcloud compute tpus tpu-vm ssh $TPU_NAME --zone=$ZONE --worker=0

2. Installa le dipendenze con lo stesso file dei requisiti della sezione Installazione dei requisiti.

   pip install -r requirements.txt
   

3. Avvia la procedura di Ray.

   ray start --head --port=6379
   

   Nota: se viene visualizzato un errore come Command 'ray' not found, potrebbe essere necessario aggiungere la directory contenente l'eseguibile ray alla variabile di ambiente PATH del sistema. Puoi anche specificare il percorso completo quando esegui i comandi ray (ad esempio, ~/.local/bin/ray).
   Nota: il comando ray start ti chiede di attivare o disattivare la raccolta delle statistiche di utilizzo. Per saperne di più, consulta la sezione Raccolta delle statistiche di utilizzo nella documentazione di Ray.

   L'output è simile al seguente:

   Enable usage stats collection? This prompt will auto-proceed in 10 seconds to avoid blocking cluster startup. Confirm [Y/n]: y
   Usage stats collection is enabled. To disable this, add `--disable-usage-stats` to the command that starts the cluster, or run the following command: `ray disable-usage-stats` before starting the cluster. See https://docs.ray.io/en/master/cluster/usage-stats.html for more details.
   
   Local node IP: 10.130.0.76
   
   --------------------
   Ray runtime started.
   --------------------
   
   Next steps
   To add another node to this Ray cluster, run
       ray start --address='10.130.0.76:6379'
   
   To connect to this Ray cluster:
       import ray
       ray.init()
   
   To terminate the Ray runtime, run
       ray stop
   
   To view the status of the cluster, use
       ray status
   

   Questo host TPU è ora il nodo head di Ray. Prendi nota delle righe che mostrano come aggiungere un altro nodo al cluster Ray, in modo simile a quanto segue:

   To add another node to this Ray cluster, run
       ray start --address='10.130.0.76:6379'
   

   Utilizzerai questo comando in un passaggio successivo.

4. Controlla lo stato del cluster Ray:

   ray status
   

   L'output è simile al seguente:

   ======== Autoscaler status: 2025-01-14 22:03:39.385610 ========
   Node status
   ---------------------------------------------------------------
   Active:
   1 node_bc0c62819ddc0507462352b76cc06b462f0e7f4898a77e5133c16f79
   Pending:
   (no pending nodes)
   Recent failures:
   (no failures)
   
   Resources
   ---------------------------------------------------------------
   Usage:
   0.0/208.0 CPU
   0.0/4.0 TPU
   0.0/1.0 TPU-v5p-16-head
   0B/268.44GiB memory
   0B/119.04GiB object_store_memory
   0.0/1.0 your-tpu-name
   
   Demands:
   (no resource demands)
   

   Il cluster contiene solo 4 TPU (0.0/4.0 TPU) perché finora hai aggiunto solo il nodo head.

   Ora che il nodo principale è in esecuzione, puoi aggiungere il secondo host al cluster.

5. Connettiti al secondo host utilizzando SSH.

   gcloud compute tpus tpu-vm ssh $TPU_NAME --zone=$ZONE --worker=1

6. Installa le dipendenze con lo stesso file dei requisiti della sezione Installare i requisiti.

   pip install -r requirements.txt
   

7. Avvia la procedura di Ray. Per aggiungere questo nodo al cluster Ray esistente, utilizza il comando dall'output del comando ray start. Assicurati di sostituire l'indirizzo IP e la porta nel comando seguente:

   ray start --address='10.130.0.76:6379'

   L'output è simile al seguente:

   Local node IP: 10.130.0.80
   [2025-01-14 22:30:07,397 W 75572 75572] global_state_accessor.cc:463: Retrying to get node with node ID 35f9ac0675c91429805cdc1b97c3713422d97eee783ccb0c0304f5c1
   
   --------------------
   Ray runtime started.
   --------------------
   
   To terminate the Ray runtime, run
   ray stop
   

8. Controlla di nuovo lo stato di Ray:

   ray status
   

   L'output è simile al seguente:

   ======== Autoscaler status: 2025-01-14 22:45:21.485617 ========
   Node status
   ---------------------------------------------------------------
   Active:
   1 node_bc0c62819ddc0507462352b76cc06b462f0e7f4898a77e5133c16f79
   1 node_35f9ac0675c91429805cdc1b97c3713422d97eee783ccb0c0304f5c1
   Pending:
   (no pending nodes)
   Recent failures:
   (no failures)
   
   Resources
   ---------------------------------------------------------------
   Usage:
   0.0/416.0 CPU
   0.0/8.0 TPU
   0.0/1.0 TPU-v5p-16-head
   0B/546.83GiB memory
   0B/238.35GiB object_store_memory
   0.0/2.0 your-tpu-name
   
   Demands:
   (no resource demands)
   

   Il secondo host TPU è ora un nodo del cluster. L'elenco delle risorse disponibili ora mostra 8 TPU (0.0/8.0 TPU).

   Nota: devi avviare il processo Ray su ogni VM TPU prima di eseguire il workload.

Esegui un workload Ray

1. Aggiorna lo snippet di codice da eseguire sul cluster Ray:

   import os
   import torch
   import torch_xla
   import torch_xla.core.xla_model as xm
   import ray
   
   import torch.distributed as dist
   import torch_xla.runtime as xr
   from torch_xla._internal import pjrt
   
   # Defines the local PJRT world size, the number of processes per host.
   LOCAL_WORLD_SIZE = 4
   # Defines the number of hosts in the Ray cluster.
   NUM_OF_HOSTS = 4
   GLOBAL_WORLD_SIZE = LOCAL_WORLD_SIZE * NUM_OF_HOSTS
   
   def init_env():
       local_rank = int(os.environ['TPU_VISIBLE_CHIPS'])
   
       pjrt.initialize_multiprocess(local_rank, LOCAL_WORLD_SIZE)
       xr._init_world_size_ordinal()
   
   # This decorator signals to Ray that the `print_tensor()` function should be run on a single TPU chip.
   @ray.remote(resources={"TPU": 1})
   def print_tensor():
       # Initializes the runtime environment on each Ray worker. Equivalent to
       # the `torch_xla.launch call` in the Run PyTorch/XLA on multiple devices section.
       init_env()
   
       t = torch.randn(2, 2, device=xm.xla_device())
       print(t.device)
       print(t)
   
   ray.init()
   
   # Uses Ray to dispatch the function call across available nodes in the cluster.
   tasks = [print_tensor.remote() for _ in range(GLOBAL_WORLD_SIZE)]
   ray.get(tasks)
   
   ray.shutdown()
   

2. Esegui lo script sul nodo head di Ray. Sostituisci ray-workload.py con il percorso dello script.

   python ray-workload.py

   L'output è simile al seguente:

   WARNING:root:libtpu.so and TPU device found. Setting PJRT_DEVICE=TPU.
   xla:0
   xla:0
   xla:0
   xla:0
   xla:0
   tensor([[ 0.6220, -1.4707],
           [-1.2112,  0.7024]], device='xla:0')
   tensor([[ 0.6220, -1.4707],
           [-1.2112,  0.7024]], device='xla:0')
   xla:0
   xla:0
   tensor([[ 0.6220, -1.4707],
           [-1.2112,  0.7024]], device='xla:0')
   tensor([[ 0.6220, -1.4707],
           [-1.2112,  0.7024]], device='xla:0')
   tensor([[ 0.6220, -1.4707],
           [-1.2112,  0.7024]], device='xla:0')
   tensor([[ 0.6220, -1.4707],
           [-1.2112,  0.7024]], device='xla:0')
   tensor([[ 0.6220, -1.4707],
           [-1.2112,  0.7024]], device='xla:0')
   xla:0
   tensor([[ 0.6220, -1.4707],
           [-1.2112,  0.7024]], device='xla:0')
   

   L'output indica che la funzione è stata chiamata correttamente su ogni dispositivo XLA (8 dispositivi in questo esempio) nella sezione TPU multi-host.

Modalità incentrata sull'host (JAX)

Le sezioni seguenti descrivono la modalità incentrata sull'host con JAX. JAX utilizza un paradigma di programmazione funzionale e supporta la semantica di un singolo programma a più dati (SPMD) di livello superiore. Invece di far interagire ogni processo con un singolo dispositivo XLA, il codice JAX è progettato per funzionare su più dispositivi su un singolo host contemporaneamente.

JAX è progettato per il computing ad alte prestazioni e può utilizzare in modo efficiente le TPU per l'addestramento e l'inferenza su larga scala. Questa modalità è ideale se hai familiarità con i concetti di programmazione funzionale, in modo da poter sfruttare tutto il potenziale di JAX.

Queste istruzioni presuppongono che tu abbia già configurato un ambiente Ray e TPU, incluso un ambiente software che include JAX e altri pacchetti correlati. Per creare un cluster Ray TPU, segui le istruzioni riportate in Avvia Google Cloud cluster GKE con TPU per KubeRay. Per saperne di più sull'utilizzo delle TPU con KubeRay, consulta Utilizzare le TPU con KubeRay.

Esegui un carico di lavoro JAX su una TPU a host singolo

Nota:questa sezione si applica all'API Cloud TPU. Per Ironwood (TPU7x), devi utilizzare Google Kubernetes Engine (GKE). Per maggiori informazioni, consulta Informazioni sulle TPU in GKE.

Il seguente script di esempio mostra come eseguire una funzione JAX su un cluster Ray con una TPU a host singolo, ad esempio v6e-4. Se hai una TPU multi-host, questo script smette di rispondere a causa del modello di esecuzione multi-controller di JAX. Per ulteriori informazioni sull'esecuzione di Ray su una TPU multi-host, consulta Eseguire un carico di lavoro JAX su una TPU multi-host.

1. Crea variabili di ambiente per i parametri di creazione della TPU.

   export PROJECT_ID=your-project-id
   export TPU_NAME=your-tpu-name
   export ZONE=europe-west4-a
   export ACCELERATOR_TYPE=v6e-4
   export RUNTIME_VERSION=v2-alpha-tpuv6e

   Descrizioni delle variabili di ambiente

   Variabile	Descrizione
   PROJECT_ID	L'ID progetto Google Cloud . Utilizza un progetto esistente o creane uno nuovo.
   TPU_NAME	Il nome della TPU.
   ZONE	La zona in cui creare la VM TPU. Per saperne di più sulle zone supportate, consulta Regioni e zone TPU.
   ACCELERATOR_TYPE	Il tipo di acceleratore specifica la versione e le dimensioni della Cloud TPU che vuoi creare. Per saperne di più sui tipi di acceleratore supportati per ogni versione di TPU, consulta la sezione Versioni di TPU.
   RUNTIME_VERSION	La versione software di Cloud TPU.

2. Utilizza il seguente comando per creare una VM TPU v6e con 4 core:

   gcloud compute tpus tpu-vm create $TPU_NAME \
      --zone=$ZONE \
      --accelerator-type=$ACCELERATOR_TYPE  \
      --version=$RUNTIME_VERSION

3. Connettiti alla VM TPU utilizzando il seguente comando:

   gcloud compute tpus tpu-vm ssh $TPU_NAME --zone=$ZONE

4. Installa JAX e Ray sulla TPU.

   pip install ray jax[tpu] -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
   

5. Salva il seguente codice in un file. Ad esempio, ray-jax-single-host.py.

   import ray
   import jax
   
   @ray.remote(resources={"TPU": 4})
   def my_function() -> int:
       return jax.device_count()
   
   h = my_function.remote()
   print(ray.get(h)) # => 4
   

   Se hai l'abitudine di eseguire Ray con le GPU, ci sono alcune differenze fondamentali quando utilizzi le TPU:

   • Anziché impostare num_gpus, specifica TPU come risorsa personalizzata e imposta il numero di chip TPU.
   • Specifica la TPU utilizzando il numero di chip per nodo worker Ray. Ad esempio, se utilizzi una v6e-4, l'esecuzione di una funzione remota con TPU impostato su 4 consuma l'intero host TPU.
   • Questo è diverso dal modo in cui vengono eseguite in genere le GPU, con un processo per host. L'impostazione di TPU su un numero diverso da 4 non è consigliata.
     • Eccezione: se hai un singolo host v6e-8 o v5litepod-8, devi impostare questo valore su 8.

6. Esegui lo script.

   python ray-jax-single-host.py

Esegui un workload JAX su una TPU multi-host

Nota:questa sezione si applica all'API Cloud TPU. Per Ironwood (TPU7x), devi utilizzare Google Kubernetes Engine (GKE). Per maggiori informazioni, consulta Informazioni sulle TPU in GKE.

Il seguente script di esempio mostra come eseguire una funzione JAX su un cluster Ray con una TPU multi-host. Lo script di esempio utilizza un v6e-16.

1. Crea variabili di ambiente per i parametri di creazione della TPU.

   export PROJECT_ID=your-project-id
   export TPU_NAME=your-tpu-name
   export ZONE=europe-west4-a
   export ACCELERATOR_TYPE=v6e-16
   export RUNTIME_VERSION=v2-alpha-tpuv6e

   Descrizioni delle variabili di ambiente

   Variabile	Descrizione
   PROJECT_ID	L'ID progetto Google Cloud . Utilizza un progetto esistente o creane uno nuovo.
   TPU_NAME	Il nome della TPU.
   ZONE	La zona in cui creare la VM TPU. Per saperne di più sulle zone supportate, consulta Regioni e zone TPU.
   ACCELERATOR_TYPE	Il tipo di acceleratore specifica la versione e le dimensioni della Cloud TPU che vuoi creare. Per saperne di più sui tipi di acceleratore supportati per ogni versione di TPU, consulta la sezione Versioni di TPU.
   RUNTIME_VERSION	La versione software di Cloud TPU.

2. Utilizza il seguente comando per creare una VM TPU v6e con 16 core:

   gcloud compute tpus tpu-vm create $TPU_NAME \
      --zone=$ZONE \
      --accelerator-type=$ACCELERATOR_TYPE  \
      --version=$RUNTIME_VERSION

3. Installa JAX e Ray su tutti i worker TPU.

   gcloud compute tpus tpu-vm ssh $TPU_NAME \
      --zone=$ZONE \
      --worker=all \
      --command="pip install ray jax[tpu] -f https://storage.googleapis.com/jax-releases/libtpu_releases.html"

4. Salva il seguente codice in un file. Ad esempio, ray-jax-multi-host.py.

   import ray
   import jax
   
   @ray.remote(resources={"TPU": 4})
   def my_function() -> int:
       return jax.device_count()
   
   ray.init()
   num_tpus = ray.available_resources()["TPU"]
   num_hosts = int(num_tpus) # 4
   h = [my_function.remote() for _ in range(num_hosts)]
   print(ray.get(h)) # [16, 16, 16, 16]
   

   Se hai l'abitudine di eseguire Ray con le GPU, ci sono alcune differenze fondamentali quando utilizzi le TPU:

   • Simile ai carichi di lavoro PyTorch sulle GPU:
     • I carichi di lavoro JAX sulle TPU vengono eseguiti in modalità multi-controller, single program, multiple data (SPMD).
     • I collettivi tra dispositivi vengono gestiti dal framework di machine learning.
   • A differenza dei carichi di lavoro PyTorch sulle GPU, JAX ha una visione globale dei dispositivi disponibili nel cluster.

5. Copia lo script in tutti i worker TPU.

   gcloud compute tpus tpu-vm scp ray-jax-multi-host.py $TPU_NAME: --zone=$ZONE --worker=all

6. Esegui lo script.

   gcloud compute tpus tpu-vm ssh $TPU_NAME \
      --zone=$ZONE \
      --worker=all \
      --command="python ray-jax-multi-host.py"

Esegui un carico di lavoro JAX multislice

Multislice ti consente di eseguire carichi di lavoro che si estendono su più slice TPU all'interno di un singolopod di TPUU o in più pod sulla rete del data center.

Puoi utilizzare il pacchetto ray-tpu per semplificare le interazioni di Ray con gli slice TPU.

Installa ray-tpu utilizzando pip.

pip install ray-tpu

Per maggiori informazioni sull'utilizzo del pacchetto ray-tpu, consulta la sezione Guida introduttiva nel repository GitHub. Per un esempio di utilizzo di Multislice, consulta Esecuzione su Multislice.

Orchestra i carichi di lavoro utilizzando Ray e MaxText

Per saperne di più sull'utilizzo di Ray con MaxText, consulta Esegui un job di addestramento con MaxText.

Risorse TPU e Ray

Ray tratta le TPU in modo diverso dalle GPU per tenere conto della differenza di utilizzo. Nell'esempio seguente, ci sono nove nodi Ray in totale:

• Il nodo head di Ray è in esecuzione su una VM n1-standard-16.
• I nodi worker Ray vengono eseguiti su due TPU v6e-16. Ogni TPU è costituita da quattro worker.

$ ray status
======== Autoscaler status: 2024-10-17 09:30:00.854415 ========
Node status
---------------------------------------------------------------
Active:
 1 node_e54a65b81456cee40fcab16ce7b96f85406637eeb314517d9572dab2
 1 node_9a8931136f8d2ab905b07d23375768f41f27cc42f348e9f228dcb1a2
 1 node_c865cf8c0f7d03d4d6cae12781c68a840e113c6c9b8e26daeac23d63
 1 node_435b1f8f1fbcd6a4649c09690915b692a5bac468598e9049a2fac9f1
 1 node_3ed19176e9ecc2ac240c818eeb3bd4888fbc0812afebabd2d32f0a91
 1 node_6a88fe1b74f252a332b08da229781c3c62d8bf00a5ec2b90c0d9b867
 1 node_5ead13d0d60befd3a7081ef8b03ca0920834e5c25c376822b6307393
 1 node_b93cb79c06943c1beb155d421bbd895e161ba13bccf32128a9be901a
 1 node_9072795b8604ead901c5268ffcc8cc8602c662116ac0a0272a7c4e04
Pending:
 (no pending nodes)
Recent failures:
 (no failures)

Resources
---------------------------------------------------------------
Usage:
 0.0/727.0 CPU
 0.0/32.0 TPU
 0.0/2.0 TPU-v6e-16-head
 0B/5.13TiB memory
 0B/1.47TiB object_store_memory
 0.0/4.0 tpu-group-0
 0.0/4.0 tpu-group-1

Demands:
 (no resource demands)

Descrizioni dei campi di utilizzo delle risorse:

• CPU: il numero totale di CPU disponibili nel cluster.
• TPU: il numero di chip TPU nel cluster.
• TPU-v6e-16-head: un identificatore speciale per la risorsa che corrisponde al worker 0 di una sezione TPU. Questo è importante per accedere alle singole sezioni TPU.
• memory: Memoria heap worker utilizzata dall'applicazione.
• object_store_memory: memoria utilizzata quando l'applicazione crea oggetti nell'archivio di oggetti utilizzando ray.put e quando restituisce valori da funzioni remote.
• tpu-group-0 e tpu-group-1: identificatori univoci per le singole sezioni della TPU. Questo è importante per l'esecuzione dei job sulle sezioni. Questi campi sono impostati su 4 perché ci sono 4 host per slice TPU in una v6e-16.

Nota: tutte queste risorse vengono rilevate automaticamente all'avvio del runtime Ray (ray start), quindi non è necessario impostarle in modo esplicito.