Effectuer le scaling des charges de travail de ML avec Ray

Ce document explique comment exécuter des charges de travail de machine learning (ML) avec Ray et JAX sur des TPU. Il existe deux modes d'utilisation des TPU avec Ray : le mode axé sur l'appareil (PyTorch/XLA) et le mode axé sur l'hôte (JAX).

Dans ce document, nous partons du principe que vous avez déjà configuré un environnement TPU. Pour en savoir plus, consultez les ressources suivantes :

• Cloud TPU : Configurer l'environnement Cloud TPU et Gérer les ressources TPU
• Google Kubernetes Engine (GKE) : Déployer des charges de travail TPU dans GKE Autopilot ou Déployer des charges de travail TPU dans GKE Standard

Mode axé sur l'appareil (PyTorch/XLA)

Le mode centré sur l'appareil conserve l'essentiel du style de programmation de PyTorch classique. Dans ce mode, vous ajoutez un nouveau type d'appareil XLA qui fonctionne comme n'importe quel autre appareil PyTorch. Chaque processus interagit avec un appareil XLA.

Ce mode est idéal si vous maîtrisez déjà PyTorch avec les GPU et que vous souhaitez utiliser des abstractions de codage similaires.

Les sections suivantes décrivent comment exécuter une charge de travail PyTorch/XLA sur un ou plusieurs appareils sans utiliser Ray, puis comment exécuter la même charge de travail sur plusieurs hôtes avec Ray.

Créer un TPU

1. Créez des variables d'environnement pour les paramètres de création de 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

   Descriptions des variables d'environnement

   Variable	Description
   PROJECT_ID	ID de votre projet Google Cloud . Utilisez un projet existant ou créez-en un.
   TPU_NAME	Nom du TPU.
   ZONE	Zone dans laquelle créer la VM TPU. Pour en savoir plus sur les zones compatibles, consultez Régions et zones des TPU.
   ACCELERATOR_TYPE	Le type d'accélérateur spécifie la version et la taille du Cloud TPU que vous souhaitez créer. Pour en savoir plus sur les types d'accélérateurs compatibles avec chaque version de TPU, consultez Versions de TPU.
   RUNTIME_VERSION	Version logicielle de Cloud TPU.

2. Exécutez la commande suivante pour créer une VM TPU v5p à huit cœurs :

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

3. Connectez-vous à la VM TPU à l'aide de la commande suivante :

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

Si vous utilisez GKE, consultez le guide KubeRay sur GKE pour obtenir des informations sur la configuration.

Installer les éléments requis

Exécutez les commandes suivantes sur votre VM TPU pour installer les dépendances requises :

1. Enregistrez le code suivant dans un fichier que vous nommerez par exemple 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. Pour installer les dépendances requises, exécutez la commande suivante :

   pip install -r requirements.txt
   

Si vous exécutez votre charge de travail sur GKE, nous vous recommandons de créer un fichier Dockerfile qui installe les dépendances requises. Consultez Exécuter votre charge de travail sur des nœuds de tranche de TPU dans la documentation GKE pour obtenir un exemple.

Exécuter une charge de travail PyTorch/XLA sur un seul appareil

L'exemple suivant montre comment créer un tenseur XLA sur un seul appareil, à savoir une puce TPU. Cette méthode se rapproche de la façon dont PyTorch traite les autres types d'appareils.

1. Enregistrez l'extrait de code suivant dans un fichier que vous nommerez par exemple 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'instruction d'importation import torch_xla initialise PyTorch/XLA, et la fonction xm.xla_device() affiche l'appareil XLA actuel, à savoir une puce TPU.

2. Définissez la variable d'environnement PJRT_DEVICE sur TPU.

   export PJRT_DEVICE=TPU
   

3. Exécutez le script.

   python workload.py
   

   Le résultat doit ressembler à ce qui suit. Assurez-vous qu'il indique bien que l'appareil XLA a été trouvé.

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

Exécuter PyTorch/XLA sur plusieurs appareils

