Crea una porción de TPU de varios hosts

Crea una plantilla de instancias

Para crear una plantilla de instancias para las VMs de TPU v6e, usa el comando gcloud compute instance-templates create:

gcloud compute instance-templates create quickstart-tpu-instance-template \
    --machine-type=ct6e-standard-4t \
    --maintenance-policy=TERMINATE \
    --image-family=ubuntu-accel-2204-amd64-tpu-v5e-v5p-v6e \
    --image-project=ubuntu-os-accelerator-images \
    --region=us-east5

Crear una política de cargas de trabajo

Una política de carga de trabajo define las propiedades físicas de tus instancias de procesamiento. En las porciones de TPU, la topología del acelerador define la disposición física de los chips de TPU. Se requiere especificar una topología de acelerador para las particiones de TPU interconectadas de varios hosts.

Para crear una política de carga de trabajo para una porción de TPU de varios hosts, usa el comando gcloud compute resource-policies create workload-policy con la marca --accelerator-topology. El siguiente comando crea una política de cargas de trabajo con una topología de 2x4:

gcloud compute resource-policies create workload-policy quickstart-tpu-workload-policy \
    --type=high-throughput \
    --accelerator-topology=2x4 \
    --region=us-east5

Crear un MIG

Ejecuta los siguientes comandos para crear un MIG que forme una porción de TPU de varios hosts.

1. Para crear un MIG que forme una porción de TPU de varios hosts, usa el comando gcloud compute instance-groups managed create:

   gcloud compute instance-groups managed create quickstart-tpu-mig \
       --size=2 \
       --target-size-policy-mode=bulk \
       --template=quickstart-tpu-instance-template \
       --region=us-east5 \
       --target-distribution-shape=any-single-zone \
       --instance-redistribution-type=none \
       --default-action-on-vm-failure=do-nothing \
       --workload-policy=projects/PROJECT_ID/regions/us-east5/resourcePolicies/quickstart-tpu-workload-policy
   

   Reemplaza PROJECT_ID por el ID del proyecto de Cloud de Confiance by S3NS .

2. De manera opcional, verifica que las instancias administradas se estén ejecutando con los siguientes comandos:

   • Para ver el estado general del MIG, usa el comando gcloud compute instance-groups managed describe:

     gcloud compute instance-groups managed describe quickstart-tpu-mig \
         --region=us-east5
     

   • Para ver el estado de cada instancia, usa el comando gcloud compute instance-groups managed list-instances:

     gcloud compute instance-groups managed list-instances quickstart-tpu-mig \
         --region=us-east5
     

Instala JAX

Instala las dependencias y el framework de JAX en un entorno virtual en todas las instancias de VM de TPU del MIG. Si tus VMs de TPU tienen instalada una versión de Python anterior a la 3.11, debes instalar Python 3.11 para ejecutar la versión más reciente de JAX.

1. Verifica qué versión de Python se ejecuta en tus VMs de TPU:

   gcloud compute instance-groups managed list-instances quickstart-tpu-mig \
       --region=us-east5 \
       --uri \
   | xargs -I {} -P 0 gcloud compute ssh {} \
       --command='python3 --version'
   

   Si la versión es anterior a Python 3.11, instala Python 3.11:

   gcloud compute instance-groups managed list-instances quickstart-tpu-mig \
       --region=us-east5 \
       --uri \
   | xargs -I {} -P 0 gcloud compute ssh {} \
       --command='sudo apt update && \
       sudo apt install -y software-properties-common && \
       sudo add-apt-repository -y ppa:deadsnakes/ppa && \
       sudo apt update && \
       sudo apt install -y python3.11 python3.11-dev'
   

2. Crea un entorno virtual:

   gcloud compute instance-groups managed list-instances quickstart-tpu-mig \
       --region=us-east5 \
       --uri \
   | xargs -I {} -P 0 gcloud compute ssh {} \
       --command='sudo apt install -y python3.11-venv && \
       python3.11 -m venv ~/jax_venv'
   

3. Instala JAX en el entorno virtual:

   gcloud compute instance-groups managed list-instances quickstart-tpu-mig \
       --region=us-east5 \
       --uri \
   | xargs -I {} -P 0 gcloud compute ssh {} \
       --command='source ~/jax_venv/bin/activate && \
       pip install --upgrade pip -q && \
       pip install jax[tpu] -f https://storage.googleapis.com/jax-releases/libtpu_releases.html -q'
   

Ejecuta el código JAX en la porción

Para ejecutar el código JAX en una porción de TPU, debes ejecutar el código en cada host de la porción de TPU. La llamada a la función jax.device_count() deja de responder hasta que se la llama en cada host de la porción. En el siguiente ejemplo, se muestra cómo ejecutar un cálculo de JAX en una porción de TPU.

Prepara el código

Crea un archivo llamado example.py en cada instancia:

gcloud compute instance-groups managed list-instances quickstart-tpu-mig \
    --region=us-east5 \
    --uri \
| xargs -I {} -P 0 gcloud compute ssh {} \
    --command="cat << 'EOF' > ~/example.py
import jax

# Initialize the slice
jax.distributed.initialize()

# The total number of TPU cores in the slice
device_count = jax.device_count()

# The number of TPU cores attached to this host
local_device_count = jax.local_device_count()

# The psum is performed over all mapped devices across the slice
xs = jax.numpy.ones(jax.local_device_count())
r = jax.pmap(lambda x: jax.lax.psum(x, 'i'), axis_name='i')(xs)

# Print from a single host to avoid duplicated output
if jax.process_index() == 0:
    print('global device count:', jax.device_count())
    print('local device count:', jax.local_device_count())
    print('pmap result:', r)
EOF"

Ejecuta el código en la porción

Ejecuta el programa example.py en cada TPU VM de la porción:

gcloud compute instance-groups managed list-instances quickstart-tpu-mig \
    --region=us-east5 \
    --uri \
| xargs -I {} -P 0 gcloud compute ssh {} \
    --command='source ~/jax_venv/bin/activate && python3 ~/example.py'

El resultado debería ser similar al siguiente ejemplo:

global device count: 8
local device count: 4
pmap result: [8. 8. 8. 8.]