Prepare um modelo com a TPU v6e

Este documento explica como preparar modelos no Cloud TPU v6e (também denominado Trillium), abordando a configuração do ambiente, a otimização do desempenho e exemplos de preparação práticos com JAX e PyTorch/XLA.

A TPU v6e, também denominada Trillium, é a 6.ª geração de TPUs da Google. Em todas as superfícies técnicas, como a API e os registos, e ao longo deste documento, o Trillium é referido como v6e. Com 256 chips por Pod, a arquitetura da TPU v6e partilha muitas semelhanças com a v5e. A TPU v6e está otimizada para a preparação, o ajuste fino e a publicação de transformadores, conversão de texto em imagem e redes neurais convolucionais (CNNs). Para mais informações sobre a arquitetura e as configurações do sistema TPU v6e, consulte o artigo TPU v6e.

Para informações sobre a execução da inferência na Cloud TPU v6e, consulte os seguintes tutoriais:

• Inferência do MaxDiffusion na v6e
• Inferência de vLLM no v6e
• Realize a inferência em vários anfitriões com o Pathways

Antes de começar

Antes de começar, tem de:

• Crie uma Google Cloud conta e um projeto com a faturação ativada
• Instale os componentes alfa da CLI Google Cloud
• Ative a API Cloud TPU
• Crie um agente de serviço da Cloud TPU
• Crie uma conta de serviço da Cloud TPU e conceda autorizações

Para mais informações, consulte o artigo Configure o ambiente do Cloud TPU.

Valide a quota e as autorizações

Verifique se o seu projeto tem as seguintes quotas:

• Quota de TPU v6e preemptível ou a pedido
• Quota de endereços IP

• Quota para o Hyperdisk Balanced e para quaisquer outros tipos de discos que queira usar

Se estiver a usar o GKE com o XPK, precisa de autorizações adicionais na Google Cloud consola. Para mais informações, consulte o artigo Autorizações necessárias na Google Cloud consola .

Aprovisione TPUs

Pode aprovisionar e gerir TPUs v6e através dos seguintes métodos:

• GKE: pode usar o GKE para aprovisionar e gerir TPUs como um conjunto de aceleradores para as suas cargas de trabalho de aprendizagem automática contentorizadas. Para mais informações, consulte o artigo Acerca das TPUs no GKE.
• GKE e XPK: o XPK é uma ferramenta de linha de comandos que simplifica a criação de clusters e a execução de cargas de trabalho no GKE. Foi concebido para que os profissionais de ML possam aprovisionar TPUs e executar tarefas de preparação sem precisarem de conhecimentos profundos de Kubernetes. Para mais informações, consulte o repositório do GitHub do XPK.
• Recursos em fila da Cloud TPU: os recursos em fila permitem-lhe pedir capacidade da TPU que é aprovisionada quando fica disponível. É ideal para tarefas de lotes e cargas de trabalho com tolerância a falhas que podem esperar numa fila. Pode especificar um período para o seu pedido. Para mais informações, consulte o artigo Faça a gestão dos recursos em fila.

Aprovisione Cloud TPUs v6e com o GKE e o XPK

Se estiver a usar comandos do GKE com a v6e, pode usar comandos do Kubernetes ou XPK para aprovisionar TPUs na nuvem e preparar ou publicar modelos. Consulte o artigo Planeie TPUs na nuvem no GKE para saber como planear as suas configurações de Cloud TPU em clusters do GKE. As secções seguintes fornecem comandos para criar um cluster XPK com suporte de placa de rede única e suporte de várias placas de rede.

Crie um cluster XPK com suporte de placa de rede única

export CLUSTER_NAME=xpk-cluster-name
export ZONE=us-east1-d
export PROJECT_ID=your-project-id
export TPU_TYPE=v6e-256
export NUM_SLICES=2

export NETWORK_NAME=${CLUSTER_NAME}-mtu9k
export NETWORK_FW_NAME=${NETWORK_NAME}-fw

gcloud compute networks create ${NETWORK_NAME} \
   --mtu=8896 \
   --project=${PROJECT_ID} \
   --subnet-mode=auto \
   --bgp-routing-mode=regional

gcloud compute firewall-rules create ${NETWORK_FW_NAME} \
   --network=${NETWORK_NAME} \
   --allow tcp,icmp,udp \
   --project=${PROJECT_ID}

export CLUSTER_ARGUMENTS="--network=${NETWORK_NAME} --subnetwork=${NETWORK_NAME}"

