Treinar um modelo usando a TPU7x (Ironwood)

Pré-lançamento

Este produto ou recurso está sujeito aos "Termos das Soluções Pré-GA" na seção Termos Gerais de Serviço dos Termos Específicos de Serviço. Os produtos e recursos pré-GA estão disponíveis no estado em que se encontram e podem ter assistência limitada. Para mais informações, consulte as descrições das fases de lançamento.

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

A TPU7x é a primeira versão da família Ironwood,a TPU de sétima geração do Google Cloud. 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 TPU7x, consulte Receitas de treinamento para TPU Ironwood no GitHub.

Provisionar TPUs

É possível provisionar e gerenciar TPUs v7x 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 precisa ou expansão de 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 precisar ter grande familiaridade com o Kubernetes. Use o XPK para criar rapidamente clusters do GKE e executar cargas de trabalho para prova de conceito e testes. Para mais informações, consulte o repositório do GitHub do XPK.
• GKE e TPU Cluster Director: o TPU Cluster Director está disponível por uma reserva no modo "Toda a capacidade", que oferece acesso total a toda a capacidade reservada (sem retenções) e visibilidade completa da topologia de hardware, do status de utilização e do status de integridade da TPU. Para mais informações, consulte Visão geral do modo de capacidade total.

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

Use o kit de processamento acelerado (XPK) para criar clusters do GKE para prova de conceito e testes. O XPK é uma ferramenta de linha de comando projetada para simplificar o provisionamento, o gerenciamento e a execução de cargas de trabalho de machine learning.

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, siga estas etapas:

• Verifique se você tem um projeto do Google Cloud com o faturamento ativado.
• Acessar o TPU7x. Para mais informações, entre em contato com sua equipe de conta.
• Verifique se a conta que você está usando com o XPK tem as funções listadas no repositório do XPK no GitHub.

Instalar o XPK e as dependências

1. Instale o XPK. Siga as instruções no repositório do GitHub do XPK.

2. Instale o Docker seguindo as instruções fornecidas pelo administrador ou as instruções de instalação oficiais. Depois de instalado, execute os seguintes comandos 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 projeto do Google Cloud .
   • YOUR_ZONE: a zona em que o cluster será criado. Na prévia, somente us-central1-c é compatível.
   • YOUR_CLUSTER_NAME: o nome do novo cluster.
   • YOUR_ACCELERATOR_TYPE: a versão e a topologia da TPU. Por exemplo, tpu7x-4x4x8. Para ver uma lista de topologias compatíveis, consulte Configurações aceitas.
   • YOUR_BUCKET_NAME: o nome do seu bucket do Cloud Storage, que será o diretório de saída para o treinamento do 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 NIC única e fração única

Escolha uma das seguintes opções para criar o cluster. Para ter uma performance ideal,recomendamos usar uma rede personalizada com MTU de 8.896 bytes.

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 sua 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 que você criou:

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

5. Crie o cluster do 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, crie 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 JobSet, evitando erros. Por padrão, o XPK usa e2-standard-16. Algumas zonas só aceitam tipos específicos de CPU. Portanto, talvez seja necessário alternar entre os tipos n1, n2 e e2. Caso contrário, você poderá encontrar erros de cota.

Criar ou fazer upload da imagem Docker do MaxText

Observação: o MaxText é compatível apenas com o Python 3.12. Crie seu ambiente virtual com 3.12 para instalar as dependências corretas.

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

Criar localmente

Os comandos a seguir copiam seu 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 dos comandos, uma imagem chamada maxtext_base_image será criada localmente. É possível usar a imagem local diretamente no comando xpk workload.

Fazer upload de uma imagem (opcional)

Depois de criar a imagem do Docker localmente usando as instruções na seção anterior, faça upload dela para o registro com o seguinte comando:

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 MaxText vai aparecer 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 este modelo para testar pequenas escalas de chips. 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 especificar a flag --docker-image e a imagem que você quer usar. Você pode 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 Diagnóstico de ML para otimizar e diagnosticar suas cargas de trabalho
• Executar uma carga de trabalho de treinamento usando uma receita otimizada para TPU7x
• Executar uma comparação de microbenchmark TPU7x