Escala cargas de trabajo de AA con Ray

En este documento, se brindan detalles para ejecutar cargas de trabajo de aprendizaje automático (AA) con Ray y JAX en TPU. Existen dos modos diferentes para usar las TPUs con Ray: modo centrado en el dispositivo (PyTorch/XLA) y modo centrado en el host (JAX).

En este documento, se supone que ya tienes configurado un entorno de TPU. Para obtener más información, consulta los siguientes recursos:

• Cloud TPU: Configura el entorno de Cloud TPU y Administra los recursos TPU.
• Google Kubernetes Engine (GKE): Implementa cargas de trabajo de TPU en GKE Autopilot o Implementa cargas de trabajo de TPU en GKE Standard.

Modo centrado en el dispositivo (PyTorch/XLA)

Este modo retiene gran parte del estilo programático de PyTorch clásico. Aquí, agregas un nuevo tipo de dispositivo XLA, que funciona como cualquier otro dispositivo de PyTorch. Cada proceso individual interactúa con un dispositivo XLA.

Este modo es ideal si ya conoces PyTorch con GPU y quieres usar abstracciones de programación similares.

En las siguientes secciones, se describe cómo ejecutar una carga de trabajo de PyTorch/XLA en uno o más dispositivos sin usar Ray, además de cómo ejecutar la misma carga de trabajo en varios hosts usando Ray.

Crea una TPU

1. Crea variables de entorno para los parámetros de creación de la 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

   Descripciones de las variables de entorno

   Variable	Descripción
   PROJECT_ID	Es el ID de tu proyecto de Google Cloud . Usa un proyecto existente o crea uno nuevo.
   TPU_NAME	Es el nombre de la TPU.
   ZONE	Es la zona en la que se creará la VM de TPU. Para obtener más información sobre las zonas compatibles, consulta Regiones y zonas de TPU.
   ACCELERATOR_TYPE	El tipo de acelerador especifica la versión y el tamaño de Cloud TPU que deseas crear. Para obtener más información sobre los tipos de aceleradores compatibles con cada versión de TPU, consulta Versiones de TPU.
   RUNTIME_VERSION	Es la versión de software de Cloud TPU.

2. Usa el siguiente comando para crear una VM de la TPU v5p con 8 núcleos:

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

3. Conéctate a la VM de la TPU con el siguiente comando:

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

Si usas GKE, consulta la guía de KubeRay en GKE para obtener información sobre la configuración.

Requisitos de instalación

Ejecuta los siguientes comandos en la VM de tu TPU para instalar las dependencias requeridas:

1. Guarda lo siguiente en un archivo. Por ejemplo, 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. Para instalar las dependencias requeridas, ejecuta el siguiente comando:

   pip install -r requirements.txt
   

Si ejecutas tu carga de trabajo en GKE, te recomendamos que crees un Dockerfile que instale las dependencias requeridas. Para ver un ejemplo, consulta Ejecuta tu carga de trabajo en nodos en porción de TPU en la documentación de GKE.

Ejecuta una carga de trabajo de PyTorch/XLA en un solo dispositivo

En el siguiente ejemplo, se muestra cómo crear un tensor de XLA en un solo dispositivo, que es un chip TPU. Esto es similar a la forma en que PyTorch controla otros tipos de dispositivos.

1. Guarda el siguiente fragmento de código en un archivo. Por ejemplo, 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)
   

   La sentencia de importación import torch_xla inicializa PyTorch/XLA, y la función xm.xla_device() devuelve el dispositivo XLA actual, que es un chip TPU.

2. Establece la variable de entorno PJRT_DEVICE en TPU.

   export PJRT_DEVICE=TPU
   

3. Ejecuta la secuencia de comandos.

   python workload.py
   

   El resultado es similar al siguiente. Asegúrate de que el resultado indique que se encontró el dispositivo XLA.

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

Ejecuta PyTorch/XLA en varios dispositivos