1. Modifiez l'extrait de code de la section précédente pour qu'il s'exécute sur plusieurs appareils.

   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. Exécutez le script.

   python workload.py
   

   Si vous exécutez l'extrait de code sur un TPU v5p-8, le résultat ressemble à ceci :

   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() accepte deux arguments : une fonction et une liste de paramètres. Elle crée un processus pour chaque appareil XLA disponible et appelle la fonction spécifiée dans les arguments. Comme il y a quatre appareils TPU disponibles dans cet exemple, torch_xla.launch() crée quatre processus et appelle _mp_fn() sur chaque appareil. Chaque processus n'a accès qu'à un seul appareil. Chaque appareil possède donc l'index 0, et xla:0 est imprimé pour tous les processus.

Exécuter PyTorch/XLA sur plusieurs hôtes avec Ray

Les sections suivantes montrent comment exécuter le même extrait de code sur une tranche de TPU multihôte plus grande. Pour en savoir plus sur l'architecture TPU multi-hôtes, consultez Architecture du système.

Dans cet exemple, vous allez configurer Ray manuellement. Si vous savez déjà comment configurer Ray, vous pouvez passer directement à la dernière section, Exécuter une charge de travail Ray. Pour en savoir plus sur la configuration de Ray pour un environnement de production, consultez les ressources suivantes :

• Lancer des clusters Ray
• Démarrer un cluster GKE Google Cloud avec des TPU pour KubeRay

Créer une VM TPU multi-hôtes

1. Créez des variables d'environnement pour les paramètres de création de 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

   Descriptions des variables d'environnement

   Variable	Description
   PROJECT_ID	ID de votre projet Google Cloud . Utilisez un projet existant ou créez-en un.
   TPU_NAME	Nom du TPU.
   ZONE	Zone dans laquelle créer la VM TPU. Pour en savoir plus sur les zones compatibles, consultez Régions et zones des TPU.
   ACCELERATOR_TYPE	Le type d'accélérateur spécifie la version et la taille du Cloud TPU que vous souhaitez créer. Pour en savoir plus sur les types d'accélérateurs compatibles avec chaque version de TPU, consultez Versions de TPU.
   RUNTIME_VERSION	Version logicielle de Cloud TPU.

2. Créez un TPU v5p multi-hôte avec deux hôtes (un v5p-16 avec quatre puces TPU sur chaque hôte) avec la commande suivante :

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

Configurer Ray

Un TPU v5p-16 comporte deux hôtes TPU, chacun contenant quatre puces TPU. Dans cet exemple, vous allez démarrer le nœud principal Ray sur un hôte, puis ajouter le deuxième hôte au cluster Ray en tant que nœud de calcul.

1. Connectez-vous au premier hôte via SSH.

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

2. Installez les dépendances en utilisant le même fichier d'exigences que celui de la section Installer les éléments requis.

   pip install -r requirements.txt
   

3. Démarrez le processus Ray.

   ray start --head --port=6379
   

   Remarque : Si vous obtenez une erreur de type Command 'ray' not found, vous devrez peut-être ajouter le répertoire contenant l'exécutable ray à la variable d'environnement PATH de votre système. Vous pouvez aussi spécifier le chemin d'accès complet lorsque vous exécutez des commandes ray (par exemple, ~/.local/bin/ray).
   Remarque : La commande ray start vous invite à activer ou désactiver la collecte des statistiques d'utilisation. Pour en savoir plus, consultez Collecte des statistiques d'utilisation dans la documentation Ray.

   Le résultat ressemble à ceci :

   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
   

   Cet hôte TPU est désormais le nœud principal Ray. Notez les lignes qui indiquent comment ajouter un autre nœud au cluster Ray, par exepmle :

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

   Vous utiliserez cette commande dans une prochaine étape.

4. Vérifiez l'état du cluster Ray :

   ray status
   

   Le résultat ressemble à ceci :

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

   Le cluster ne contient que quatre TPU (0.0/4.0 TPU), car seul le nœud principal a été ajouté jusqu'à présent.

   Maintenant que le nœud principal est en cours d'exécution, vous pouvez ajouter le deuxième hôte au cluster.

5. Connectez-vous au deuxième hôte via SSH.

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

6. Installez les dépendances en utilisant le même fichier d'exigences que celui de la section Installer les éléments requis.

   pip install -r requirements.txt
   

