Treinar um modelo usando a TPU v6e

Este documento orienta você no treinamento de modelos na Cloud TPU v6e (também chamada de Trillium), abordando a configuração do ambiente, a otimização de desempenho e exemplos práticos de treinamento usando JAX e PyTorch/XLA.

A TPU v6e, também chamada de Trillium, é a 6ª geração de TPUs do Google. Em todas as plataformas técnicas, como a API e os registros, e ao longo deste documento, o Trillium será chamado de v6e. Com 256 chips por Pod, a arquitetura da TPU v6e compartilha muitas semelhanças com a v5e. A TPU v6e é otimizada para treinamento, ajuste e veiculação de transformadores, texto para 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 informações sobre como executar inferências no Cloud TPU v6e, consulte os seguintes tutoriais:

• Inferência do JetStream MaxText na v6e
• Inferência do PyTorch do JetStream na v6e
• Inferência do MaxDiffusion na v6e
• Inferência de vLLM na v6e
• Realizar inferência multihost usando o Pathways

Antes de começar

Antes de começar, faça o seguinte:

• Crie uma Google Cloud conta e um projeto com o faturamento ativado
• Instalar os componentes Alfa da CLI Google Cloud
• Ativar a API Cloud TPU
• Criar um agente de serviço da Cloud TPU
• Criar uma conta de serviço do Cloud TPU e conceder permissões

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

Verificar cota e permissões

Verifique se o projeto tem as seguintes cotas:

• Cota preemptiva ou sob demanda da TPU v6e
• Cota de endereços IP

• Cota para o Hyperdisk Balanced e para qualquer outro tipo de disco que você queira usar

  Observação: o v6e é compatível com o Hyperdisk Balanced e o Hyperdisk ML. Ele não é compatível com Persistent Disk (DP). Para mais informações, consulte Opções de armazenamento para dados do Cloud TPU.

Se você estiver usando o GKE com o XPK, precisará de outras permissões no console Google Cloud . Para mais informações, consulte Permissões necessárias no Google Cloud console .

Provisionar TPUs

É possível provisionar e gerenciar TPUs v6e usando os seguintes métodos:

• GKE: é possível usar o GKE para provisionar e gerenciar TPUs como um pool de aceleradores para cargas de trabalho de aprendizado de máquina em contêineres. Para mais informações, consulte Sobre TPUs no GKE.
• GKE e XPK: o XPK é uma ferramenta de linha de comando que simplifica a criação de clusters e a execução de cargas de trabalho no GKE. Ele foi criado para profissionais de ML provisionarem TPUs e executarem jobs de treinamento sem precisar de experiência no uso de Kubernetes. Para mais informações, consulte o repositório do XPK no GitHub.
• Recursos enfileirados da Cloud TPU: com os recursos enfileirados, é possível solicitar capacidade de TPU provisionada quando fica disponível. É ideal para jobs em lote e cargas de trabalho tolerantes a falhas que podem esperar em uma fila. É possível especificar uma janela de tempo para sua solicitação. Para mais informações, consulte Gerenciar recursos em fila.

Provisionar Cloud TPUs v6e com o GKE e o XPK

Se você estiver usando comandos do GKE com v6e, use comandos do Kubernetes ou XPK para provisionar Cloud TPUs e treinar ou veicular modelos. Consulte Planejar o uso de Cloud TPUs no GKE para saber como planejar as configurações da Cloud TPU em clusters do GKE. As seções a seguir fornecem comandos para criar um cluster XPK com suporte a uma única NIC e várias NICs.

Criar um cluster XPK com suporte a NIC ú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

Observação: use uma rede personalizada para melhorar a performance e evitar a sobrecarga da rede padrão.

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 de sinalizações de comando

Variável	Descrição
CLUSTER_NAME	O nome atribuído pelo usuário para o cluster XPK.
PROJECT_ID	Google Cloud nome do projeto. Use um projeto atual ou crie um novo. Para mais informações, consulte Configurar seu Google Cloud projeto.
ZONA	Consulte o documento Regiões e zonas do Cloud TPU para saber quais são as zonas compatíveis.
TPU_TYPE	Consulte Tipos de aceleradores.
NUM_SLICES	O número de intervalos que você quer criar
CLUSTER_ARGUMENTS	A rede e a sub-rede a serem usadas. Por exemplo: --network=${NETWORK_NAME} --subnetwork=${NETWORK_NAME}
NUM_SLICES	O número de intervalos a serem criados.
NETWORK_NAME	O nome de uma rede secundária a ser usada.
NETWORK_FW_NAME	O nome de um firewall de rede secundário a ser usado.

