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 TPU v6e.

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

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 do 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:

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 do 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 em lote 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

Variável Descrição
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 da 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}

NUM_SLICES O número de divisões a criar.
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

Variável Descrição
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 da 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}

NUM_SLICES O número de fatias a criar (necessário apenas para o Multislice).
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 usar:

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 de 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 o 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 ambos os cartões de interface de rede (NICs) estão a ser usados 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 Google Cloud registos da consola.

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 do Cloud TPU 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 diminuição da velocidade. Nestes casos, pode reverter para a função malloc padrão anulando a definição da variável LD_PRELOAD na sessão de 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 do software noturno:
    • JAX noturno 0.4.32.dev20240912
    • LibTPU noturno 0.1.dev20240912+nightly
  • Versões de software estáveis:
    • JAX + JAX Lib of 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 gcloud:

    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 JAX, consulte o artigo Imagens de IA 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

    O MaxText é um LLM de código aberto de alto desempenho e altamente escalável escrito em Python e JAX puros e que segmenta TPUs e GPUs para preparação e inferência. Google Cloud 

    JAX_PLATFORMS=tpu,cpu \
ENABLE_PJRT_COMPATIBILITY=true \
TPU_SLICE_BUILDER_DUMP_CHIP_FORCE=true \
TPU_SLICE_BUILDER_DUMP_ICI=true && \
python3 -m MaxText.train MaxText/configs/base.yml \
     base_output_directory=${BASE_OUTPUT_DIR} \
     dataset_type=synthetic \
     per_device_batch_size=${PER_DEVICE_BATCH_SIZE} \
     enable_checkpointing=false \
     gcs_metrics=true \
     profiler=xplane \
     skip_first_n_steps_for_profiler=5 \
     steps=${NUM_STEPS}  # attention='dot_product'"

    Gemma2

    O Gemma é uma família de MDIs/CEs de pesos abertos desenvolvidos pelo Google DeepMind, com base na investigação e tecnologia do Gemini.

    python3 -m MaxText.train MaxText/configs/base.yml \
    model_name=gemma2-27b \
    run_name=gemma2-27b-run \
    base_output_directory=${BASE_OUTPUT_DIR} \
    max_target_length=${MAX_TARGET_LENGTH} \
    per_device_batch_size=${PER_DEVICE_BATCH_SIZE} \
    steps=${NUM_STEPS} \
    enable_checkpointing=false \
    use_iota_embed=true \
    gcs_metrics=true \
    dataset_type=synthetic \
    profiler=xplane \
    attention=flash

    Mixtral 8x7b

    O Mixtral é um modelo de IA de vanguarda desenvolvido pela Mistral AI, que usa uma arquitetura de mistura de especialistas (MoE) esparsa.

    python3 -m MaxText.train MaxText/configs/base.yml \
    base_output_directory=${BASE_OUTPUT_DIR} \
    per_device_batch_size=${PER_DEVICE_BATCH_SIZE} \
    model_name=mixtral-8x7b \
    steps=${NUM_STEPS} \
    max_target_length=${MAX_TARGET_LENGTH} \
    tokenizer_path=assets/tokenizer.mistral-v1 \
    attention=flash \
    dtype=bfloat16 \
    dataset_type=synthetic \
    profiler=xplane

    Llama3-8b

    O Llama é uma família de MDIs de pesos abertos desenvolvidos pela Meta.

    Para ver um exemplo de como executar o Llama3 no PyTorch, consulte os modelos torch_xla no repositório torchprime.

    MaxDiffusion

    O MaxDiffusion é uma coleção de implementações de referência de vários modelos de difusão latente escritos em Python puro e JAX que são executados em dispositivos XLA, incluindo GPUs e TPUs na nuvem. O Stable Diffusion é um modelo de texto para imagem latente que gera imagens fotorrealistas a partir de qualquer entrada de texto.

    Tem de instalar um ramo Git específico para executar o MaxDiffusion, conforme mostrado no seguinte script de preparação.

    git clone https://github.com/google/maxdiffusion.git
&& cd maxdiffusion
&& git checkout 4a8155ec0129512812b31930f0a91c6d5a141103
&& pip install -r requirements.txt && pip install .
&& pip install huggingface_hub==0.30.2 && OUT_DIR=${BASE_OUTPUT_DIR}
&& python src/maxdiffusion/train_sdxl.py \
    src/maxdiffusion/configs/base_xl.yml \
    revision=refs/pr/95 \
    activations_dtype=bfloat16 \
    weights_dtype=bfloat16 \
    resolution=1024 \
    per_device_batch_size=1 \
    output_dir=${OUT_DIR} \
    jax_cache_dir=${OUT_DIR}/cache_dir/ \
    max_train_steps=200 \
    attention=flash \
    run_name=sdxl-ddp-v6e

  3. Exporte as seguintes variáveis:

    export CLUSTER_NAME=CLUSTER_NAME