python3 xpk.py cluster create --cluster=${CLUSTER_NAME} \
   --cluster-cpu-machine-type=e2-standard-8 \
   --num-slices=${NUM_SLICES} \
   --tpu-type=${TPU_TYPE} \
   --zone=${ZONE} \
   --project=${PROJECT_ID} \
   --on-demand \
   --custom-cluster-arguments="${CLUSTER_ARGUMENTS}" \
   --create-vertex-tensorboard

Descrições das flags de comando

• CLUSTER_NAME: o nome atribuído pelo utilizador para o cluster XPK.
• PROJECT_ID: Google Cloud nome do projeto. Use um projeto existente ou crie um novo. Para mais informações, consulte o artigo Configure o seu Google Cloud projeto.
• ZONE: consulte o documento Regiões e zonas do Cloud TPU para ver as zonas suportadas.
• TPU_TYPE: consulte os tipos de aceleradores.
• NUM_SLICES: o número de divisões que quer criar
• CLUSTER_ARGUMENTS: a rede e a sub-rede a usar.

  Por exemplo: --network=${NETWORK_NAME} --subnetwork=${NETWORK_NAME}

• NETWORK_NAME: o nome de uma rede secundária a usar.
• NETWORK_FW_NAME: o nome de uma firewall de rede secundária a usar.

Crie um cluster XPK com suporte para várias NICs

export CLUSTER_NAME=xpk-cluster-name
export REGION=your-region
export ZONE=us-east1-d
export PROJECT_ID=your-project-id
export TPU_TYPE=v6e-256
export NUM_SLICES=2

export NETWORK_NAME_1=${CLUSTER_NAME}-mtu9k-1-${ZONE}
export SUBNET_NAME_1=${CLUSTER_NAME}-privatesubnet-1-${ZONE}
export NETWORK_FW_NAME_1=${NETWORK_NAME_1}-fw-1-${ZONE}
export FIREWALL_RULE_NAME=${CLUSTER_NAME}-privatefirewall-1-${ZONE}
export ROUTER_NAME=${CLUSTER_NAME}-network-1-${ZONE}
export NAT_CONFIG=${CLUSTER_NAME}-natconfig-1-${ZONE}

gcloud compute networks create ${NETWORK_NAME_1} \
   --mtu=8896 \
   --bgp-routing-mode=regional \
   --subnet-mode=custom \
   --project=${PROJECT_ID}

gcloud compute networks subnets create ${SUBNET_NAME_1} \
   --network=${NETWORK_NAME_1} \
   --range=10.11.0.0/18 \
   --region=${REGION} \
   --project=${PROJECT_ID}

gcloud compute firewall-rules create ${FIREWALL_RULE_NAME} \
   --network=${NETWORK_NAME_1} \
   --allow tcp,icmp,udp \
   --project=${PROJECT_ID}

gcloud compute routers create ${ROUTER_NAME} \
   --project=${PROJECT_ID} \
   --network=${NETWORK_NAME_1} \
   --region=${REGION}

gcloud compute routers nats create ${NAT_CONFIG} \
   --router=${ROUTER_NAME} \
   --region=${REGION} \
   --auto-allocate-nat-external-ips \
   --nat-all-subnet-ip-ranges \
   --project=${PROJECT_ID} \
   --enable-logging

# Secondary subnet for multi-nic experience.
# Need custom IP routing to be different from the first network's subnet.

export NETWORK_NAME_2=${CLUSTER_NAME}-privatenetwork-2-${ZONE}
export SUBNET_NAME_2=${CLUSTER_NAME}-privatesubnet-2-${ZONE}
export FIREWALL_RULE_NAME=${CLUSTER_NAME}-privatefirewall-2-${ZONE}
export ROUTER_NAME=${CLUSTER_NAME}-network-2-${ZONE}
export NAT_CONFIG=${CLUSTER_NAME}-natconfig-2-${ZONE}

gcloud compute networks create ${NETWORK_NAME_2} \
   --mtu=8896 \
   --bgp-routing-mode=regional \
   --subnet-mode=custom \
   --project=${PROJECT_ID}

gcloud compute networks subnets create ${SUBNET_NAME_2} \
   --network=${NETWORK_NAME_2} \
   --range=10.10.0.0/18 \
   --region=${REGION} \
   --project=${PROJECT_ID}

gcloud compute firewall-rules create ${FIREWALL_RULE_NAME} \
   --network=${NETWORK_NAME_2} \
   --allow tcp,icmp,udp \
   --project=${PROJECT_ID}

