Esta página foi traduzida pela API Cloud Translation.

Prepare um modelo com a TPU7x (Ironwood)

Este documento descreve como aprovisionar recursos de TPUv7 e dá um exemplo de implementação de uma carga de trabalho de preparação usando o MaxText e o XPK.

O TPU7x é o primeiro lançamento da família Ironwood, Google Cloudo TPU de sétima geração. A geração Ironwood foi concebida para a inferência e a preparação de IA em grande escala. Para mais informações, consulte o artigo TPU7x.

Para mais exemplos otimizados para a TPU7x, consulte Receitas de preparação para a TPU Ironwood no GitHub.

Aprovisione TPUs

Pode aprovisionar e gerir TPUs v7 usando os seguintes métodos:

GKE: pode usar o GKE para aprovisionar e gerir TPUs como um conjunto de aceleradores para as suas cargas de trabalho de aprendizagem automática contentorizadas. Use a Google Cloud CLI para criar manualmente a instância do cluster do GKE para uma personalização precisa ou uma expansão dos ambientes de produção do GKE existentes. Para mais informações, consulte o artigo Acerca das TPUs no GKE.
GKE e XPK: o XPK é uma ferramenta de linha de comandos que simplifica a criação de clusters e a execução de cargas de trabalho no GKE. Foi concebido para que os profissionais de ML possam aprovisionar TPUs e executar tarefas de preparação sem precisar de conhecimentos profundos do Kubernetes. Use o XPK para criar rapidamente clusters do GKE e executar cargas de trabalho para testes e provas de conceito. Para mais informações, consulte o repositório do GitHub do XPK.
GKE e TPU Cluster Director: o TPU Cluster Director está disponível através de uma reserva no modo de capacidade total, que lhe dá acesso total a toda a capacidade reservada (sem restrições) e visibilidade total da topologia do hardware da TPU, do estado de utilização e do estado de funcionamento. Para mais informações, consulte a vista geral do modo de capacidade total.

Implemente uma carga de trabalho de preparação com o MaxText e o XPK

Use o Accelerated Processing Kit (XPK) para criar clusters do GKE para validação de conceito e testes. O XPK é uma ferramenta de linha de comandos concebida para simplificar o aprovisionamento, a gestão e a execução de cargas de trabalho de aprendizagem automática.

As secções seguintes mostram como implementar uma carga de trabalho de preparação usando o MaxText e o XPK.

Antes de começar

Antes de começar, conclua os seguintes passos:

Certifique-se de que tem um Google Cloud projeto com a faturação ativada.
Aceda ao TPU7x. Para mais informações, contacte a equipa da conta.
Certifique-se de que a conta que está a usar com o XPK tem as funções indicadas no repositório GitHub do XPK.

Instale o XPK e as dependências

Instale o XPK. Siga as instruções no repositório do GitHub do XPK.
Instale o Docker seguindo as instruções fornecidas pelo seu administrador ou as instruções de instalação oficiais. Após a instalação, 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
```
Defina as seguintes variáveis de ambiente:
```
export PROJECT_ID=YOUR_PROJECT_ID
export ZONE=YOUR_ZONE
export CLUSTER_NAME=YOUR_CLUSTER_NAME
export ACCELERATOR_TYPE=YOUR_ACCELERATOR_TYPE
export RESERVATION_NAME=YOUR_RESERVATION_NAME
export BASE_OUTPUT_DIR="gs://YOUR_BUCKET_NAME"
```
Substitua o seguinte:
- YOUR_PROJECT_ID: o ID do seu Google Cloud projeto.
- YOUR_ZONE: a zona na qual criar o cluster. Para a pré-visualização, apenas é suportado o formato us-central1-c.
- 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 das topologias suportadas, consulte o artigo Configurações suportadas.
- YOUR_RESERVATION_NAME: o nome da sua reserva. Para reservas partilhadas, use projects/YOUR_PROJECT_NUMBER/reservations/YOUR_RESERVATION_NAME.
- YOUR_BUCKET_NAME: o nome do seu contentor do Cloud Storage, que será o diretório de saída para a preparação do modelo.

Se não tiver um contentor do Cloud Storage, crie um com o seguinte comando:

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

Crie um cluster de fatia única com uma única NIC

Siga as instruções na secção Configure a MTU para otimizar a configuração da rede.

Preencha a variável ${CLUSTER_ARGUMENTS}, que vai usar no comando xpk cluster create:

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

Crie o cluster do GKE com node pools TPU7x através do comando xpk cluster create:
```
xpk cluster create \
    --project=${PROJECT_ID} \
    --zone=${ZONE} \
    --cluster ${CLUSTER_NAME} \
    --cluster-cpu-machine-type=n1-standard-8 \
    --tpu-type=${ACCELERATOR_TYPE} \
    --reservation=${RESERVATION_NAME} \
    --custom-cluster-arguments="${CLUSTER_ARGUMENTS}"