export ACCELERATOR_TYPE=ACCELERATOR_TYPE
export NUM_SLICES=NUM_SLICES
export YOUR_MODEL_SCRIPT=YOUR_MODEL_SCRIPT

    Descrições das variáveis de ambiente

    Variável Descrição
    CLUSTER_NAME O nome do seu cluster XPK.
    ACCELERATOR_TYPE O tipo de acelerador especifica a versão e o tamanho da Cloud TPU que quer criar. Para mais informações sobre os tipos de aceleradores suportados para cada versão da TPU, consulte o artigo Versões da TPU.
    NUM_SLICES O número de fatias de TPU.
    YOUR_MODEL_SCRIPT O script do modelo a executar como um comando de preparação.

  4. Execute o modelo com o script que criou no passo anterior. Tem de especificar a sinalização --base-docker-image para usar a imagem base MaxText ou especificar a sinalização --docker-image e a imagem que quer usar.

    Pode optar por adicionar as seguintes flags opcionais:

    • Pode ativar o registo de depuração incluindo a flag --enable-debug-logs. Para mais informações, consulte o artigo Depure o JAX no MaxText.
    • Pode criar uma experiência do Vertex AI para carregar dados para o Vertex AI TensorBoard incluindo a flag --use-vertex-tensorboard. Para mais informações, consulte o artigo Monitorize o JAX no MaxText com a Vertex AI.
    python3 xpk.py workload create \
  --cluster ${CLUSTER_NAME} \
  {--base-docker-image maxtext_base_image | --docker-image gcr.io/${PROJECT_ID}/${CLOUD_IMAGE_NAME}:latest} \
  --workload=${USER}-xpk-${ACCELERATOR_TYPE}-${NUM_SLICES} \
  --tpu-type=${ACCELERATOR_TYPE} \
  --num-slices=${NUM_SLICES}  \
  --on-demand \
  --zone=${ZONE} \
  --project=${PROJECT_ID} \
  --command="${YOUR_MODEL_SCRIPT}"

    A saída inclui um link para acompanhar a sua carga de trabalho. Abra o link e clique no separador Registos para monitorizar a sua carga de trabalho em tempo real.

Depure o JAX no MaxText

Use comandos XPK suplementares para diagnosticar o motivo pelo qual o cluster ou a carga de trabalho não está a ser executado:

Monitorize o JAX no MaxText com o Vertex AI

Para usar o TensorBoard, a sua Google Cloud conta de utilizador tem de ter a função aiplatform.user Execute o seguinte comando para conceder esta função:

gcloud projects add-iam-policy-binding your-project-id \
   --member='user:your-email' \
   --role='roles/aiplatform.user'

Veja dados escalares e de perfil através do TensorBoard gerido do Vertex AI.

  1. Aumentar os pedidos de gestão de recursos (CRUD) para a zona que está a usar de 600 para 5000. Isto pode não ser um problema para cargas de trabalho pequenas que usam menos de 16 VMs.

  2. Instale dependências, como cloud-accelerator-diagnostics, para o Vertex AI:

    # xpk dependencies will install cloud-accelerator-diagnostics for Vertex AI
cd ~/xpk
pip install .

  3. Crie o seu cluster XPK com a flag --create-vertex-tensorboard, conforme documentado em Crie o Vertex AI TensorBoard. Também pode executar este comando em clusters existentes.

  4. Crie a sua experiência do Vertex AI quando executar a carga de trabalho XPK usando a flag --use-vertex-tensorboard e a flag --experiment-name opcional. Para ver a lista completa de passos, consulte o artigo Crie uma experiência do Vertex AI para carregar dados para o Vertex AI TensorBoard.

Os registos incluem um link para um Vertex AI TensorBoard, semelhante ao seguinte:

View your TensorBoard at https://us-central1.tensorboard.googleusercontent.com/experiment/project_id+locations+us-central1+tensorboards+hash+experiments+name

Também pode encontrar o link do Vertex AI TensorBoard na Google Cloud consola. Aceda ao Vertex AI Experiments na Google Cloud consola. Selecione a região adequada no menu pendente.

O diretório do TensorBoard também é escrito no contentor do Cloud Storage que especificou com ${BASE_OUTPUT_DIR}.

Elimine a sua carga de trabalho XPK