gcloud compute routers create ${ROUTER_NAME} \
   --project=${PROJECT_ID} \
   --network=${NETWORK_NAME_2} \
   --region=${REGION}

gcloud compute routers nats create ${NAT_CONFIG} \
   --router=${ROUTER_NAME} \
   --region=${REGION} \
   --auto-allocate-nat-external-ips \
   --nat-all-subnet-ip-ranges \
   --project=${PROJECT_ID} \
   --enable-logging

export CLUSTER_ARGUMENTS="--enable-dataplane-v2 --enable-ip-alias --enable-multi-networking --network=${NETWORK_NAME_1} --subnetwork=${SUBNET_NAME_1}"
export NODE_POOL_ARGUMENTS="--additional-node-network network=${NETWORK_NAME_2},subnetwork=${SUBNET_NAME_2}"

python3 xpk.py cluster create \
   --cluster=${CLUSTER_NAME} \
   --cluster-cpu-machine-type=e2-standard-8 \
   --num-slices=${NUM_SLICES} \
   --tpu-type=${TPU_TYPE} \
   --zone=${ZONE}  \
   --project=${PROJECT_ID} \
   --on-demand \
   --custom-cluster-arguments="${CLUSTER_ARGUMENTS}" \
   --custom-nodepool-arguments="${NODE_POOL_ARGUMENTS}" \
   --create-vertex-tensorboard

Descrições das flags de comando

• CLUSTER_NAME: o nome atribuído pelo utilizador para o cluster XPK.
• PROJECT_ID: Google Cloud nome do projeto. Use um projeto existente ou crie um novo. Para mais informações, consulte o artigo Configure o seu Google Cloud projeto.
• ZONE: consulte o documento Regiões e zonas do Cloud TPU para ver as zonas suportadas.
• TPU_TYPE: consulte os tipos de aceleradores.
• NUM_SLICES: o número de divisões que quer criar
• CLUSTER_ARGUMENTS: a rede e a sub-rede a usar.

  Por exemplo: --enable-dataplane-v2 --enable-ip-alias --enable-multi-networking --network=${NETWORK_NAME_1} --subnetwork=${SUBNET_NAME_1}

• NODE_POOL_ARGUMENTS: rede de nós adicional a usar.

  Por exemplo: --additional-node-network network=${NETWORK_NAME_2},subnetwork=${SUBNET_NAME_2}

• NETWORK_NAME: o nome de uma rede secundária a usar.
• NETWORK_FW_NAME: o nome de uma firewall de rede secundária a usar.

Configure o JAX ou o PyTorch

Os seguintes recursos mostram como configurar o JAX ou o PyTorch no seu Cloud TPU, consoante o método de aprovisionamento e gestão que usa:

• GKE Autopilot: prepare a sua aplicação de TPU
• GKE Standard: prepare as suas cargas de trabalho
• GKE e XPK: XPK README
• Cloud TPU de anfitrião único com JAX: execute um cálculo numa VM do Cloud TPU com JAX
• TPU do Google Cloud com vários anfitriões através do JAX: execute código JAX em fatias de TPU
• Cloud TPU de anfitrião único com o PyTorch: execute um cálculo numa VM da Cloud TPU com o PyTorch
• Cloud TPU multi-host com PyTorch: execute código PyTorch em fatias de TPU

Para configurar e executar o XPK com o MaxText, consulte o artigo Executar o MaxText em grande escala com o XPK .

Otimize o desempenho da rede

Esta secção descreve como otimizar o desempenho da sua rede configurando a unidade de transmissão máxima (MTU), usando várias NICs para ambientes de divisão múltipla e melhorando as definições de TCP.

Configure a MTU

Para o melhor desempenho da rede, use uma rede com 8896 MTU (unidade de transmissão máxima).

Por predefinição, uma nuvem virtual privada (VPC) só fornece uma MTU de 1460 bytes, o que oferece um desempenho de rede abaixo do ideal. Pode definir a MTU de uma rede VPC para qualquer valor entre 1300 bytes e 8896 bytes (inclusive). Os tamanhos de MTU personalizados comuns são 1500 bytes (Ethernet padrão) ou 8896 bytes (o máximo possível). Para mais informações, consulte o artigo Tamanhos de MTU de rede VPC válidos.

Para mais informações sobre como alterar a definição de MTU para uma rede existente ou predefinida, consulte o artigo Altere a definição de MTU de uma rede VPC.

O exemplo seguinte cria uma rede com uma MTU de 8896 e uma regra de firewall correspondente que permite o tráfego TCP, ICMP e UDP na rede.