Criar um cluster XPK com suporte a 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}

Observação: use uma rede personalizada para melhorar a performance e evitar a sobrecarga da rede padrão.

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 de sinalizações de comando

Variável	Descrição
CLUSTER_NAME	O nome atribuído pelo usuário para o cluster XPK.
PROJECT_ID	Google Cloud nome do projeto. Use um projeto atual ou crie um novo. Para mais informações, consulte Configurar seu Google Cloud projeto.
ZONA	Consulte o documento Regiões e zonas do Cloud TPU para saber quais são as zonas compatíveis.
TPU_TYPE	Consulte Tipos de aceleradores.
NUM_SLICES	O número de intervalos que você quer criar
CLUSTER_ARGUMENTS	A rede e a sub-rede a serem usadas. 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 ser usada. Por exemplo: --additional-node-network network=${NETWORK_NAME_2},subnetwork=${SUBNET_NAME_2}
NUM_SLICES	O número de intervalos a serem criados (necessário apenas para Multislice).
NETWORK_NAME	O nome de uma rede secundária a ser usada.
NETWORK_FW_NAME	O nome de um firewall de rede secundário a ser usado.

Configurar o JAX ou o PyTorch

Os recursos a seguir mostram como configurar o JAX ou o PyTorch no Cloud TPU, dependendo do método de provisionamento e gerenciamento usado:

• GKE Autopilot: preparar o aplicativo de TPU
• GKE Standard: preparar as cargas de trabalho
• GKE e XPK: README do XPK
• Cloud TPU de host único usando JAX: executar um cálculo em uma VM da Cloud TPU usando JAX
• Cloud TPU de vários hosts usando JAX: executar código JAX em frações de TPU
• Cloud TPU de host único usando PyTorch: executar um cálculo em uma VM do Cloud TPU usando o PyTorch
• Cloud TPU de vários hosts usando o PyTorch: executar código do PyTorch em frações de TPU

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

Otimizar o desempenho da rede

Esta seção descreve como otimizar o desempenho da rede configurando a unidade máxima de transmissão (MTU), usando várias NICs para ambientes Multislice e melhorando as configurações de TCP.

Configurar a MTU

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

Por padrão, uma nuvem privada virtual (VPC) só fornece uma MTU de 1.460 bytes, o que resulta em um desempenho de rede abaixo do ideal. É possível definir a MTU de uma rede VPC como qualquer valor entre 1.300 e 8.896 bytes (inclusive). Os tamanhos comuns de MTU personalizados são 1.500 bytes (Ethernet padrão) ou 8.896 bytes (o máximo possível). Para mais informações, consulte Tamanhos válidos de MTU da rede VPC.

Para mais informações sobre como mudar a configuração de MTU de uma rede padrão ou atual, consulte Alterar a configuração de MTU de uma rede VPC.

O exemplo a seguir cria uma rede com MTU de 8.896 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 o firewall.

Usar a opção multi-NIC para Multislice

Se você estiver usando um ambiente Multislice, 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 comandos a seguir para criar um roteamento IP personalizado para a rede e a sub-rede.

Observação: as duas sub-redes precisam ter intervalos de IP mutuamente exclusivos. Isso inclui os intervalos CIDR de IP primário e secundário na 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 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, é possível validar se as duas placas de interface de rede (NICs) estão sendo usadas. Para isso, configure um cluster XPK e adicione a flag --command ifconfig ao comando de criação de carga de trabalho XPK.

