Treinar um modelo usando a TPU7x (Ironwood)

Este documento descreve como provisionar recursos da TPU7x e oferece um exemplo de implantação de uma carga de trabalho de treinamento usando o MaxText e o XPK.

A TPU7x é a primeira versão da família Ironwood, Google Clouda TPU de sétima geração do Google. A geração Ironwood foi projetada para treinamento e inferência de IA em grande escala. Para mais informações, consulte TPU7x.

Para mais exemplos otimizados para a TPU7x, consulte Receitas de treinamento para a Ironwood TPU no GitHub (link em inglês).

Provisionar TPUs

É possível provisionar e gerenciar a TPU7x 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 conteinerizadas de machine learning. Use a Google Cloud CLI para criar manualmente a instância do cluster do GKE para personalização ou expansão precisa dos ambientes de produção do GKE. 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. Com ele, os profissionais de ML podem provisionar TPUs e executar jobs de treinamento sem ter grande familiaridade com o Kubernetes. Use o XPK para criar clusters do GKE e executar cargas de trabalho rapidamente para prova de conceito e testes. Para mais informações, consulte o repositório do GitHub do XPK (link em inglês).
• GKE e TPU Cluster Director: o TPU Cluster Director está disponível por meio de uma reserva no modo "All Capacity", que oferece acesso total a toda a capacidade reservada (sem retenções) e visibilidade completa da topologia de hardware da TPU, do status de utilização e da integridade. Para mais informações, consulte Visão geral do modo "All Capacity".

Implantar uma carga de trabalho de treinamento com o MaxText e o XPK

Use o Accelerated Processing Kit (XPK) para criar clusters do GKE para prova de conceito e testes.

As seções a seguir mostram como implantar uma carga de trabalho de treinamento usando MaxText e XPK.

Antes de começar

Antes de começar, conclua as seguintes etapas:

• Verifique se você tem um Google Cloud projeto com o faturamento ativado.
• Acesse a TPU7x. Para mais informações, entre em contato com a equipe de conta.
• Verifique se a conta que você está usando com o XPK tem os papéis listados no repositório do GitHub do XPK.

Instalar o XPK e as dependências

1. Instale o XPK. Siga as instruções no repositório do GitHub do XPK (link em inglês).

2. Instale o Docker usando as instruções fornecidas pelo administrador ou siga as instruções de instalação oficiais. Depois de instalado, execute os comandos a seguir para configurar o Docker e testar a instalação:

   gcloud auth configure-docker
   sudo usermod -aG docker $USER # relaunch the terminal and activate venv after running this command
   docker run hello-world # Test Docker
   

3. Configure as variáveis de ambiente a seguir:

   export PROJECT_ID=YOUR_PROJECT_ID
   export ZONE=YOUR_ZONE
   export CLUSTER_NAME=YOUR_CLUSTER_NAME
   export ACCELERATOR_TYPE=YOUR_ACCELERATOR_TYPE
   export BASE_OUTPUT_DIR="gs://YOUR_BUCKET_NAME"

   Substitua:

   • YOUR_PROJECT_ID: o ID do Google Cloud projeto do.
   • YOUR_ZONE: a zona em que o cluster será criado.
   • YOUR_CLUSTER_NAME: o nome do novo cluster.
   • YOUR_ACCELERATOR_TYPE: a versão e a topologia da TPU. Por exemplo, tpu7x-4x4x8. Para uma lista de topologias compatíveis, consulte Configurações compatíveis.
   • YOUR_BUCKET_NAME: o nome do bucket do Cloud Storage, que será o diretório de saída para o treinamento de modelo.

4. Se você não tiver um bucket do Cloud Storage, crie um usando o comando a seguir:

   gcloud storage buckets create ${BASE_OUTPUT_DIR} \
       --project=${PROJECT_ID} \
       --location=US \
       --default-storage-class=STANDARD \
       --uniform-bucket-level-access
   

Criar um cluster de fatia única e NIC único

Escolha uma das opções a seguir para criar o cluster. Recomendamos o uso de uma rede personalizada com MTU de 8.896 para um desempenho ideal.

Rede personalizada

Para criar uma rede personalizada com MTU de 8.896 e usá-la no cluster, siga estas etapas:

1. Defina variáveis de ambiente para os nomes da rede e do firewall:

   export NETWORK_NAME=NETWORK_NAME
   export NETWORK_FW_NAME=FIREWALL_NAME

   Substitua:

   • NETWORK_NAME: um nome para a rede.
   • FIREWALL_NAME: um nome para a regra de firewall de rede.

2. Crie uma rede personalizada com uma MTU de 8.896:

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

3. Crie uma regra de firewall que permita o tráfego TCP, ICMP e UDP na rede:

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

4. Defina uma variável de ambiente para os argumentos do cluster XPK para usar a rede criada:

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