7. Démarrez le processus Ray. Pour ajouter ce nœud au cluster Ray existant, utilisez la commande issue du résultat de la commande ray start. Veillez à remplacer l'adresse IP et le port dans la commande suivante :

   ray start --address='10.130.0.76:6379'

   Le résultat ressemble à ceci :

   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. Vérifiez à nouveau l'état de Ray :

   ray status
   

   Le résultat ressemble à ceci :

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

   Le deuxième hôte TPU est désormais un nœud du cluster. La liste des ressources disponibles affiche désormais huit TPU (0.0/8.0 TPU).

   Remarque : Vous devez démarrer le processus Ray sur chacune de vos VM TPU avant d'exécuter la charge de travail.

Exécuter une charge de travail Ray

1. Modifiez l'extrait de code pour qu'il s'exécute sur le 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. Exécutez le script sur le nœud principal Ray. Remplacez ray-workload.py par le chemin d'accès à votre script.

   python ray-workload.py

   Le résultat ressemble à ceci :

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

   Le résultat indique que la fonction a été appelée avec succès sur chaque appareil XLA (huit appareils dans cet exemple) de la tranche TPU multi-hôte.

Mode centré sur l'hôte (JAX)

Les sections suivantes décrivent le mode centré sur l'hôte avec JAX. JAX utilise un paradigme de programmation fonctionnelle et prend en charge la sémantique SPMD (Single Program, Multiple Data) de niveau supérieur. Au lieu de faire interagir chaque processus avec un seul appareil XLA, le code JAX est conçu pour fonctionner simultanément sur plusieurs appareils d'un même hôte.

JAX est conçu pour le calcul hautes performances et peut utiliser efficacement les TPU pour l'entraînement et l'inférence à grande échelle. Si vous connaissez déjà les concepts de la programmation fonctionnelle, ce mode vous permettra de tirer pleinement parti du potentiel de JAX.

Ces instructions supposent que vous avez déjà configuré un environnement Ray et TPU, y compris un environnement logiciel qui inclut JAX et d'autres packages associés. Pour créer un cluster Ray TPU, suivez les instructions de la section Démarrer un cluster GKE Google Cloud avec des TPU pour KubeRay. Pour en savoir plus sur l'utilisation des TPU avec KubeRay, consultez Utiliser des TPU avec KubeRay.

Exécuter une charge de travail JAX sur un TPU à hôte unique

L'exemple de script suivant montre comment exécuter une fonction JAX sur un cluster Ray avec un TPU à hôte unique, tel qu'un v6e-4. Si votre TPU est multi-hôte, le modèle d'exécution multi-contrôleur de JAX empêchera ce script de répondre. Pour en savoir plus sur l'exécution de Ray sur un TPU multi-hôte, consultez Exécuter une charge de travail JAX sur un TPU multi-hôte.

1. Créez des variables d'environnement pour les paramètres de création de 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

   Descriptions des variables d'environnement

   Variable	Description
   PROJECT_ID	ID de votre projet Google Cloud . Utilisez un projet existant ou créez-en un.
   TPU_NAME	Nom du TPU.
   ZONE	Zone dans laquelle créer la VM TPU. Pour en savoir plus sur les zones compatibles, consultez Régions et zones des TPU.
   ACCELERATOR_TYPE	Le type d'accélérateur spécifie la version et la taille du Cloud TPU que vous souhaitez créer. Pour en savoir plus sur les types d'accélérateurs compatibles avec chaque version de TPU, consultez Versions de TPU.
   RUNTIME_VERSION	Version logicielle de Cloud TPU.

2. Créez une VM TPU v6e avec quatre cœurs en utilisant la commande suivante :

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

3. Connectez-vous à la VM TPU à l'aide de la commande suivante :

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

4. Installez JAX et Ray sur votre TPU.

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

5. Enregistrez le code suivant dans un fichier que vous nommerez par exemple 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
   

   Si vous avez l'habitude d'exécuter Ray avec des GPU, sachez que l'utilisation de TPU présente des différences fondamentales :

   • Au lieu de définir num_gpus, spécifiez TPU comme ressource personnalisée et définissez le nombre de puces TPU.
   • Spécifiez le TPU en indiquant le nombre de puces par nœud de calcul Ray. Par exemple, si vous utilisez un TPU v6e-4, l'exécution d'une fonction distante avec TPU défini sur 4 consomme l'intégralité de l'hôte TPU.
   • L'exécution des GPU est généralement différente, avec un processus par hôte. Évitez de définir TPU sur un nombre différent de 4.
     • Exception : Si vous disposez d'un v6e-8 ou v5litepod-8 à hôte unique, vous devez définir cette valeur sur 8.