export RESOURCE_NAME=your-resource-name
export NETWORK_NAME=${RESOURCE_NAME}-privatenetwork
export NETWORK_FW_NAME=${RESOURCE_NAME}-privatefirewall
gcloud compute networks create ${NETWORK_NAME} --mtu=8896 --project=${PROJECT_ID} \
    --subnet-mode=auto --bgp-routing-mode=regional
gcloud compute firewall-rules create ${NETWORK_FW_NAME} --network=${NETWORK_NAME} \
    --allow tcp,icmp,udp --project=${PROJECT_ID}

Substitua your-resource-name por um nome base para a rede e a firewall.

Use a opção de várias NICs para o Multislice

Se estiver a usar um ambiente de divisão múltipla, defina as seguintes variáveis de ambiente, que são necessárias para uma sub-rede secundária:

export NETWORK_NAME_2=${RESOURCE_NAME}
export SUBNET_NAME_2=${RESOURCE_NAME}
export FIREWALL_RULE_NAME=${RESOURCE_NAME}
export ROUTER_NAME=${RESOURCE_NAME}-network-2
export NAT_CONFIG=${RESOURCE_NAME}-natconfig-2
export REGION=your-region

Use os seguintes comandos para criar um encaminhamento IP personalizado para a rede e a sub-rede.

1. Crie a rede secundária.

   gcloud compute networks create ${NETWORK_NAME_2} --mtu=8896 \
   --bgp-routing-mode=regional --subnet-mode=custom --project=${PROJECT_ID}
   

2. Crie uma sub-rede para a rede secundária.

   gcloud compute networks subnets create ${SUBNET_NAME_2} \
   --network=${NETWORK_NAME_2} \
   --range=10.10.0.0/18 --region=${REGION} \
   --project=${PROJECT_ID}
   

3. Crie uma regra de firewall para permitir o tráfego na nova sub-rede.

   gcloud compute firewall-rules create ${FIREWALL_RULE_NAME} \
   --network=${NETWORK_NAME_2} --allow tcp,icmp,udp \
   --source-ranges 10.10.0.0/18 --project=${PROJECT_ID}
   

4. Crie um Cloud Router para a rede secundária.

   gcloud compute routers create ${ROUTER_NAME} \
   --project=${PROJECT_ID} \
   --network=${NETWORK_NAME_2} \
   --region=${REGION}
   

5. Crie uma configuração de NAT para o Cloud Router.

   gcloud compute routers nats create ${NAT_CONFIG} \
   --router=${ROUTER_NAME} \
   --region=${REGION} \
   --auto-allocate-nat-external-ips \
   --nat-all-subnet-ip-ranges \
   --project=${PROJECT_ID} \
   --enable-logging
   

Depois de criar uma fatia de várias redes, pode validar se ambas as placas de interface de rede (NICs) estão a ser usadas configurando um cluster XPK e adicionando a flag --command ifconfig ao comando de criação da carga de trabalho XPK.

1. Use o seguinte comando workload create para apresentar o resultado do comando ifconfig nos registos da consola e verifique se eth0 e eth1 têm o MTU definido como 8896. Google Cloud

   python3 xpk.py workload create \
       --cluster CLUSTER_NAME \
       {--base-docker-image maxtext_base_image | --docker-image your-cloud-image-name} \
       --workload=${USER}-xpk-${ACCELERATOR_TYPE}-${NUM_SLICES} \
       --tpu-type=${ACCELERATOR_TYPE} \
       --num-slices=${NUM_SLICES}  \
       --on-demand \
       --zone=${ZONE} \
       --project=${PROJECT_ID} \
       --command "ifconfig"

   Se quiser ativar os registos de depuração ou usar o Vertex AI TensorBoard, adicione os seguintes argumentos opcionais ao comando:

   --enable-debug-logs \
   --use-vertex-tensorboard

2. Verifique se eth0 e eth1 têm o MTU definido como 8896,verificando o resultado da carga de trabalho XPK nos registos da consola. Google Cloud

Melhore as definições de TCP

Se aprovisionou as suas TPUs na nuvem através de recursos em fila, pode executar o seguinte comando para melhorar o desempenho da rede aumentando os limites do buffer de receção de TCP.

gcloud alpha compute tpus queued-resources ssh "${QUEUED_RESOURCE_ID}" \
    --project "${PROJECT_ID}" \
    --zone "${ZONE}" \
    --node=all \
    --worker=all \
    --command='
    sudo sh -c "echo \"4096 41943040 314572800\" > /proc/sys/net/ipv4/tcp_rmem"'