5. Crie o cluster XPK. O comando a seguir provisiona a capacidade sob demanda:

   xpk cluster create --cluster=${CLUSTER_NAME} \
       --cluster-cpu-machine-type=n1-standard-8 \
       --num-slices=${NUM_SLICES} \
       --tpu-type=${ACCELERATOR_TYPE} \
       --zone=${ZONE} \
       --project=${PROJECT_ID} \
       --on-demand \
       --custom-cluster-arguments="${CLUSTER_ARGUMENTS}"

   Para usar a capacidade reservada, substitua --on-demand por --reservation=RESERVATION_NAME. Para usar VMs spot de TPU, substitua --on-demand por --spot.

Rede padrão

Se você não precisar de uma rede de MTU alta, poderá criar um cluster que use a rede VPC padrão. O comando a seguir provisiona a capacidade sob demanda:

xpk cluster create --cluster=${CLUSTER_NAME} \
    --cluster-cpu-machine-type=n1-standard-8 \
    --num-slices=${NUM_SLICES} \
    --tpu-type=${ACCELERATOR_TYPE} \
    --zone=${ZONE} \
    --project=${PROJECT_ID} \
    --on-demand

Para usar a capacidade reservada, substitua --on-demand por --reservation=RESERVATION_NAME. Para usar VMs spot de TPU, substitua --on-demand por --spot.

Observação: definir a flag --cluster-cpu-machine-type como n1-standard-8 (ou maior) garante que o pool de nós padrão tenha CPU suficiente para pods do sistema, por exemplo, o webhook do JobSet, evitando erros. Por padrão, o XPK usa e2-standard-16. Algumas zonas oferecem suporte apenas a tipos de CPU específicos. Portanto, talvez seja necessário alternar entre os tipos n1, n2 e e2. Caso contrário, poderão ocorrer erros de cota.

Criar ou fazer upload da imagem do Docker do MaxText

Observação: o MaxText oferece suporte apenas ao Python 3.12. Crie seu ambiente virtual com a versão 3.12 para instalar as dependências corretas.

É possível criar uma imagem Docker localmente usando scripts fornecidos pelo MaxText ou usar uma imagem pré-criada.

Criar localmente

Os comandos a seguir copiam o diretório local para o contêiner:

# Make sure you're running on a virtual environment with python3.12. If nothing is printed, you have the correct version.
[[ "$(python3 -c 'import sys; print(f"{sys.version_info.major}.{sys.version_info.minor}")' 2>/dev/null)" == "3.12" ]] || { >&2 echo "Error: Python version must be 3.12."; false; }

# Clone MaxText
git clone https://github.com/AI-Hypercomputer/maxtext.git
cd maxtext
git checkout maxtext-tutorial-v1.0.0

# Build the Docker image
bash docker_build_dependency_image.sh MODE=stable JAX_VERSION=0.8.2

Após a execução bem-sucedida dos comandos, uma imagem chamada maxtext_base_image será criada localmente. É possível usar a imagem local diretamente no comando de carga de trabalho do xpk.

Fazer upload da imagem (opcional)

Depois de criar a imagem Docker localmente usando as instruções na seção anterior, é possível fazer upload da imagem Docker do MaxText para o registro usando o comando a seguir:

export CLOUD_IMAGE_NAME="${USER}-maxtext-runner"
bash docker_upload_runner.sh CLOUD_IMAGE_NAME=${CLOUD_IMAGE_NAME}

Após a execução bem-sucedida desse comando, a imagem do MaxText será exibida em gcr.io com o nome gcr.io/PROJECT_ID/CLOUD_IMAGE_NAME.

Definir o comando de treinamento do MaxText

Prepare o comando para executar o script de treinamento no contêiner do Docker.

O modelo MaxText 1B é uma configuração no framework MaxText projetada para treinar um modelo de linguagem com aproximadamente 1 bilhão de parâmetros. Use esse modelo para testar escalas de chips pequenos. O desempenho não é otimizado.

export MAXTEXT_COMMAND="JAX_PLATFORMS=tpu,cpu \
    ENABLE_PJRT_COMPATIBILITY=true \
    python3 src/MaxText/train.py src/MaxText/configs/base.yml \
        base_output_directory=${BASE_OUTPUT_DIR} \
        dataset_type=synthetic \
        per_device_batch_size=2 \
        enable_checkpointing=false \
        gcs_metrics=true \
        run_name=maxtext_xpk \
        steps=30"

Implantar a carga de trabalho de treinamento

Execute o comando xpk workload create para implantar o job de treinamento. 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. É possível incluir a flag --enable-debug-logs para ativar o registro de depuração.

xpk workload create \
    --cluster ${CLUSTER_NAME} \
    --base-docker-image maxtext_base_image \
    --workload maxtext-1b-$(date +%H%M) \
    --tpu-type=${ACCELERATOR_TYPE} \
    --zone ${ZONE} \
    --project ${PROJECT_ID} \
    --command "${MAXTEXT_COMMAND}"
    # [--enable-debug-logs]

Os nomes das cargas de trabalho precisam ser exclusivos no cluster. Neste exemplo, $(date +%H%M) é anexado ao nome da carga de trabalho para garantir a exclusividade.

A seguir

• Use a plataforma Google Cloud ML Diagnostics para otimizar e diagnosticar suas cargas de trabalho
• Execute uma carga de trabalho de treinamento usando uma receita otimizada para a TPU7x.
• Execute um microbenchmark da TPU7x.