Criar uma fração de TPU com vários hosts

Criar um modelo de instância

Para criar um modelo de instância para VMs de TPU v6e, use o 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

Criar uma política de carga de trabalho

Uma política de carga de trabalho define propriedades físicas das suas instâncias de computação. Em frações de TPU, a topologia do acelerador define a disposição física dos chips de TPU. É necessário especificar uma topologia de acelerador para frações de TPU interconectadas de vários hosts.

Para criar uma política de carga de trabalho para uma fração de TPU de vários hosts, use o comando gcloud compute resource-policies create workload-policy com a flag --accelerator-topology. O comando a seguir cria uma política de carga de trabalho com uma topologia 2x4:

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

Criar um MIG

Execute os comandos a seguir para criar uma MIG que forma uma fração de TPU de vários hosts.

1. Para criar um MIG que forma uma fração de TPU de vários hosts, use o 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
   

   Substitua PROJECT_ID pelo ID do projeto Cloud de Confiance by S3NS .

2. Se quiser, verifique se as instâncias gerenciadas estão em execução usando os seguintes comandos:

   • Para conferir o status geral do MIG, use o comando gcloud compute instance-groups managed describe:

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

   • Para conferir o status de cada instância, use o comando gcloud compute instance-groups managed list-instances:

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

Instalar o JAX

Instale as dependências e o framework JAX em um ambiente virtual em todas as instâncias de VM de TPU no MIG. Se as VMs da TPU tiverem uma versão do Python anterior à 3.11 instalada, instale o Python 3.11 para executar a versão mais recente do JAX.

1. Verifique qual versão do Python está sendo executada nas suas VMs da 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'
   

   Se a versão for anterior ao Python 3.11, instale o 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. Crie um ambiente 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. Instale o JAX no ambiente 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'
   

Executar o código JAX na fração

Para executar o código JAX em uma fração de TPU, é preciso executá-lo em cada host dessa fração. A chamada de função jax.device_count() para de responder até ser chamada em cada host na fração. O exemplo a seguir mostra como executar um cálculo JAX em uma fração de TPU.

Preparar o código

Crie um arquivo chamado example.py em cada instância:

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"

Executar o código na fração

Execute o programa example.py em cada VM de TPU na fração:

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'

A saída será semelhante a esta:

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