1. Use o comando workload create a seguir para mostrar a saída do comando ifconfig nos registros do console Google Cloud e verifique se eth0 e eth1 têm MTU definido como 8.896.

   Observação: no comando workload create, especifique a flag --docker-image para usar uma imagem do Docker atual ou a flag --base-docker-image para construir uma nova imagem de contêiner em que executar sua carga de trabalho XPK. A flag --base-docker-image usa a imagem de base especificada e adiciona o conteúdo do seu diretório atual.
   Observação: antes de executar esse comando, verifique se você criou um cluster XPK.

   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 você quiser ativar os registros de depuração ou usar o TensorBoard da Vertex AI, adicione os seguintes argumentos opcionais ao comando:

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

2. Verifique se eth0 e eth1 têm MTU definido como 8.896 conferindo a saída da carga de trabalho XPK nos registros do console Google Cloud .

Melhorar as configurações de TCP

Se você provisionou as Cloud TPUs usando recursos enfileirados, execute o comando a seguir para melhorar o desempenho da rede aumentando os limites do buffer de recebimento 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"'

Otimizar a performance da alocação de memória

A biblioteca tcmalloc é usada por padrão em VMs da Cloud TPU para melhorar o desempenho de modelos com alocações de memória consideráveis e frequentes. Isso é configurado pela variável de ambiente LD_PRELOAD.

No entanto, para algumas cargas de trabalho (por exemplo, DLRM com alocações de tabela de incorporação muito grandes), tcmalloc pode causar uma lentidão. Nesses casos, é possível reverter para a função malloc padrão removendo a definição da variável LD_PRELOAD na sessão do shell antes de executar o script de treinamento:

unset LD_PRELOAD

Usar o SkyPilot

Você pode usar a Cloud TPU v6e com o SkyPilot. O SkyPilot é um framework de código aberto que simplifica o processo de execução, gerenciamento e escalonamento de cargas de trabalho de IA. É possível adicionar informações de localização e preços relacionadas ao v6e ao SkyPilot. Para mais informações, consulte o exemplo de TPU v6e do SkyPilot.

Exemplos de treinamento

As seções a seguir fornecem exemplos de treinamento de modelos MaxText, MaxDiffusion e PyTorch na Cloud TPU v6e.

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

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

Treinar o MaxText e o MaxDiffusion na Cloud TPU v6e

As seções a seguir abordam o ciclo de vida de treinamento dos modelos MaxText e MaxDiffusion.

Em geral, as etapas de alto nível são:

1. Crie a imagem de base da carga de trabalho.
2. Execute a carga de trabalho usando XPK.
   1. Crie o comando de treinamento para a carga de trabalho.
   2. Implante a carga de trabalho.
3. Acompanhe a carga de trabalho e confira as métricas.
4. Exclua a carga de trabalho XPK se ela não for necessária.
5. Exclua o cluster XPK quando ele não for mais necessário.

Criar imagem de base

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

1. Clone o repositório que você quer usar e mude para o diretório dele:

   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 Google Cloud CLI:

   gcloud auth configure-docker
   

3. Crie a imagem do Docker usando o comando a seguir ou uma imagem de IA do JAX. Para mais informações sobre imagens de IA do JAX, consulte 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 ativa da CLI gcloud:

   gcloud config set project ${PROJECT_ID}
   

5. Se você estiver iniciando a carga de trabalho em uma máquina que não tem a imagem criada localmente, faça upload dela.

   Observação: verifique se você está usando o bash versão 4.2 ou mais recente.

   1. Defina a variável de ambiente CLOUD_IMAGE_NAME:

      export CLOUD_IMAGE_NAME=${USER}_runner
      

   2. Faça upload da imagem:

      bash docker_upload_runner.sh ${CLOUD_IMAGE_NAME}
      

Executar a carga de trabalho usando XPK