1. Actualiza el fragmento de código de la sección anterior para que se ejecute en varios dispositivos.

   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. Ejecuta la secuencia de comandos.

   python workload.py
   

   Si ejecutas el fragmento de código en una TPU v5p-8, el resultado será similar al siguiente:

   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() toma dos argumentos: una función y una lista de parámetros. Crea un proceso para cada dispositivo XLA disponible y llama a la función especificada en los argumentos. En este ejemplo, hay 4 dispositivos de TPU disponibles, por lo que torch_xla.launch() crea 4 procesos y llama a _mp_fn() en cada dispositivo. Cada proceso tiene acceso a un solo dispositivo, por lo que cada dispositivo tiene el índice 0 y se imprime xla:0 para todos los procesos.

Ejecuta PyTorch/XLA en varios hosts con Ray

En las siguientes secciones, se muestra cómo ejecutar el mismo fragmento de código en una porción de TPU de varios hosts más grande. Para obtener más información sobre la arquitectura de TPU de varios hosts, consulta Arquitectura del sistema.

En este ejemplo, vas a configurar Ray de forma manual. Si ya sabes cómo configurar Ray, puedes pasar a la última sección, Ejecuta una carga de trabajo de Ray. Si deseas obtener más información sobre cómo configurar Ray para un entorno de producción, consulta los siguientes recursos:

• Inicia clústeres de Ray.
• Inicia un clúster de GKE de Google Cloud con TPU para KubeRay.

Crea una VM de TPU de varios hosts

1. Crea variables de entorno para los parámetros de creación de la 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

   Descripciones de las variables de entorno

   Variable	Descripción
   PROJECT_ID	Es el ID de tu proyecto de Google Cloud . Usa un proyecto existente o crea uno nuevo.
   TPU_NAME	Es el nombre de la TPU.
   ZONE	Es la zona en la que se creará la VM de TPU. Para obtener más información sobre las zonas compatibles, consulta Regiones y zonas de TPU.
   ACCELERATOR_TYPE	El tipo de acelerador especifica la versión y el tamaño de Cloud TPU que deseas crear. Para obtener más información sobre los tipos de aceleradores compatibles con cada versión de TPU, consulta Versiones de TPU.
   RUNTIME_VERSION	Es la versión de software de Cloud TPU.

2. Crea una TPU v5p de varios hosts con 2 hosts (una v5p-16, con 4 chips TPU en cada host) con el siguiente comando:

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

Configura Ray

Una TPU v5p-16 tiene 2 hosts de TPU, cada uno con 4 chips TPU. En este ejemplo, iniciarás el nodo principal de Ray en un host y agregarás el segundo como un nodo trabajador al clúster de Ray.

1. Conéctate al primer host con SSH.

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

2. Instala las dependencias con el mismo archivo de requisitos que en la sección Requisitos de instalación.

   pip install -r requirements.txt
   

3. Inicia el proceso de Ray.

   ray start --head --port=6379
   

   Nota: Si recibes un error como Command 'ray' not found, es posible que debas agregar el directorio que contiene el ejecutable ray a la variable de entorno PATH de tu sistema. También puedes especificar la ruta de acceso completa cuando ejecutes comandos ray (por ejemplo, ~/.local/bin/ray).
   Nota: El comando ray start te solicita que habilites o inhabilites la recopilación de estadísticas de uso. Para obtener más información, consulta Recopilación de estadísticas de uso en la documentación de Ray.

   El resultado es similar al siguiente:

   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
   

   Ahora, este host de TPU es el nodo principal de Ray. Toma nota de las líneas que muestran cómo agregar otro nodo al clúster de Ray, de manera similar a la siguiente:

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

   Usarás este comando en otro paso más adelante.

4. Verifica el estado del clúster de Ray.

   ray status
   

   El resultado es similar al siguiente:

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

   El clúster solo contiene 4 TPU (0.0/4.0 TPU) porque solo agregaste el nodo principal hasta el momento.

   Ahora que el nodo principal está en ejecución, puedes agregar el segundo host al clúster.

5. Conéctate al segundo host con SSH.

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

6. Instala las dependencias con el mismo archivo de requisitos que en la sección Requisitos de instalación.

   pip install -r requirements.txt
   