Use o comando xpk workload delete para eliminar uma ou mais cargas de trabalho com base no prefixo da tarefa ou no estado da tarefa. Este comando pode ser útil se tiver enviado cargas de trabalho XPK que já não precisam de ser executadas ou se tiver tarefas bloqueadas na fila.

Elimine o seu cluster XPK

Use o comando xpk cluster delete para eliminar o cluster:

python3 xpk.py cluster delete --cluster ${CLUSTER_NAME} \
    --zone=${ZONE} --project=${PROJECT_ID}

Resultados dos testes de referência do MaxDiffusion

Executámos o script de preparação do MaxDiffusion num v6e-4, num v6e-16 e em dois v6e-16. A tabela seguinte mostra as taxas de transferência medidas.

v6e-4 v6e-16 Dois v6e-16
Passos de preparação 0,069 0,073 0,13
Tamanho do lote global 8 32 64
Tráfego transmitido (exemplos/seg) 115,9 438,4 492,3

Forme modelos Llama com PyTorch/XLA na Cloud TPU v6e

Esta secção descreve como preparar modelos Llama com o PyTorch/XLA na Cloud TPU v6e usando o conjunto de dados WikiText.

Aceda ao Hugging Face e ao modelo Llama 3

Precisa de uma chave de acesso de utilizador da Hugging Face para este exemplo. Para obter informações sobre como criar chaves de acesso de utilizadores, consulte a documentação do Hugging Face sobre chaves de acesso de utilizadores.

Também precisa de autorização para aceder ao modelo Llama-3-8B no Hugging Face. Para obter acesso, aceda ao modelo Meta-Llama-3-8B no HuggingFace e peça acesso.

Crie uma VM da Cloud TPU

Crie uma Cloud TPU v6e com 8 chips para este exemplo.

  1. Configure variáveis de ambiente:

    export PROJECT_ID=your-project-id
export TPU_NAME=your-tpu-name
export ZONE=us-east1-d
export ACCELERATOR_TYPE=v6e-8
export RUNTIME_VERSION=v2-alpha-tpuv6e

    Descrições das variáveis de ambiente

    Variável Descrição
    PROJECT_ID O seu Google Cloud ID do projeto. Use um projeto existente ou crie um novo.
    TPU_NAME O nome da TPU.
    ZONE A zona na qual criar a VM da TPU. Para mais informações sobre as zonas suportadas, consulte o artigo Regiões e zonas de TPUs.
    ACCELERATOR_TYPE O tipo de acelerador especifica a versão e o tamanho do Cloud TPU que quer criar. Para mais informações sobre os tipos de aceleradores suportados para cada versão da TPU, consulte o artigo Versões da TPU.
    RUNTIME_VERSION A versão do software do Cloud TPU.

  2. Crie uma VM da Cloud TPU:

    gcloud alpha compute tpus tpu-vm create ${TPU_NAME} --version=${RUNTIME_VERSION} \
   --accelerator-type=${ACCELERATOR_TYPE} \
   --zone=${ZONE} \
   --project=${PROJECT_ID}

Instalação

Instale a ramificação pytorch-tpu/transformers dos transformadores e das dependências do Hugging Face. Este exemplo foi testado com as seguintes versões de dependências:

  • torch: compatível com a versão 2.5.0
  • torch_xla[tpu]: compatível com a versão 2.5.0
  • jax: 0.4.33
  • jaxlib: 0.4.33
gcloud alpha compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone ${ZONE} \
   --worker=all \
   --command='git clone -b flash_attention https://github.com/pytorch-tpu/transformers.git
   cd transformers
   sudo pip3 install -e .
   pip3 install datasets
   pip3 install evaluate
   pip3 install scikit-learn
   pip3 install accelerate
   pip install torch~=2.6.0 torch_xla[tpu]~=2.6.0 -f https://storage.googleapis.com/libtpu-releases/index.html -f https://storage.googleapis.com/libtpu-wheels/index.html
   pip install jax==0.4.38 jaxlib==0.4.38 -i https://us-python.pkg.dev/ml-oss-artifacts-published/jax/simple/'

Configure ficheiros de configuração do modelo