1. Defina as seguintes variáveis de ambiente se você não estiver usando os valores padrão definidos pelo MaxText ou pelo 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. Esse script será copiado como um comando de treinamento em uma etapa posterior.

   Não execute o script do modelo ainda.

   Observação: a criação automática de pontos de verificação não está ativada. Se você encontrar erros como Could not find a known environment for initializing distributed JAX, use JAX_FORCE_TPU_INIT=true.

   MaxText

   O MaxText é um LLM de código aberto de alto desempenho e altamente escalonável escrito em Python e JAX puros e destinado a TPUs e GPUs para treinamento 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 LLMs de peso aberto desenvolvidos pelo Google DeepMind com base na pesquisa e tecnologia do Gemini.

   Observação: para executar esse modelo, provisione uma TPU v6e-256 ou maior.

   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 última geração desenvolvido pela Mistral AI, que usa uma arquitetura esparsa de combinação de especialistas (MoE, na sigla em inglês).

   Observação: para executar esse modelo, provisione uma TPU v6e-256 ou maior.

   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 LLMs de peso aberto desenvolvidos pela Meta.

   Observação: para executar esse modelo, provisione uma TPU v6e-256 ou maior.

   Para ver um exemplo de como executar o Llama3 no PyTorch, consulte 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 e JAX puros que são executados em dispositivos XLA, incluindo TPUs e GPUs do Cloud. O Stable Diffusion é um modelo de texto latente para imagem que gera imagens fotorrealistas com base em qualquer entrada de texto.

   Você precisa instalar uma ramificação específica do Git para executar o MaxDiffusion, conforme mostrado no script de treinamento a seguir.

   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
   

   Observação: opcionalmente, defina o dcn_data_parallelism para execuções de várias fatias.

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 de variáveis de ambiente

   Variável	Descrição
   CLUSTER_NAME	O nome do cluster do XPK.
   ACCELERATOR_TYPE	O tipo de acelerador especifica a versão e o tamanho da Cloud TPU que você quer criar. Para mais informações sobre os tipos de aceleradores compatíveis com cada versão de TPU, consulte Versões de TPU.
   NUM_SLICES	O número de fatias de TPU.
   YOUR_MODEL_SCRIPT	O script do modelo a ser executado como um comando de treinamento.

4. Execute o modelo usando o script criado na etapa anterior. Você precisa especificar a flag --base-docker-image para usar a imagem de base do MaxText ou a flag --docker-image e a imagem que você quer usar.

   Você pode adicionar as seguintes flags opcionais:

   • Para ativar o registro de depuração, inclua a flag --enable-debug-logs. Para mais informações, consulte Depurar JAX no MaxText.
   • É possível criar um experimento da Vertex AI para fazer upload de dados para o TensorBoard da Vertex AI incluindo a flag --use-vertex-tensorboard. Para mais informações, consulte Monitorar o JAX no MaxText usando 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 sua carga de trabalho. Abra o link e clique na guia Registros para acompanhar sua carga de trabalho em tempo real.

Depurar o JAX no MaxText

Use comandos XPK complementares para diagnosticar por que o cluster ou a carga de trabalho não está em execução:

• Lista de cargas de trabalho do XPK
• Inspetor de XPK
• Ative o registro detalhado nos registros de carga de trabalho usando a flag --enable-debug-logs ao criar a carga de trabalho XPK.

Monitorar o JAX no MaxText usando a Vertex AI

Para usar o TensorBoard, sua conta de usuário do Google Cloud precisa ter a função aiplatform.user. Execute o seguinte comando para conceder esse papel:

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

Confira dados escalares e de perfil no TensorBoard gerenciado da Vertex AI.

1. Aumente as solicitações de gerenciamento de recursos (CRUD) para a zona que você está usando de 600 para 5.000. Isso 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 a Vertex AI:

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

3. Crie o cluster XPK usando a flag --create-vertex-tensorboard, conforme documentado em Criar o TensorBoard da Vertex AI. Você também pode executar esse comando em clusters atuais.

4. Crie seu experimento da Vertex AI ao executar a carga de trabalho do XPK usando a flag --use-vertex-tensorboard e a flag opcional --experiment-name. Para conferir a lista completa de etapas, consulte Criar um experimento da Vertex AI para fazer upload de dados no TensorBoard da Vertex AI.

Observação: se você estiver executando uma carga de trabalho que não é do MaxText, use este exemplo do MaxText para fazer as mudanças correspondentes no upload de dados.

Os registros incluem um link para um TensorBoard da Vertex AI, 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 é possível encontrar o link do TensorBoard da Vertex AI no console Google Cloud . Acesse Experimentos da Vertex AI no console Google Cloud . Selecione a região apropriada no menu suspenso.

O diretório do TensorBoard também é gravado no bucket do Cloud Storage especificado com ${BASE_OUTPUT_DIR}.

Excluir sua carga de trabalho XPK