```
Definir a flag --cluster-cpu-machine-type como n1-standard-8 (ou superior) garante que o conjunto de nós predefinido tem CPU suficiente para pods do sistema, por exemplo, o webhook JobSet, o que evita erros. Por predefinição, o XPK usa e2-standard-16. Algumas zonas só suportam tipos de CPU específicos, pelo que pode ter de alternar entre os tipos n1, n2 e e2. Caso contrário, pode encontrar erros de quota.

Adicione uma exclusão de manutenção para impedir atualizações para o cluster:

gcloud container clusters update ${CLUSTER_NAME} \
    --zone=${ZONE} \
    --add-maintenance-exclusion-name="no-upgrade-next-month" \
    --add-maintenance-exclusion-start="EXCLUSION_START_TIME" \
    --add-maintenance-exclusion-end="EXCLUSION_END_TIME" \
    --add-maintenance-exclusion-scope="no_upgrades"

Substitua o seguinte:

EXCLUSION_START_TIME: a hora de início selecionada para a exclusão de manutenção no formato YYYY-MM-DDTHH:MM:SSZ.
EXCLUSION_END_TIME: a hora de fim selecionada para a exclusão de manutenção no formato YYYY-MM-DDTHH:MM:SSZ.

Crie ou carregue a imagem do Docker do MaxText

Pode criar uma imagem do Docker localmente através de scripts fornecidos pelo MaxText ou usar uma imagem pré-criada.

Crie localmente

Os seguintes comandos copiam o seu diretório local para o contentor:

# 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

# Custom Jax and LibTPU wheels
pip download libtpu==0.0.28.dev20251104+nightly -f "https://storage.googleapis.com/jax-releases/libtpu_releases.html"
pip download --pre jax==0.8.1.dev20251104 jaxlib==0.8.1.dev20251104 --index https://us-python.pkg.dev/ml-oss-artifacts-published/jax/simple/

# Build the Docker image
bash docker_build_dependency_image.sh MODE=custom_wheels

Após a execução bem-sucedida dos comandos, deve ver uma imagem denominada maxtext_base_image criada localmente. Pode usar a sua imagem local diretamente no comando de carga de trabalho xpk.

Carregue uma imagem (opcional)

Depois de criar a imagem Docker localmente através das instruções na secção anterior, pode carregar a imagem Docker do MaxText para o registo através do 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 deste comando, deve ver a imagem MaxText em gcr.io com o nome gcr.io/PROJECT_ID/CLOUD_IMAGE_NAME.

Defina o comando de preparação MaxText

Prepare o comando para executar o script de preparação no contentor Docker.

O modelo MaxText 1B é uma configuração na estrutura MaxText concebida para formar um modelo de linguagem com aproximadamente 1000 milhões de parâmetros. Use este modelo para experimentar escalas de chips pequenas. 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_xpk \
        steps=30"

Implemente a carga de trabalho de preparação

Execute o comando xpk workload create para implementar a tarefa de preparação. Tem de especificar o parâmetro --base-docker-image para usar a imagem base MaxText ou especificar o parâmetro --docker-image e a imagem que quer usar. Pode optar por incluir a flag --enable-debug-logs para ativar o registo 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 têm de ser exclusivos no cluster. Neste exemplo, $(date +%H%M) é anexado ao nome da carga de trabalho para garantir a exclusividade.