7. Inicia el proceso de Ray. Para agregar este nodo al clúster de Ray existente, usa el comando desde el resultado del comando ray start. Asegúrate de reemplazar el puerto y la dirección IP en el siguiente comando:

   ray start --address='10.130.0.76:6379'

   El resultado es similar al siguiente:

   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. Vuelve a verificar el estado de Ray:

   ray status
   

   El resultado es similar al siguiente:

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

   El segundo host de TPU ahora es un nodo del clúster. En la lista de recursos disponibles, ahora se muestran 8 TPU (0.0/8.0 TPU).

   Nota: Debes iniciar el proceso de Ray en cada una de las VMs de tu TPU antes de ejecutar la carga de trabajo.

Ejecuta una carga de trabajo de Ray

1. Actualiza el fragmento de código para que se ejecute en el clúster de 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. Ejecuta la secuencia de comandos en el nodo principal de Ray. Reemplaza ray-workload.py por la ruta de acceso a tu secuencia de comandos.

   python ray-workload.py

   El resultado es similar al siguiente:

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

   El resultado indica que la función se llamó de forma correcta en cada dispositivo XLA (8 dispositivos en este ejemplo) en la porción de TPU de varios hosts.

Modo centrado en el host (JAX)

En las siguientes secciones, se describe el modo centrado en el host con JAX. JAX usa un paradigma de programación funcional y admite semánticas de datos múltiples y programa único (SPMD) de nivel superior. En lugar de que cada proceso interactúe con un solo dispositivo XLA, el código de JAX está diseñado para operar en varios dispositivos en un solo host de forma simultánea.

JAX está diseñado para la computación de alto rendimiento y puede utilizar las TPU de manera eficiente para el entrenamiento y la inferencia a gran escala. Este modo es ideal si conoces los conceptos de programación funcional, ya que te permite aprovechar todo el potencial de JAX.

En estas instrucciones, se supone que ya tienes configurado un entorno de Ray y TPU, incluido uno de software que incluye JAX y otros paquetes relacionados. Para crear un clúster de TPU de Ray, sigue las instrucciones que se indican en Inicia un clúster de GKE de Google Cloud con TPU para KubeRay. Para obtener más información sobre el uso de las TPU con KubeRay, consulta Usa TPU con KubeRay.

Ejecuta una carga de trabajo de JAX en una TPU de host único

La siguiente secuencia de comandos de ejemplo muestra cómo ejecutar una función de JAX en un clúster de Ray con una TPU de host único, como v6e-4. Si tienes una TPU de varios hosts, esta secuencia de comandos dejará de responder debido al modelo de ejecución de varios controladores de JAX. Para obtener más información sobre cómo ejecutar Ray en una TPU de varios hosts, consulta Ejecuta una carga de trabajo de JAX en una TPU de varios hosts.

1. Crea variables de entorno para los parámetros de creación de la 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

   Descripciones de las variables de entorno

   Variable	Descripción
   PROJECT_ID	Es el ID de tu proyecto de Google Cloud . Usa un proyecto existente o crea uno nuevo.
   TPU_NAME	Es el nombre de la TPU.
   ZONE	Es la zona en la que se creará la VM de TPU. Para obtener más información sobre las zonas compatibles, consulta Regiones y zonas de TPU.
   ACCELERATOR_TYPE	El tipo de acelerador especifica la versión y el tamaño de Cloud TPU que deseas crear. Para obtener más información sobre los tipos de aceleradores compatibles con cada versión de TPU, consulta Versiones de TPU.
   RUNTIME_VERSION	Es la versión de software de Cloud TPU.

2. Usa el siguiente comando para crear una VM de TPU v6e con 4 núcleos:

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

3. Conéctate a la VM de la TPU con el siguiente comando:

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

4. Instala JAX y Ray en tu TPU.

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

5. Guarda el siguiente código en un archivo. Por ejemplo, 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 tienes la costumbre de ejecutar Ray con GPU, hay algunas diferencias clave cuando usas TPU:

   • En lugar de configurar num_gpus, especifica TPU como un recurso personalizado y establece la cantidad de chips TPU.
   • Especifica la TPU con la cantidad de chips por nodo trabajador de Ray. Por ejemplo, si usas un v6e-4, ejecutar una función remota con TPU establecido en 4 consume todo el host de la TPU.
   • Esto es diferente de cómo suelen ejecutarse las GPU, que realizan un proceso por host. No se recomienda establecer TPU en un número que no sea 4.
     • Excepción: Si tienes un v6e-8 o un v5litepod-8 de un solo host, debes establecer este valor en 8.