Otimize o desempenho da atribuição de memória

A biblioteca tcmalloc é usada por predefinição nas VMs de TPU do Google Cloud para melhorar o desempenho dos modelos com alocações de memória frequentes e consideráveis. Isto é configurado através da variável de ambiente LD_PRELOAD.

No entanto, para algumas cargas de trabalho (por exemplo, DLRM com alocações de tabelas de incorporação muito grandes), o tcmalloc pode causar uma desaceleração. Nestes casos, pode reverter para a função malloc padrão anulando a definição da variável LD_PRELOAD na sessão do shell antes de executar o script de preparação:

unset LD_PRELOAD

Use o SkyPilot

Pode usar a Cloud TPU v6e com o SkyPilot. O SkyPilot é um framework de código aberto que simplifica o processo de execução, gestão e escalabilidade de cargas de trabalho de IA. Pode adicionar informações de localização e preços relacionadas com o v6e ao SkyPilot. Para mais informações, consulte o exemplo da TPU v6e do SkyPilot.

Exemplos de preparação

As secções seguintes fornecem exemplos para preparar modelos MaxText, MaxDiffusion e PyTorch na Cloud TPU v6e.

Estes exemplos foram testados com as seguintes versões de software:

• Python 3.10 ou posterior
• Versões de software noturnas:
  • JAX noturno 0.4.32.dev20240912
  • LibTPU noturno 0.1.dev20240912+nightly
• Versões de software estáveis:
  • JAX + JAX Lib v0.4.37

Forme o MaxText e o MaxDiffusion na Cloud TPU v6e

As secções seguintes abordam o ciclo de vida da preparação dos modelos MaxText e MaxDiffusion.

Em geral, os passos de alto nível são os seguintes:

1. Crie a imagem base da carga de trabalho.
2. Execute a sua carga de trabalho com o XPK.
   1. Crie o comando de preparação para a carga de trabalho.
   2. Implemente a carga de trabalho.
3. Acompanhe a carga de trabalho e veja as métricas.
4. Elimine a carga de trabalho XPK se não for necessária.
5. Elimine o cluster XPK quando já não for necessário.

Crie uma imagem base

Instale o MaxText ou o MaxDiffusion e crie a imagem do Docker:

1. Clone o repositório que quer usar e altere para o diretório do repositório:

   MaxText:

   git clone https://github.com/google/maxtext.git && cd maxtext
   

   MaxDiffusion:

   git clone https://github.com/google/maxdiffusion.git && cd maxdiffusion && git checkout 4a8155ec0129512812b31930f0a91c6d5a141103
   

2. Configure o Docker para usar a CLI do Google Cloud:

   gcloud auth configure-docker
   

3. Crie a imagem do Docker com o seguinte comando ou com uma imagem de IA do JAX. Para mais informações sobre as imagens de IA do JAX, consulte o artigo Imagens de IA do JAX.

   MaxText:

   bash docker_build_dependency_image.sh MODE=stable JAX_VERSION=0.4.35
   

   MaxDiffusion:

   bash .github/workflows/build_and_upload_images.sh CLOUD_IMAGE_NAME=maxdiffusion_jax_stable_stack MODE=jax_ai_image PROJECT=${PROJECT_ID} LOCAL_IMAGE_NAME=maxdiffusion_jax_stable_stack BASEIMAGE=us-docker.pkg.dev/cloud-tpu-images/jax-ai-image/tpu:latest
   

4. Defina o ID do projeto na configuração da CLI gcloud ativa:

   gcloud config set project ${PROJECT_ID}
   

5. Se estiver a iniciar a carga de trabalho a partir de uma máquina que não tenha a imagem criada localmente, carregue a imagem.

   1. Defina a variável de ambiente CLOUD_IMAGE_NAME:

      export CLOUD_IMAGE_NAME=${USER}_runner
      

   2. Carregue a imagem:

      bash docker_upload_runner.sh ${CLOUD_IMAGE_NAME}
      

Execute a sua carga de trabalho com o XPK

1. Defina as seguintes variáveis de ambiente se não estiver a usar os valores predefinidos definidos pelo MaxText ou MaxDiffusion:

   export BASE_OUTPUT_DIR=gs://YOUR_BUCKET
   export PER_DEVICE_BATCH_SIZE=2
   export NUM_STEPS=30
   export MAX_TARGET_LENGTH=8192

2. Crie o script do modelo. Este script vai ser copiado como um comando de preparação num passo posterior.

   Não execute ainda o script do modelo.

   MaxText