Use o comando xpk workload delete para excluir uma ou mais cargas de trabalho com base no prefixo ou status do job. Esse comando pode ser útil se você enviou cargas de trabalho XPK que não precisam mais ser executadas ou se tiver jobs presos na fila.

Excluir o cluster do XPK

Use o comando xpk cluster delete para excluir o cluster:

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

Resultados de comparativo de mercado do MaxDiffusion

Executamos o script de treinamento do MaxDiffusion em uma v6e-4, uma v6e-16 e duas v6e-16. A tabela a seguir mostra as capacidades de processamento medidas.

	v6e-4	v6e-16	Dois v6e-16
Etapas de treinamento	0,069	0,073	0,13
Tamanho global do lote	8	32	64
Capacidade (exemplos/segundo)	115,9	438,4	492,3

Treinar modelos Llama usando PyTorch/XLA na Cloud TPU v6e

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

Acessar o Hugging Face e o modelo Llama 3

Você precisa de um token de acesso de usuário do Hugging Face para este exemplo. Para informações sobre como criar tokens de acesso do usuário, consulte a documentação do Hugging Face sobre tokens de acesso do usuário.

Você também precisa de permissão para acessar o modelo Llama-3-8B no Hugging Face. Para ter acesso, acesse o modelo Meta-Llama-3-8B no HuggingFace e solicite acesso.

Criar uma VM da Cloud TPU

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

1. Configure as 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 de variáveis de ambiente

   Variável	Descrição
   PROJECT_ID	O ID do seu projeto Google Cloud . Use um projeto atual ou crie um novo.
   TPU_NAME	O nome da TPU.
   ZONE	A zona em que a VM da TPU será criada. Para mais informações sobre as zonas compatíveis, consulte Regiões e zonas de TPU.
   ACCELERATOR_TYPE	O tipo de acelerador especifica a versão e o tamanho da Cloud TPU que você quer criar. Para mais informações sobre os tipos de aceleradores compatíveis com cada versão de TPU, consulte Versões de TPU.
   RUNTIME_VERSION	A versão do software da Cloud TPU.

2. Crie uma VM do 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 o forkpytorch-tpu/transformers dos transformadores e dependências do Hugging Face. Este exemplo foi testado com as seguintes versões de dependência:

• torch: compatível com 2.5.0
• torch_xla[tpu]: compatível com 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/'

Configurar arquivos de configuração do modelo

O comando de treinamento na próxima seção, Executar o modelo, usa dois arquivos de configuração JSON para definir parâmetros do modelo e configuração de paralelismo de dados totalmente fragmentados (FSDP, na sigla em inglês). O sharding do FSDP permite usar um tamanho de lote maior durante o treinamento, fragmentando os pesos do modelo em várias TPUs. Ao treinar com modelos menores, pode ser suficiente usar o paralelismo de dados e replicar os pesos em cada dispositivo. Para mais informações sobre como fragmentar tensores em dispositivos no PyTorch/XLA, consulte o guia do usuário do PyTorch/XLA SPMD.

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

   Importante: defina use_cache como false para treinamento.

   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 arquivo 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 o FSDP, consulte Paralelismo de dados totalmente fragmentados usando SPMD .

3. Faça upload dos arquivos de configuração para as VMs da Cloud TPU usando o seguinte comando:

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

Executar o modelo

Usando os arquivos de configuração criados na seção anterior, execute o script run_clm.py para treinar o modelo Llama-3-8B no conjunto de dados WikiText. O script de treinamento leva aproximadamente 10 minutos para ser executado em uma Cloud TPU v6e-8.

1. Faça login no Hugging Face na sua Cloud TPU usando o 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 o treinamento de 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'

Solução de problemas do PyTorch/XLA

Se você definiu as variáveis opcionais para depuração na seção anterior, o perfil do modelo será armazenado no local especificado pela variável PROFILE_LOGDIR. É possível extrair o arquivo xplane.pb armazenado nesse local e usar tensorboard para conferir os perfis no navegador seguindo as instruções do TensorBoard.

Se o PyTorch/XLA não estiver funcionando como esperado, consulte o Guia de solução de problemas, que tem sugestões para depurar, criar perfis e otimizar seu modelo.