O comando de preparação na secção seguinte, Executar o modelo, usa dois ficheiros de configuração JSON para definir os parâmetros do modelo e a configuração de paralelismo de dados totalmente fragmentados (FSDP). A divisão do FSDP permite-lhe usar um tamanho do lote maior durante o treino dividindo as ponderações do modelo em várias UTPs. Quando a preparação é feita com modelos mais pequenos, pode ser suficiente usar o paralelismo de dados e replicar os pesos em cada dispositivo. Para mais informações sobre como dividir tensores em vários dispositivos no PyTorch/XLA, consulte o guia do utilizador SPMD do PyTorch/XLA.

  1. Crie o ficheiro de configuração dos parâmetros do modelo. Segue-se a configuração dos parâmetros do modelo para o Llama-3-8B. Para outros modelos, encontre o ficheiro de configuração no Hugging Face. Por exemplo, consulte a configuração do Llama-2-7B.

    cat > llama-config.json << EOF
{
  "architectures": [
    "LlamaForCausalLM"
  ],
  "attention_bias": false,
  "attention_dropout": 0.0,
  "bos_token_id": 128000,
  "eos_token_id": 128001,
  "hidden_act": "silu",
  "hidden_size": 4096,
  "initializer_range": 0.02,
  "intermediate_size": 14336,
  "max_position_embeddings": 8192,
  "model_type": "llama",
  "num_attention_heads": 32,
  "num_hidden_layers": 32,
  "num_key_value_heads": 8,
  "pretraining_tp": 1,
  "rms_norm_eps": 1e-05,
  "rope_scaling": null,
  "rope_theta": 500000.0,
  "tie_word_embeddings": false,
  "torch_dtype": "bfloat16",
  "transformers_version": "4.40.0.dev0",
  "use_cache": false,
  "vocab_size": 128256
}
EOF

  2. Crie o ficheiro de configuração do FSDP:

    cat > fsdp-config.json << EOF
{
  "fsdp_transformer_layer_cls_to_wrap": [
    "LlamaDecoderLayer"
  ],
  "xla": true,
  "xla_fsdp_v2": true,
  "xla_fsdp_grad_ckpt": true
}
EOF

    Para mais informações sobre a FSDP, consulte o artigo Paralelismo de dados totalmente fragmentados com SPMD .

  3. Carregue os ficheiros de configuração para as VMs do Cloud TPU através do seguinte comando:

    gcloud alpha compute tpus tpu-vm scp llama-config.json fsdp-config.json ${TPU_NAME}:. \
   --worker=all \
   --project=${PROJECT_ID} \
   --zone=${ZONE}

Execute o modelo

Usando os ficheiros de configuração que criou na secção anterior, execute o script run_clm.py para preparar o modelo Llama-3-8B no conjunto de dados WikiText. O script de treino demora aproximadamente 10 minutos a ser executado num Cloud TPU v6e-8.

  1. Inicie sessão no Hugging Face na sua Cloud TPU através do seguinte comando:

    gcloud alpha compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone ${ZONE} \
   --worker=all \
   --command='
   pip3 install "huggingface_hub[cli]"
   huggingface-cli login --token HUGGING_FACE_TOKEN'

  2. Execute a preparação do modelo:

    gcloud alpha compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone ${ZONE} \
   --worker=all \
   --command='
   export PJRT_DEVICE=TPU
   export XLA_USE_SPMD=1
   export ENABLE_PJRT_COMPATIBILITY=true
   # Optional variables for debugging:
   export XLA_IR_DEBUG=1
   export XLA_HLO_DEBUG=1
   export PROFILE_EPOCH=0
   export PROFILE_STEP=3
   export PROFILE_DURATION_MS=100000
   # Set PROFILE_LOGDIR to a local VM path or gs://my-bucket/profile_path
   export PROFILE_LOGDIR=PROFILE_PATH
   python3 transformers/examples/pytorch/language-modeling/run_clm.py \
     --dataset_name wikitext \
     --dataset_config_name wikitext-2-raw-v1 \
     --per_device_train_batch_size 16 \
     --do_train \
     --output_dir /home/$USER/tmp/test-clm \
     --overwrite_output_dir \
     --config_name /home/$USER/llama-config.json \
     --cache_dir /home/$USER/cache \
     --tokenizer_name meta-llama/Meta-Llama-3-8B \
     --block_size 8192 \
     --optim adafactor \
     --save_strategy no \
     --logging_strategy no \
     --fsdp "full_shard" \
     --fsdp_config /home/$USER/fsdp-config.json \
     --torch_dtype bfloat16 \
     --dataloader_drop_last yes \
     --flash_attention \
     --max_steps 20'

Resolução de problemas do PyTorch/XLA

Se definir as variáveis opcionais para depuração na secção anterior, o perfil do modelo é armazenado na localização especificada pela variável PROFILE_LOGDIR. Pode extrair o ficheiro xplane.pb armazenado nesta localização e usar tensorboard para ver os perfis no seu navegador através das instruções do TensorBoard.

Se o PyTorch/XLA não estiver a ter o desempenho esperado, consulte o guia de resolução de problemas, que tem sugestões para depurar, criar perfis e otimizar o seu modelo.