Treinar um modelo usando a TPU7x (Ironwood)

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

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.

Implantar uma carga de trabalho de treinamento com o MaxText e o Cluster Toolkit

Use o Cluster Toolkit para criar clusters do GKE prontos para produção e executar cargas de trabalho de treinamento.

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

Antes de começar

Antes de começar, siga estas etapas:

  1. 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 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, que será o diretório de saída para o treinamento de modelo.

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

    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)

Implante um cluster do GKE TPU7x seguindo as instruções em Implantar um cluster do GKE TPU 7x. Verifique se a topologia do pool de nós do cluster corresponde aos requisitos da carga de trabalho de treinamento, como 4x4x8 para a carga de trabalho de exemplo do MaxText abaixo.

Crie e faça 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 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 pequenas escalas de chips. O desempenho não está 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"

Implante a carga de trabalho de treinamento

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. Neste exemplo, $(date +%H%M) é anexado ao nome da carga de trabalho para garantir a exclusividade.

A seguir