Treinar um modelo usando a TPU7x (Ironwood)

Neste documento, descrevemos como provisionar recursos da TPU v7x com o Cluster Toolkit e implantar uma carga de trabalho de treinamento usando o MaxText (JAX) ou o TorchTitan na TPU (PyTorch).

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 os seguintes repositórios no GitHub:

Antes de começar

Antes de começar, verifique se você atende aos seguintes requisitos:

Em seguida, siga estas etapas de configuração:

  1. Defina as variáveis de ambiente do projeto e do cluster:

    export PROJECT_ID=YOUR_PROJECT_ID
export ZONE=YOUR_ZONE
export CLUSTER_NAME=YOUR_CLUSTER_NAME
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.
    • YOUR_CLUSTER_NAME: o nome do novo cluster.
    • YOUR_BUCKET_NAME: o nome do bucket do Cloud Storage para a saída de treinamento.

  2. Se você não tiver um bucket do Cloud Storage, crie um:

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

Implantar o cluster TPU7x (Ironwood)

Siga as instruções em Implantar um cluster de TPU 7x do GKE para implantar um cluster de TPU 7x do Google Kubernetes Engine (GKE). Verifique se a topologia do pool de nós do cluster corresponde aos requisitos da carga de trabalho de treinamento, como 4x4x8 para o exemplo do MaxText ou 2x2x4 para o exemplo do TorchTitan.

Implantar uma carga de trabalho de treinamento com o MaxText

As seções a seguir mostram como criar uma imagem de contêiner e enviar um job de treinamento do MaxText com o Cluster Toolkit.

Criar e fazer upload da imagem Docker do MaxText

Crie e envie a imagem Docker do MaxText:

# 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

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

Definir o comando de treinamento do MaxText

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

A configuração do modelo MaxText 1B treina 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_training \
        steps=30"

Enviar o job de treinamento do MaxText

Envie o job de treinamento usando gcluster job submit. Para detalhes sobre as opções de envio de jobs, consulte o Guia de jobs do Cluster Toolkit:

gcluster job submit \
    --name="maxtext-1b-$(date +%H%M)" \
    --cluster=${CLUSTER_NAME} \
    --project=${PROJECT_ID} \
    --location=${ZONE} \
    --compute-type=tpu7x-standard-4t \
    --topology=4x4x8 \
    --image="gcr.io/${PROJECT_ID}/${CLOUD_IMAGE_NAME}" \
    --command="${MAXTEXT_COMMAND}"

Os nomes das cargas de trabalho precisam ser exclusivos no cluster. Este exemplo anexa $(date +%H%M) ao nome da carga de trabalho.

Implantar uma carga de trabalho de treinamento com o TorchTitan

O TorchTitan usa o TorchTPU para executar o treinamento distribuído do PyTorch no Cloud TPU. Para mais informações, consulte o README do TorchTitan na TPU.

A seguir