6. Exécutez le script.

   python ray-jax-single-host.py

Exécuter une charge de travail JAX sur un TPU multi-hôte

L'exemple de script suivant montre comment exécuter une fonction JAX sur un cluster Ray avec un TPU multi-hôte. L'exemple de script utilise un v6e-16.

1. Créez des variables d'environnement pour les paramètres de création de 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

   Descriptions des variables d'environnement

   Variable	Description
   PROJECT_ID	ID de votre projet Google Cloud . Utilisez un projet existant ou créez-en un.
   TPU_NAME	Nom du TPU.
   ZONE	Zone dans laquelle créer la VM TPU. Pour en savoir plus sur les zones compatibles, consultez Régions et zones des TPU.
   ACCELERATOR_TYPE	Le type d'accélérateur spécifie la version et la taille du Cloud TPU que vous souhaitez créer. Pour en savoir plus sur les types d'accélérateurs compatibles avec chaque version de TPU, consultez Versions de TPU.
   RUNTIME_VERSION	Version logicielle de Cloud TPU.

2. Créez une VM TPU v6e avec 16 cœurs en utilisant la commande suivante :

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

3. Installez JAX et Ray sur tous les travailleurs 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. Enregistrez le code suivant dans un fichier que vous nommerez par exemple 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]
   

   Si vous avez l'habitude d'exécuter Ray avec des GPU, sachez que l'utilisation de TPU présente des différences fondamentales :

   • Comme pour les charges de travail PyTorch sur les GPU :
     • Les charges de travail JAX sur les TPU s'exécutent en mode multi-contrôleur, programme unique, données multiples (SPMD).
     • Les collectifs entre appareils sont gérés par le framework de machine learning.
   • Contrairement aux charges de travail PyTorch sur les GPU, JAX offre une vue globale des appareils disponibles dans le cluster.

5. Copiez le script sur tous les travailleurs TPU.

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

6. Exécutez le script.

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

Exécuter une charge de travail JAX Multislice

Multislice vous permet d'exécuter des charges de travail couvrant plusieurs tranches de TPU dans un seul pod TPU ou dans plusieurs pods sur le réseau du centre de données.

Pour simplifier les interactions de Ray avec les tranches TPU, vous pouvez utiliser le package ray-tpu.

Installez ray-tpu à l'aide de pip.

pip install ray-tpu

Pour en savoir plus sur l'utilisation du package ray-tpu, consultez Premiers pas dans le dépôt GitHub. Pour un exemple d'utilisation de Multislice, consultez Exécuter sur Multislice.

Orchestrer des charges de travail à l'aide de Ray et MaxText

Pour en savoir plus sur l'utilisation de Ray avec MaxText, consultez Exécuter un job d'entraînement avec MaxText.

Ressources TPU et Ray

Ray traite les TPU différemment des GPU pour refléter les différences d'utilisation entre ces deux circuits. Dans l'exemple suivant, il y a un total de neuf nœuds Ray :

• Le nœud principal Ray s'exécute sur une VM n1-standard-16.
• Les nœuds de calcul Ray s'exécutent sur deux TPU v6e-16. Chaque TPU forme quatre nœuds de calcul.

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

Descriptions des champs d'utilisation des ressources :

• CPU : nombre total de processeurs disponibles dans le cluster.
• TPU : nombre de puces TPU dans le cluster.
• TPU-v6e-16-head : identifiant spécial de la ressource qui correspond au nœud de calcul 0 d'une tranche TPU. Cette étape est importante pour accéder aux tranches de TPU individuelles.
• memory : tas de mémoire du nœud de calcul utilisé par votre application.
• object_store_memory : mémoire utilisée lorsque votre application crée des objets dans le store d'objets à l'aide de ray.put et lorsqu'elle renvoie des valeurs à partir de fonctions à distance.
• tpu-group-0 et tpu-group-1 : identifiants uniques pour les tranches de TPU individuelles. Ils sont importants pour exécuter des jobs sur des tranches. Ces champs sont définis sur 4, car il y a quatre hôtes par tranche de TPU dans un v6e-16.

Remarque : Toutes ces ressources sont détectées automatiquement lorsque l'environnement d'exécution Ray est démarré (ray start). Vous n'avez donc pas besoin de les définir explicitement.