6. Ejecuta la secuencia de comandos.

   python ray-jax-single-host.py

Ejecuta una carga de trabajo de JAX en una TPU de varios hosts

En el siguiente ejemplo de secuencia de comandos, se muestra cómo ejecutar una función de JAX en un clúster de Ray con una TPU de varios hosts. La secuencia de comandos de ejemplo usa una v6e-16.

1. Crea variables de entorno para los parámetros de creación de la 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

   Descripciones de las variables de entorno

   Variable	Descripción
   PROJECT_ID	Es el ID de tu proyecto de Google Cloud . Usa un proyecto existente o crea uno nuevo.
   TPU_NAME	Es el nombre de la TPU.
   ZONE	Es la zona en la que se creará la VM de TPU. Para obtener más información sobre las zonas compatibles, consulta Regiones y zonas de TPU.
   ACCELERATOR_TYPE	El tipo de acelerador especifica la versión y el tamaño de Cloud TPU que deseas crear. Para obtener más información sobre los tipos de aceleradores compatibles con cada versión de TPU, consulta Versiones de TPU.
   RUNTIME_VERSION	Es la versión de software de Cloud TPU.

2. Usa el siguiente comando para crear una VM de la TPU v6e con 16 núcleos:

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

3. Instala JAX y Ray en todos los trabajadores 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. Guarda el siguiente código en un archivo. Por ejemplo, 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 tienes la costumbre de ejecutar Ray con GPU, hay algunas diferencias clave cuando usas TPU:

   • De manera similar a las cargas de trabajo de PyTorch en GPU:
     • Las cargas de trabajo de JAX en las TPU se ejecutan en un modo de datos múltiples y programa único con varios controladores.
     • El framework de aprendizaje automático controla las colecciones entre dispositivos.
   • A diferencia de las cargas de trabajo de PyTorch en GPU, JAX tiene una vista global de los dispositivos disponibles en el clúster.

5. Copia la secuencia de comandos en todos los trabajadores TPU.

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

6. Ejecuta la secuencia de comandos.

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

Ejecuta una carga de trabajo de JAX de Multislice

Multislice te permite ejecutar cargas de trabajo que abarcan varias porciones de TPU dentro de un solo pod de TPU o en varios Pods a través de la red del centro de datos.

Puedes usar el paquete ray-tpu para simplificar las interacciones de Ray con las porciones de TPU.

Instala ray-tpu con pip.

pip install ray-tpu

Para obtener más información sobre el uso del paquete ray-tpu, consulta Cómo comenzar en el repositorio de GitHub. Para ver un ejemplo del uso de Multislice, consulta Ejecuta en Multislice.

Organiza cargas de trabajo con Ray y MaxText

Para obtener más información sobre cómo usar Ray con MaxText, consulta Ejecuta un trabajo de entrenamiento con MaxText.

Recursos de TPU y Ray

Ray trata a las TPU de manera diferente a las GPU para adaptarse a la diferencia en el uso. En el siguiente ejemplo, hay nueve nodos de Ray en total:

• El nodo principal de Ray se ejecuta en una VM n1-standard-16.
• Los nodos trabajadores de Ray se ejecutan en dos TPU v6e-16. Cada TPU constituye cuatro trabajadores.

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

Descripciones de los campos de uso de recursos:

• CPU: la cantidad total de CPU disponibles en el clúster
• TPU: la cantidad de chips TPU en el clúster
• TPU-v6e-16-head: un identificador especial para el recurso que corresponde al trabajador 0 de una porción de TPU. Esto es importante para acceder a las porciones individuales de TPU
• memory: la memoria del montón del trabajador que usa tu aplicación
• object_store_memory: la memoria que se usa cuando tu aplicación crea objetos en el almacén de objetos con ray.put y cuando devuelve valores de funciones remotas
• tpu-group-0 y tpu-group-1: identificadores únicos para las porciones individuales de TPU. Esto es importante para ejecutar trabajos en porciones. Estos campos se establecen en 4 porque hay 4 hosts por porción de TPU en una v6e-16

Nota: Todos estos recursos se detectan de forma automática cuando se inicia el entorno de ejecución de Ray (ray start), por lo que no es necesario configurarlos de forma explícita.